Learning Sparsity Itself:
Soft Threshold Reparameterization

Chapter 1: The Key Insight

The whole paper turns on one move. Pruning has always been an external operation: train the weights, then apply some rule to zero out the small ones. That rule has a knob — a threshold α, where any weight smaller than α in magnitude gets killed. Traditionally a human or a heuristic sets α.

STR's insight: make α a learnable parameter, one per layer, and write the forward pass so that the loss is a smooth function of α. Then α gets a gradient just like any weight, and SGD tunes it. Each layer discovers its own threshold — which means each layer discovers its own sparsity.

The operator at the heart of it

The trick is to use the soft-threshold operator — a function statisticians have used for sparse regression since 1995. For a single weight w and threshold α:

In words: if |w| is below α, the output is exactly zero (pruned). If |w| is above α, you keep its sign but shrink its magnitude by α. The forward pass uses S(W, α) instead of the raw weight W. We will dissect this operator in Chapter 3 — for now, notice it is continuous in α, which is exactly what makes α trainable.

What this buys you

Because every layer learns its own α, the per-layer sparsity is non-uniform and optimized. Layers that can afford to lose weights become very sparse; layers that can't stay denser. And because the gradient update is multiplied by an indicator that the weight survived, the gradient itself is sparse — STR fixes the two classic pains of threshold-based pruning at once. We'll prove both claims.

Chapter 2: Background — Hard vs Soft Thresholding

Pruning is, mathematically, an old friend: sparse optimization. You want to minimize a loss while keeping the number of non-zero parameters small. Formally:

where ‖W‖₀ is the L0 "norm" — literally the count of non-zero entries — and k is your parameter budget. This is the cleanest way to write "keep at most k weights." The trouble: the L0 constraint is combinatorial and non-convex, so you can't just run gradient descent on it.

The classical fix: projected gradient descent

The decades-old workaround is projected gradient descent (PGD): take a normal gradient step, then project back onto the feasible set (snap the iterate to something sparse). Two projections dominate the literature, depending on which sparsity-encouraging set you project onto:

• Hard thresholding — project onto the L0 ball. Keep the k largest-magnitude weights, zero the rest. Each survivor's value is untouched. This is exactly "magnitude pruning."
• Soft thresholding — project onto the L1 ball (the convex relaxation of L0). Zero out small weights AND shrink the survivors toward zero by the threshold. This is the proximal operator of the L1 penalty.

Two things break when you bring PGD to deep nets

The paper is blunt about why naive PGD pruning is painful at the scale of modern networks:

1. Dense gradients and dense intermediate iterates. PGD takes a full dense gradient step then projects. So between steps the weights are dense again — you're computing and storing billion-element dense updates even though the final model is sparse.
2. The budget k is global and indivisible. You're told "keep k weights total" but not how to split k across layers. That's the exact heuristic-budget problem from Chapter 0.

STR's reparameterization will dissolve both. The first one disappears because the gradient update gets multiplied by a survival indicator (Chapter 5). The second disappears because each layer learns its own α (Chapter 6).

Chapter 3: The Soft-Threshold Operator, Dissected

Everything rides on this one function, so let's make it concrete. For a scalar weight w and threshold α ≥ 0:

Decompose it symbol by symbol, each with a plain-language role:

• |w| — the weight's magnitude. "How big is this connection, ignoring direction?"
• α — the threshold. The moat: any weight whose magnitude doesn't clear it gets drowned (set to 0).
• |w| − α — how far the weight clears the moat. Negative if it doesn't clear it.
• ReLU(·) — clamps the negatives to 0. This is what performs the actual pruning.
• sign(w) — re-attaches the original direction, since |w| threw it away.

Worked example with real numbers

Set α = 0.3. Push a few weights through:

Notice the regime: anything within ±0.3 vanishes; everything else loses 0.3 of its size. Raise α and the dead zone widens, more weights vanish, survivors shrink harder. That single dial controls the layer's sparsity. Now drag it yourself.

Chapter 4: The Reparameterization — SHOWCASE

