PDP: Parameter-free
Differentiable Pruning

Chapter 1: The Key Insight

Here is the trick, stated as plainly as the paper does. A weight w is in one of two symbolic states: Z ("to-prune") or M ("to-keep"). Instead of forcing a hard choice, give each weight a soft mask m(w) ∈ [0,1] — the probability it survives — and multiply the weight by it: ŵ = m(w)·w.

Now the network runs on the masked weights ŵ. Because m(w) is a smooth function of w, the chain rule flows straight through it. The loss can nudge w not just to fit the data, but to push itself toward "definitely keep" (m→1) or "definitely kill" (m→0).

Why does "parameter-free" matter so much that Apple put it in the title? Three reasons:

Chapter 2: Why Hard Masks Kill Weights for Good

To feel why a soft mask is worth the trouble, we first have to watch a hard mask fail. This is the failure mode PDP was designed to escape, and it's worth understanding viscerally.

Classic magnitude pruning with a threshold t uses a hard mask: keep w if |w| ≥ t, else set it to 0. The forward pass uses ŵ = \mathbb{1}[|w|≥t]·w. During backprop, the gradient that reaches w is multiplied by that same hard mask:

Read what this says. The moment a weight drops below t, its mask becomes 0, so its gradient becomes 0. It stops receiving updates. It is frozen at whatever near-zero value it had — permanently dead. Even if, later in training, that weight would have become important, it can never grow back, because it gets no signal to do so.

So hard-mask pruning has two diseases at once: (1) no recovery — pruned weights are gradient-dead, and (2) sparsity drift — when you steer the threshold, the achieved sparsity lurches unpredictably. PDP cures both: soft masks keep every weight's gradient alive, and deriving t from the target ratio r pins sparsity exactly.

Drag w toward the threshold (dashed line) and toggle the mode. With the hard mask, the gradient multiplier is a brick wall — 1 on one side, 0 on the other. With PDP's soft mask, the multiplier is a smooth S-curve: a weight just below the line still gets a fraction of its gradient, so it can climb back.

Chapter 3: Building the Soft Mask from First Principles

Let's derive the mask instead of being handed it. We want two functions of magnitude: z(|w|) = chance of being pruned (state Z), and m(|w|) = chance of being kept (state M). What must they satisfy to be a sensible magnitude-based mask?

• Bigger weights are safer: if |a| > |b|, then m(|a|) > m(|b|) and z(|a|) < z(|b|). (Monotonic.)
• They are complementary: z(w) + m(w) = 1. A weight is either kept or pruned; the two chances sum to one.

Now the pivotal observation. Because m rises monotonically from 0 to 1 and z falls from 1 to 0, and they're continuous, there must be exactly one magnitude where they cross. Call it the threshold t. At that point both are 0.5 — the weight is a perfect coin-flip:

With the boundary conditions m(0)=0 (a zero weight is surely pruned), m(∞)=1 (a huge weight is surely kept), and m(t)=½, any smooth monotone function will do. PDP picks the most natural one in deep learning: a softmax over two logits.

Writing out the two-class softmax for the keep probability gives PDP's soft mask:

Let's verify the boundary conditions by hand — this is exactly the kind of check that builds trust:

• At |w| = t: numerator = e^{t^2/τ}, denominator = e^{t^2/τ} + e^{t^2/τ} = 2e^{t^2/τ}, so m = 1/2. ✓
• As |w| → ∞: e^{w^2/τ} dwarfs e^{t^2/τ}, so m → 1. ✓
• At w = 0: numerator = e^0 = 1, denominator = 1 + e^{t^2/τ}; since t>0 this is small (→ 0 as τ→0). ✓

Chapter 4: The PDP Lab — Watch the Mask Live

This is the heart of the paper made interactive. Below is a real layer's weights, drawn from a bell-shaped distribution. You control the target prune ratio r and the temperature τ. PDP does the rest: it finds the threshold t that hits your sparsity target, computes m(w) for every weight, and shows you the masked weights ŵ=m(w)·w.

