Dynamic Model Pruning with Feedback

In plain words

A network with millions of weights is slow and large to run on small devices. Pruning removes weights, leaving a sparse model. Most methods either prune a finished model and then repair it, or prune during training and never bring back a weight they cut. If they cut the wrong weight early, the damage stays.

This method keeps two versions of the weights during training: the full set and a sparse copy chosen by a simple rule (keep the largest-magnitude weights). The sparse copy is the one used to compute each training step, so the model learns to work with the cuts. The update, however, goes to the full set. A weight that was cut by mistake keeps receiving updates, and if it grows large enough it comes back into the sparse copy. The sparse model is ready at the end of the one training run, with no separate repair stage.

TL;DR

• DPF keeps a dense weight vector x and a sparse model m ⊙ x. The mask m is recomputed by global magnitude pruning every 16 iterations alg. 1 §5.1.
• The gradient is evaluated at the sparse model and applied to the dense weights §3 eq. DPF.
• The paper reports sparse accuracy close to dense baselines on CIFAR-10 and ImageNet, with few extra hyperparameters (the paper's claim) and no retraining step §1 §5.2.
• Author-reported gains: WideResNet-28-2 at 99% sparsity, 88.92% vs 83.45% (SNIP) on CIFAR-10 tab. 2. ResNet-50 at 80% sparsity on ImageNet, −0.47 points top-1 vs dense tab. 4.
• No official code was found. Several details (schedule units, tie handling, ImageNet schedule) are not stated; see Pitfalls.

Problem

• One-shot pruning of a trained model needs a dense training run, a pruning step and fine-tuning, and it commits to one mask §1 §3 fig. 2.
• Dynamic methods often have many hyperparameters that must be tuned §1 §5.1.
• Pruning before training fixes the mask at initialisation. A wrong early choice strongly affects performance §3.
• Incremental pruning can only remove weights; it cannot recover from premature pruning §3 fig. 1.

Method

The method is a dynamic mask over the dense weights with feedback from the sparse model §3.

Notation. x_t are the dense weights and m_t a binary mask. The sparse model is x̂_t = m_t ⊙ x_t §3. The error of the compression is e_t = x̂_t − x_t. The sparsity target at epoch t follows a cubic ramp from zero to the final sparsity app. A.2.

Update. The gradient g is a stochastic gradient of the loss. It is evaluated at the pruned weights, and the step is applied to the dense weights §3:

x_{t+1} := x_t - \gamma_t\, g(m_t \odot x_t) = x_t - \gamma_t\, g(\hat x_t) \tag{DPF}

eq. DPF §3

Equivalently, the step is a gradient at the dense point plus the error, x_{t+1} = x_t − γ_t g(x_t + e_t). This links the method to error feedback, as in Stich et al. (2018) and Karimireddy et al. (2019) §3. The paper describes the method as an instance of perturbed-gradient updates to the dense model §3.

Procedure (Algorithm 1, alg. 1, with the training schedule from app. A.2 and the setup from §5.1):

1. Set the sparsity for the current epoch from the cubic schedule app. A.2.
2. Every p = 16 iterations, recompute the mask from the dense weights by magnitude, at that sparsity alg. 1 §5.1.
3. Form the pruned weights x̂ = m ⊙ x alg. 1.
4. Compute the mini-batch gradient at x̂ alg. 1.
5. Apply the gradient to the dense weights x with a general optimiser; the paper's runs use SGD with Nesterov momentum 0.9 alg. 1 §5.1.
6. Output the dense weights x_T and the sparse model x̂_T. The sparse model is the one used for inference alg. 1 §6.

# Dynamic pruning with feedback (DPF)  [alg. 1]
# Input:  dense weights x, mask period p (= 16), iterations T
# Output: dense weights x and sparse model x_hat = mask * x

def train_dpf(x, T, p):
    # Mask is recomputed from the dense weights every p iterations
    # [alg. 1] [§5.1]
    for t in range(T + 1):
        if t % p == 0:
            # Sparsity for the current epoch; the mapping from
            # iterations to epochs is not stated [app. A.2]
            s = sparsity_at(epoch_of(t))
            mask = magnitude_mask(x, s)
        # Pruned model used for the forward and backward pass
        # [alg. 1]
        x_hat = mask * x
        # Gradient is taken at the pruned weights [§3] [eq. DPF]
        grad = stochastic_gradient(x_hat)
        # Dense weights get the update from any optimiser, e.g. SGD
        # with momentum [alg. 1] [§5.1]
        x = optimizer_step(x, grad)
    return x, mask * x

# Global magnitude mask over the conv weights  [§5.1] [app. A.1]
def magnitude_mask(x, sparsity):
    # Last FC layer, biases and BatchNorm stay dense  [§5.1] [app. A.1]
    conv = conv_weights(x)
    # Count of conv weights to zero; rounding is not stated  [app. A.1]
    n_zero = sparsity * conv.size
    # One threshold over all conv layers; tie handling is not stated
    threshold = kth_smallest(abs(conv), n_zero)
    # Keep the weights whose magnitude is above the threshold
    return keep_above(x, threshold)

# Sparsity schedule: cubic ramp from 0 to the final target  [app. A.2]
def sparsity_at(epoch, s_final, n):
    # n: epoch of the second learning-rate decay; t0 = 0, dt = 1 epoch
    # [app. A.2]
    progress = epoch / n
    # Clamping for epochs beyond n is not stated  [app. A.2]
    return s_final * (1 - (1 - progress) ** 3)

Structured extension. The paper also applies the same scheme to filter pruning, using an ℓ2-norm filter criterion. The threshold is applied across layers, and the paper says this is preliminary and only for CIFAR-10 §6 tab. 7.

Quality measure (for the theory). The per-iterate quality of pruning is δ_t = ||x_t − x̂_t||² / ||x_t||² eq. 1 §4.

Hyperparameters

Results

Author-reported. Setup: CIFAR-10 (50K/10K) and ImageNet (1.28M/50K) image classification; ResNet, WideResNet, VGG; training from scratch with the schedules above; baselines SNIP, SM, DSR, One-shot+FT, Incremental (Zhu & Gupta) and the lottery ticket procedure §5.1 tab. 1 §5.2.

Table 2: CIFAR-10, top-1 test accuracy (mean ± std), unstructured pruning, from scratch tab. 2 tab. 5.

⋆ = the method did not converge; "-" = not reported tab. 2.

Additional DPF-only rows from the complete CIFAR-10 table tab. 5: WRN-28-2 at 50%, 60%, 70%, 80%: 95.11, 94.90, 94.86, 94.76. WRN-28-4 at 70%, 80%, 90%: 95.58, 95.60, 95.65. WRN-28-8 at 70% and 80%: not reported ("-").

Table 3: CIFAR-10 with extra fine-tuning tab. 3. "Incremental" and "DPF" are trained without extra fine-tuning; "+FT" adds 60 epochs.

Table 4: ImageNet, ResNet-50, top-1 and top-5 accuracy, unstructured pruning tab. 4. DSR numbers are taken from [mostafa2019parameter], not re-run. The Incremental row with Gale et al. uses layer-wise sparsity and a tuned schedule.

Other reported results §5.2 §6 app. A.3 app. A.4.

• Lottery ticket: retraining DPF's final sparse mask from scratch does not recover its accuracy; the paper concludes that the mask is not a lottery ticket §6 fig. 5. DPF matches the expensive iterative lottery-ticket procedure at mild sparsity and is better at 99% fig. 5.
• Mask convergence: for 95% sparsity on WRN-28-2 after epoch 157 (43 epochs before the end), only 5% of mask elements still change. Most masks settle early, and late training mainly fine-tunes the mask §6 fig. 4.
• Structured pruning on CIFAR-10 (ℓ2 filter criterion), DPF vs SFP, WRN-28-8 at 80%: 95.15 ± 0.03 vs 94.22 ± 0.21 tab. 7. The full structured table is below.

Table 7: CIFAR-10, structured (filter) pruning, DPF vs SFP tab. 7.

Compute overhead, for global versus layer-wise pruning on ResNet-20 (MACs, same parameter count) tab. 6: at 50% target, 22.80 M vs 20.99 M; at 90%, 6.54 M vs 4.67 M. The paper says the mask update adds trivial overhead for ResNet-50 on ImageNet fig. 12, but gives no wall-clock measurements app. A.3.3.

Paper vs code

No official code.

Searched: the paper text, the LaTeX source and the arXiv record for code links (none; the only mention is that baselines use their open-sourced code §5.1); a GitHub repository search for "dynamic model pruning with feedback" and "DPF pruning epfl" (one hit, a third-party PyTorch reimplementation with no license, owner not an author, not used) and "dpf org:epfml" (no hits).

Pitfalls

• Schedule units are ambiguous. The formula uses t and Δt = 1 epoch app. A.2, while Algorithm 1 updates the mask every 16 iterations and uses iteration t alg. 1. How the epoch schedule maps onto iterations is not stated.
• The ImageNet ramp length is not stated for DPF. For CIFAR-10 the ramp ends at the second LR decay app. A.2. For ImageNet the same rule would give 60 epochs, but the paper does not say so for DPF.
• Ties, rounding and clamping are not stated. Threshold ties, the rounding of the count of zeroed weights and the clamping of the cubic schedule past n are not described.
• "No additional hyperparameters" needs care. The paper says DPF avoids extra hyperparameters §1 §5.1, but the learning rate is grid-searched for DPF app. A.2, and the ImageNet warm-up is given as 0.1 to 0.4 without the warm-up length app. A.2.
• Theory vs practice. The convergence theorems use plain SGD with a specific step size, smooth and bounded-gradient assumptions, and a random iterate §4 app. B. The experiments use Nesterov momentum with step decays §5.1. Theorem constants (G, L, μ) are not measured. The theorem's step size for the non-convex case depends on f(x_0) − f*, which is unknown in practice.
• One-shot comparison. The paper argues the final-iterate bound for one-shot pruning is weaker §4. That is a bound, not a measured gap.
• ImageNet comparison is not like-for-like. Against Gale et al.'s Incremental (same 5.15 M remaining parameters), DPF is lower in absolute pruned top-1 (75.13 vs 75.58) and has a smaller drop (−0.82 vs −1.11), because the dense baselines differ (75.95 vs 76.69) tab. 4. The 80% DPF rows also reach different sparsities (73.5% and 79.9%) tab. 4. The paper's ImageNet claim of a "significant improvement" §5.2 rests on the relative drop from the ZG baseline.
• Baselines are not all re-run. DSR's ImageNet numbers come from another paper's run, so their dense baseline (74.90) differs from the others (75.95) tab. 4.
• Results without fine-tuning. Tables 2 and 4 use DPF without extra fine-tuning. The paper says dynamic methods benefit from fine-tuning, so comparisons with fine-tuned baselines need Table 3 tab. 3 §5.2.
• Small models at extreme sparsity. The paper notes DPF has difficulties with ResNet-20 at 95% sparsity, and attributes this partly to fine-tuning; with 60 extra epochs it reaches 90.34 §5.2 tab. 3.
• Baseline non-convergence. DSR does not converge at 99% on WRN-28-2 and at 95% on ResNet-20 tab. 2 fig. 3.
• Global pruning costs MACs. For the same parameter count, global pruning gives more MACs than layer-wise pruning tab. 6. The paper still prefers global pruning for aggressive sparsity app. A.3.3.
• Sparse storage is dense. The implementation stores sparse tensors as dense tensors times masks app. A.2. The paper did not measure speedups app. A.3.3.
• The mask is not a lottery ticket. Do not use DPF's final mask as a fixed ticket for a fresh run; the paper's retraining experiment does not recover the accuracy §6 fig. 5.
• Structured pruning is preliminary. Only CIFAR-10, with an ℓ2 filter criterion and a threshold applied across layers; the paper says a better filter selection could improve results §6 tab. 7.

Lineage

• Builds on: the lottery ticket hypothesis (Frankle & Carbin 2018, "LT" in the paper) §2; the gradual pruning schedule and Incremental baseline (ZG, Zhu & Gupta) §5.1 app. A.2; error feedback (Stich et al. 2018; Karimireddy et al. 2019) §3.
• Compared against: SNIP (LAT), one-shot pruning with fine-tuning (HPDT), Incremental (ZG), DSR (MW), Sparse Momentum (DZ), and the iterative lottery-ticket procedure (FDRC) tab. 1 §5.1.
• Related card, not discussed in this paper: Rigging the Lottery (drop-and-grow sparse training).