We now assemble the actual method. The naive idea — "set α and apply S" — still needs a human to set α. STR closes the loop by writing α as a smooth function of a free, trainable parameter s:

where g is any continuous, increasing function (the paper uses the sigmoid for unstructured sparsity). Why route α through g rather than train α directly? Two reasons:

• Keep α ≥ 0 automatically. A negative threshold is meaningless. Sigmoid maps any real s to (0,1), so α stays in a valid range without constraints.
• Control the dynamics. A "sharp" g would let the threshold explode and over-prune; a "flat" g learns too slowly. Sigmoid is a gentle, bounded choice. The paper initializes s at a value s_init so that α = g(s) starts near zero — training begins essentially dense and grows sparse.

The full training loop, in pseudocode

# STR: a 2-D conv whose weights are soft-thresholded by a LEARNED threshold
import torch, torch.nn as nn, torch.nn.functional as F

class STRConv(nn.Conv2d):
    def __init__(self, *a, s_init=-5.0, **k):
        super().__init__(*a, **k)
        # s is a SINGLE learnable scalar per layer -> one threshold per layer
        self.s = nn.Parameter(torch.tensor(s_init))   # weight-decayed like any param

    def sparse_weight(self):
        alpha = torch.sigmoid(self.s)               # g(s): threshold in (0,1), starts ~0
        w = self.weight
        return torch.sign(w) * F.relu(torch.abs(w) - alpha)  # S_g(W, s)

    def forward(self, x):
        return self._conv_forward(x, self.sparse_weight(), self.bias)

# Training is ORDINARY SGD. Sparsity is a side effect of weight decay on s + W.
opt = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9,
                      weight_decay=1e-4)   # lambda controls final sparsity

That is the entire method. No mask buffers, no per-step top-k, no pruning schedule. Just an extra scalar per layer and standard SGD. Now watch it run.

Chapter 5: The Gradients — Why the Update Is Sparse

Time to earn the claim that STR fixes the "dense gradient" pain of classical PGD. Write the loss with the reparameterization, add standard L2 decay, and take a gradient-descent step on each weight tensor. The update equation from the paper is:

Read every piece:

• (1 − ηλ)W_l — ordinary weight decay: shrink the weight a little each step.
• η — learning rate. ∇_SL — gradient of the loss w.r.t. the thresholded weight S.
• 1{S ≠ 0} — the survival indicator: 1 where the weight cleared the threshold, 0 where it was pruned. This is the crucial factor.

Where the indicator comes from

By the chain rule, the gradient w.r.t. the raw weight W is ∇_SL · ∂S/∂W. And what is ∂S/∂W? Differentiate S(w) = sign(w)·ReLU(|w| − α):

• In the dead zone (|w| < α): S is flat at 0, so ∂S/∂w = 0.
• Outside (|w| > α): S is a shifted identity with slope 1, so ∂S/∂w = 1.