The top plot shows the soft mask curve m(w) over weight magnitude, with the threshold t marked. The bottom shows each weight as a dot — colored by how much of it survives. Watch what happens to weights right at the boundary: they are the gray, uncertain ones PDP keeps fighting over.

The "Show hard-mask result" button overlays where a plain magnitude prune would chop. PDP at small τ converges to that line — but during training it approaches it gradually, giving boundary weights many chances to prove themselves before the cut is final.

Chapter 5: Finding the Threshold t (No Learning Required)

We've been treating t as given. Where does it come from? This is the "parameter-free" magic: t is computed from the current weights and the target ratio r, fresh after every SGD step. No gradient, no optimizer state — just sorting.

The recipe, using \text{topK} (the largest k magnitudes) and \text{bottomK} (the smallest k):

Why the midpoint, not just the cutoff? Because the soft mask needs m(t)=0.5 to sit between the two camps. Putting t exactly halfway means the smallest kept weight gets m>0.5 (leaning keep) and the largest pruned weight gets m<0.5 (leaning prune) — symmetric, exactly as the equal-chance derivation demands.

Where does r itself come from? Three options, all compatible with PDP: (a) sort weights globally after a few warm-up epochs and back out a per-layer r that meets the global target; (b) a budget-allocation algorithm that assigns more sparsity to less-critical layers; or (c) a human just sets it. PDP doesn't care — it only needs r per layer to derive t.

Chapter 6: The Hidden Gradient Boost

Now the deepest part of the paper — and the most beautiful. Let's differentiate the masked weight ŵ = m(w)·w with respect to w, and see what PDP secretly does to the gradient. (We treat t as fixed within a step, which it is.)

Start from ŵ = m(w)\,w. By the product rule:

We need dm/dw. A softmax of two logits has the standard derivative dm = m(1-m)\,d(\text{logit difference}). The logit difference here is (w^2-t^2)/τ, whose derivative w.r.t. w is 2w/τ. So:

Substitute back:

The gradient that actually updates w is Δw = (dŵ/dw)·Δŵ, where Δŵ=∂L/∂ŵ is the upstream gradient. So:

When is that bonus largest? The term m(1-m) peaks at m=0.5 — exactly at the threshold, where the prune decision is most uncertain. So PDP injects the strongest extra push precisely into the weights it can't yet decide about. The check-cases:

• m≈0 (clearly prune): both terms → 0. The weight barely moves — same as any pruning method.
• m≈1 (clearly keep): m(1-m)→0, so Δw ≈ Δŵ — a normal, full update.
• m≈0.5 (on the fence): m(1-m) is maximal → the boundary weight gets boosted movement toward whatever direction lowers the loss.

Chapter 7: Second Chances — Recovery and Decision Flipping

Put Chapters 2 and 6 together and you get PDP's signature behavior: a weight's prune decision can flip, repeatedly, all the way to the end of training. A near-zero weight that keeps receiving keep-ward gradients (scaled, but never zeroed) can grow back above t and rejoin the network. A large weight that keeps getting pushed toward zero can eventually fall below t and die. The loss, not a one-time magnitude snapshot, decides.

This matters because pruning decisions made early in training are made with bad information — the weights haven't settled. STR's early lock-in bakes in mistakes. PDP's perpetual second chances let it correct those mistakes as the weights mature. The cost is that train-time and test-time behavior differ slightly (test uses a hard mask), but the recovery upside dominates.

Each cell is a kernel. White = its keep/prune decision flipped this epoch; dark = stable. Press replay and compare: PDP's grid keeps sparkling with white cells late into training; STR's grid goes dark and stays dark early. White, here, is the visual signature of "still thinking."

Chapter 8: One Method, Many Shapes — Structured & Channel Pruning