That derivative is the indicator 1{|w| > α} = 1{S ≠ 0}. (At the kink S is non-differentiable, so the paper uses a sub-gradient — pick 0 or 1; it doesn't matter in practice.)

The gradient on the threshold itself

And s is trainable because L is continuous in s: as α = g(s) increases, more weights cross into the dead zone and the survivors shrink, all smoothly. The sub-gradient ∂L/∂s exists and SGD descends it, while weight decay on s nudges the threshold upward. The tug-of-war between "loss wants accuracy (low α)" and "decay wants small s" settles each layer at its own α.

Chapter 6: The Learnt Budget

This is the payoff. Because every one of ResNet50's 54 layers learns its own α, STR produces a non-uniform sparsity budget that no human designed. The paper's Figure 3 plots the final per-layer thresholds; Figure 6 plots the resulting per-layer sparsity. The shape is striking and consistent.

Why the learnt budget transfers

A second surprise: the sparsity budget STR learns on one task is a good backbone for other tasks. Take the per-layer sparsity STR discovered, freeze those budgets, and retrain — you get an efficient sparse backbone that generalizes. The budget encodes which layers structurally matter, not just task-specific quirks.

The training trajectory of sparsity

During training the overall sparsity follows a smooth S-curve (the paper's Figure 2): near-dense early, accelerating in the middle, leveling off. The authors note this looks like the hand-crafted GMP pruning schedule — except STR was never told to follow it. It learned the same shape that practitioners had been hand-engineering for years.

Chapter 7: From Unstructured to Structured Sparsity

So far every weight is pruned individually — unstructured sparsity. It gives the highest compression but the resulting scattered zeros are hard for standard hardware to exploit. The paper shows STR generalizes cleanly to structured sparsity, which removes whole structures (filters, channels, or — for RNNs — entire dimensions of rank).

The low-rank version, step by step

For an RNN weight matrix W, the paper does the following each forward pass:

The threshold α on the singular values is, once again, learned via the same reparameterization. So the network learns its own rank per layer. For structured sparsity the paper uses the exponential as g (instead of sigmoid) — it allows the threshold to grow large enough to aggressively cut singular values.

Why low-rank is "hardware-friendly"

A rank-r matrix W' ∈ R^m×n is stored as A ∈ R^m×r and B ∈ R^r×n with W' = AB. A matrix-vector product drops from mn multiplies to r(m+n) — a dense, regular computation that GPUs love, unlike scattered unstructured zeros. The paper reports STR matching or beating prior low-rank RNN methods (e.g. on Google-12 keyword spotting) while learning the rank automatically.

Chapter 8: The Experiments

STR is benchmarked on ImageNet-1K with ResNet50 and MobileNetV1, against the strongest pruning baselines of the day: GMP, DNW, DSR, SNFS, RigL (with and without the ERK budget), and DPF. The headline: across every sparsity regime from 80% to 99%, STR sits on the accuracy-vs-sparsity frontier — most accurate at each size — while using the fewest FLOPs.

ResNet50 on ImageNet (dense baseline: 77.01% top-1)

MobileNetV1: efficiency on an already-efficient net

MobileNetV1 is designed to be small, so it's a harder pruning target. STR still wins: at 90% sparsity it boosts accuracy by ~0.3% over baselines while cutting FLOPs by up to 50% — the clearest demonstration that the learnt budget is buying real compute savings, not just parameter count.

Reading the frontier yourself

Chapter 9: Connections & Cheat Sheet

STR's lasting idea isn't a particular accuracy number — it's the pattern: take a discrete, hand-tuned design choice (the sparsity budget) and make it a smooth, learnable parameter folded into the forward pass. That pattern recurs everywhere in modern model compression.

Where it sits in the landscape

• vs Magnitude Pruning / GMP — GMP applies a hand-scheduled global threshold; STR learns a per-layer threshold by gradient. STR's learnt budget recovers GMP's S-curve shape automatically and beats it on FLOPs.
• vs Lottery Ticket / RigL (sparse-to-sparse) — those keep the model sparse throughout training to save training cost; STR is dense-to-sparse, trading training memory for a better learned budget and higher final accuracy.
• vs L0 / variational dropout — those also learn sparsity, but via stochastic gates and extra memory/compute (VD costs 2× memory, 4× compute); STR is a deterministic operator with one scalar per layer.
• vs Quantization — an orthogonal axis of compression (fewer bits per weight vs fewer weights). The two compose: prune with STR, then quantize the survivors.

The cheat sheet

The whole method in five lines

# STR in five lines. That is genuinely all of it.
alpha = torch.sigmoid(s)                              # learned threshold, one s per layer
W_sparse = torch.sign(W) * F.relu(W.abs() - alpha) # soft-threshold the weights
y = conv(x, W_sparse)                              # forward pass uses the SPARSE weight
loss = criterion(y, target)
# loss.backward() + SGD(weight_decay=lambda) does the rest. Sparsity emerges.

Related on Engineermaxxing

Pair this with lessons on magnitude pruning, the Lottery Ticket Hypothesis, low-rank factorization (SVD), and quantization to see the full model-compression toolbox. STR is the "learn the budget" corner of that map.