So far we've pruned individual weights anywhere (unstructured / "random" sparsity). That's great for model size but hardware often can't exploit scattered zeros for speed. Two structured patterns are what actually accelerate inference — and PDP handles both with almost no change.

N:M sparsity (e.g. 2:4)

The constraint: in every consecutive group of M weights, exactly N may survive. NVIDIA Ampere+ GPUs have hardware that runs 2:4 sparsity at ~2× throughput — but only if the zeros follow this pattern exactly. PDP's adaptation is almost embarrassingly simple:

Channel pruning

Here you remove whole output channels (filters), which shrinks the actual tensor dimensions and speeds up any hardware. PDP's tweak: instead of using each scalar weight's magnitude, compute the L2 norm of each entire channel and use that as the "magnitude" fed into the threshold and mask. A channel with a small norm gets a small mask; the whole channel fades out together, differentiably.

where c is a channel's weight vector and ‖c‖_2 replaces |w|. Every weight in that channel is multiplied by the same channel mask m(c) — so the channel survives or dies as a unit, and the decision stays differentiable.

Chapter 9: Does It Work? The Receipts

A pruning method is only as good as the accuracy it keeps at high sparsity and the cost it saves. PDP claims state-of-the-art across vision and language, structured and unstructured. Highlights:

• MobileNet-v1, ImageNet, 86.6% sparse: 68.2% top-1 — +1.7% over the previous best at the same sparsity. A hard model (already compact) made far sparser without falling apart.
• BERT, MNLI, 90% sparse: 83.1% accuracy vs. 81.5% for the next best. Differentiable pruning that holds up on a language task at extreme sparsity.
• ResNet18, 1:4 structured: +3.6% top-1 over the prior structured-pruning state of the art — the largest gap in the paper.
• GPT-2 (163M), 75% sparse, OpenWebText: PDP matches OptG's perplexity (33.7) while OptG needs +124M extra trainable params and PDP needs 0 — and PDP trains at 42% lower GPU cost.

# PDP soft mask — the entire core idea, in ~12 lines.
import torch, torch.nn.functional as F

def pdp_mask(W, r, tau=0.03):
    # W: weight tensor; r: target prune ratio; tau: temperature
    a   = W.abs().flatten()
    k   = int(round((1 - r) * a.numel()))   # how many to keep
    s   = torch.sort(a, descending=True).values
    # t = midpoint of smallest-kept and largest-pruned magnitude
    t   = (s[k-1] + s[k]) / 2
    # two-class softmax: keep-logit w^2, prune-logit t^2
    logits = torch.stack([W**2, t**2 * torch.ones_like(W)]) / tau
    m   = F.softmax(logits, dim=0)[0]  # keep-probability in [0,1]
    return m * W                          # masked weights; gradient flows!

# Forward pass simply uses w_hat = pdp_mask(layer.weight, r). No new params,
# no optimizer changes. Recompute every step so t tracks the weights.

Chapter 10: Connections & Cheat Sheet

PDP sits at the crossroads of three ideas you may already know: magnitude pruning (which weight to cut), the straight-through estimator (how others fake a gradient through a hard mask — PDP needs no STE because its mask is genuinely smooth), and Gumbel-softmax / soft relaxations (turning discrete choices into differentiable ones — PDP is a parameter-free cousin).

Where to go next

• Magnitude & iterative pruning — the foundation PDP improves on. See the TinyML pruning lessons.
• Lottery Ticket Hypothesis — the idea that sparse subnetworks exist from the start; PDP is one way to find them by training.
• Quantization — the other half of model compression; often stacked on top of pruning for deployment.
• Knowledge distillation — an alternative compression route (and one some pruning methods lean on; PDP deliberately does not).

The one-sentence summary to carry away: PDP turns pruning into a temperature-controlled softmax between each weight and a derived threshold, so the task loss can decide which weights die — and reverse that decision — without a single extra parameter.