Source-linked AI summary
Dynamic Model Pruning with Feedback
Tao Lin, Sebastian U. Stich, Luis Barba, Daniil Dmitriev, Martin Jaggi
TL;DR
Large neural networks strain the memory and computational resources of low-end devices, motivating sparse models that retain accuracy without costly retraining. The paper introduces dynamic pruning with feedback, which jointly maintains sparse and dense models, dynamically changes masks, and can reactivate prematurely pruned weights during one training pass. Across CIFAR-10 and ImageNet, the method reports state-of-the-art sparse-model performance and outperforms prior pruning schemes, while its analysis shows convergence only to a neighborhood when compression error persists.
Problem
Large neural networks demand substantial memory and inference computation, while pruning methods face trade-offs involving retraining, mask selection, or hyperparameter tuning.
Method
DPF dynamically proposes sparsity masks, evaluates gradients at the pruned model, and applies them to a jointly maintained dense model that provides feedback for pruning errors.
Results
DPF reports state-of-the-art accuracy and sparsity, outperforming prior pruning schemes across evaluations including CIFAR-10 and structured filter pruning.
Takeaways & Limitations
DPF produces trained sparse models in one pass without mandatory retraining, with trivial dynamic-reparameterization overhead and inference performed only on sparse models.
Takeaways & Limitations
When compression error does not vanish, DPF converges only to a neighborhood of a stationary or optimal solution because dense optima may not be well approximated by sparse models.
Abstract
from arXiv · showhide
Deep neural networks often have millions of parameters. This can hinder their deployment to low-end devices, not only due to high memory requirements but also because of increased latency at inference. We propose a novel model compression method that generates a sparse trained model without additional overhead: by allowing (i) dynamic allocation of the sparsity pattern and (ii) incorporating feedback signal to reactivate prematurely pruned weights we obtain a performant sparse model in one single training pass (retraining is not needed, but can further improve the performance). We evaluate our method on CIFAR-10 and ImageNet, and show that the obtained sparse models can reach the state-of-the-art performance of dense models. Moreover, their performance surpasses that of models generated by all previously proposed pruning schemes.
1 INTRODUCTION
Pruning reduces deployment costs by removing weights, but existing approaches trade computational expense, hyperparameter tuning, or retraining against performance. The paper proposes dynamic pruning with feedback to train sparse models in one pass while maintaining strong accuracy.
- Pruning removes weights to reduce memory and inference-resource demands while often preserving dense-model test accuracy.Compressed networks can sometimes generalize better than full models.
- One-shot pruning reduces model size with little accuracy degradation but requires dense training, pruning, and refinement, and is outperformed by mask-exploring methods.Dynamic methods instead readjust sparsity patterns during training but require fine-tuning many hyperparameters.
- The proposed approach identifies important weights throughout training, jointly maintains sparse and dense models, and uses the dense model to correct pruning errors.The method targets better generalization without additional hyperparameter tuning or mandatory retraining.
- The paper claims state-of-the-art accuracy and sparsity, outperforming previously proposed pruning schemes while finding a trained sparse model in one pass.An ablation study and convergence analysis complement the empirical results.
2 RELATED WORK
Prior pruning methods differ by when pruning occurs and how masks are selected. They include post-training pipelines, dynamic during-training schemes, initialization-time pruning, and learned gating or Bayesian approaches.
- Pruning after training: Post-training pruning typically trains a dense model, applies one-shot pruning, and fine-tunes the fixed sparse model.This three-stage pipeline became standard after achieving moderate sparsity with minor quality loss.
- Pruning during training: During-training methods gradually increase sparsity, prune filters, or periodically regrow connections to adapt the mask.Examples include magnitude-based pruning, SFP, and Deep Rewiring.
- Pruning during training: Dynamic Sparse Reparameterization and Sparse Momentum redistribute layer sparsity through prune–redistribute–regrowth cycles using gradients or momentum.Sparse Momentum slightly outperforms DSR on ImageNet but shows no CIFAR difference.
- Pruning before training: Pruning-before-training methods select a mask at initialization and keep it fixed, while lottery-ticket procedures rely on repeated full training passes.SNIP selects masks using connection sensitivities for a given task.
- Further approaches: Other approaches learn gating variables or dropout probabilities to reduce nonzero weights, including ℓ0-regularized and Bayesian methods.
3 METHOD
DPF trains a sparse model by computing gradients at the currently pruned weights while updating a jointly maintained dense model. Dynamic masks and feedback allow previously pruned weights to reactivate during training.
- Update rule: SGD uses unbiased stochastic gradients of a non-convex loss, while DPF evaluates the update at the compressed model rather than the dense iterate.The stochastic gradient may represent a mini-batch gradient.
- Sparse model construction: A pruning mask m produces the sparse model ewt := m ⊙ wt by setting selected weights of wt to zero.The mask may depend on weights, such as smallest-magnitude pruning, or on training time.
- Pruning methodologies: Existing pruning categories differ in when masks are computed and applied, including before, after, or during training.
- Existing approaches: Pruning before training optimizes a fixed subnetwork, whereas one-shot pruning trains densely first and then requires fine-tuning because the sparse model is not generally at a local optimum.
- Dynamic pruning: Fully dynamic pruning can reactivate previously pruned weights, unlike incremental schemes that monotonically increase sparsity.
- Dynamic pruning with feedback: DPF evaluates a stochastic gradient at the pruned model ewt and applies it to the simultaneously maintained dense model wt.
- Feedback mechanism: Updating the full model lets accumulated gradients restore prematurely masked weights when later updates substantially change them.Equivalently, DPF applies gradients at wt + et, where et is the compression error.
4 CONVERGENCE ANALYSIS
The convergence analysis gives DPF guarantees under convex and non-convex objectives, while showing that pruning quality controls the residual error. The analysis also extends to arbitrary compressors and identifies conditions under which convergence reaches only a neighborhood of the optimum.
- Assumptions: Under smoothness and bounded stochastic-gradient assumptions, DPF's pruning quality is measured by δ_t = ∥w_t − ew_t∥_2 / ∥w_t∥_2.The guarantee assumes δ_t ∈ [0, 1] and bounds the stochastic gradients evaluated at pruned models.
- Convex objectives: For strongly convex objectives, a randomly selected pruned iterate of DPF satisfies an expected convergence guarantee under γ_t = 4 / [µ(t+2)].The selected iterate uses probability p_t = 2(t+1)/[(T+1)(T+2)].
- Convergence limitations: Unless δ_t → 0 or ∥w_t∥ → 0, the pruning error does not vanish, so DPF converges only to a neighborhood of the optimal solution.This limitation applies to both the pruned model and the jointly maintained dense model; a globally optimal dense model may not be well approximated by a sparse model.
- Convex objectives: The convergence bound depends sensitively on the final pruning quality δ_T and the norm ∥w_T∥, and includes a condition-number factor in the asymptotically decreasing term.The analysis reports a slight theoretical advantage over rounding at the end.
- Non-convex objectives: For non-convex objectives, the analysis establishes convergence to a neighborhood of a stationary point for a uniformly random pruned iterate.The result uses a learning rate proportional to 1/√T and takes expectation over stochasticity and iterate selection.
- General compression: The analysis applies beyond mask pruning because the scheme can use an arbitrary compressor whose quality is measured through δ_t parameters.The paper gives other compression schemes as examples of this extension.
5 EXPERIMENTS
Experiments compare DPF with established pruning methods across CIFAR-10 and ImageNet, architectures, and sparsity levels. DPF generally delivers stronger accuracy, including at high sparsity, while extra fine-tuning can further improve results.
- DPF shows consistent and noticeable performance benefits over competitors across neural architectures and sparsity levels.
- CIFAR-10: On CIFAR-10, DPF remains viable at very high sparsity, including 99.9% on WideResNet-28-8, while other methods may lose quality or fail to converge.
- CIFAR-10: DPF outperforms simple pruning techniques in almost all settings but faces difficulties on smaller models at extremely high sparsity.
- CIFAR-10: Extra fine-tuning enables DPF to convincingly outperform other methods regardless of network size.The reported extra fine-tuning uses the model checkpoint after normal training and 60 additional epochs with tuned optimizer and learning rate.
- CIFAR-10: DPF explores subnetworks more efficiently in larger parametric spaces than in smaller models.This conclusion comes from comparing predefined-parameter searches on models such as WideResNet-28-8 and WideResNet-28-2.
- ImageNet: On ImageNet, DPF significantly improves top-1 test accuracy at both 80% and 90% sparsity with fewer or equal parameters.The comparison uses ResNet-50 and includes dynamic reparameterization methods and the Incremental baseline.
6 DISCUSSION
DPF trains sparse models through dynamic mask updates and feedback, avoiding retraining while enabling mask exploration and competitive accuracy. The discussion also examines convergence, lottery-ticket behavior, structured pruning, and implementation choices.
- Efficiency: DPF adds only trivial computational overhead during training and performs inference using the sparse model, while fine-tuning remains optional.The paper reports that fine-tuning can further improve performance, although it is not required.
- Training dynamics: After epoch 157 at 95% sparsity, only 5% of mask elements were still changing, indicating early convergence followed by mask fine-tuning.The example uses WideResNet-28-2 on CIFAR-10 and reports continued accuracy improvement in the final epochs.
- Lottery-ticket effect: DPF masks are not lottery tickets: training the discovered sparse model from scratch does not recover the same performance, while DPF is much better at 99% sparsity.Figure 5 compares DPF with lottery-ticket training across target sparsity levels and reports lower computation for DPF.
- Structured pruning: DPF extends to structured filter pruning and outperforms the filter-norm-based SFP method on CIFAR-10, although the reported extension uses an ℓ2 filter-selection criterion.The authors leave exploration of stronger filter-selection schemes for future work.
- Method: DPF updates pruning masks during training and jointly maintains dense and sparse parameterizations, with mask updates triggered every p = 16 iterations.The procedure applies the mask for sparse forward/backward computation while updating the uncompressed weights.
A.2 IMPLEMENTATION DETAILS
The implementation evaluates DPF and pruning baselines on CIFAR-10 and ImageNet using standard image-classification datasets, gradual sparsity schedules, and tuned optimization settings.
- Implementation: Sparse tensors are represented as dense tensors multiplied by binary masks, and experiments run in PyTorch on NVIDIA Tesla V100 GPUs.The implementation description covers the computational representation used for sparse models.
- Datasets: Experiments use CIFAR-10 with 50K training and 10K test images across 10 classes, and ImageNet-1k with 1.28 million training and 50K validation images across 1K classes.Both datasets use standard preprocessing and augmentation schemes.
- Optimization: Learning rates are grid-searched, while most methods use mini-batch SGD with Nesterov momentum and fine-tuning baselines additionally consider Adam.The grid search expands when the best value lies at a boundary.
- Training schedules: The DPF mini-batch size is fixed at 128 for CIFAR-10 and 1024 for ImageNet across datasets and models.CIFAR-10 training uses architecture-specific epoch counts and learning-rate schedules, while ImageNet training follows a 90-epoch schedule.
A.3.1 COMPLETE RESULTS OF UNSTRUCTURED PRUNING ON CIFAR-10
This appendix section reports complete numerical results for training compressed state-of-the-art neural networks on CIFAR-10 from scratch under unstructured pruning.
- Results: Table 4 reports Top-1 test accuracy for compressed state-of-the-art DNNs trained from scratch on CIFAR-10.The results are averaged over three runs and use unstructured pruning.
- Results: The table marks methods that cannot converge with a star, making convergence failures explicit in the numerical comparison.The reported results reconstruct results shown in Table 1 and Figure 3.
A.3.2 UNDERSTANDING THE TRAINING DYNAMICS AND LOTTERY TICKET EFFECT
The appendix analyzes DPF training dynamics and lottery-ticket behavior, emphasizing mask evolution, structural exploration, and the effects of training budgets.
- Training dynamics: Figures 7 and 8 complement Figure 4 by detailing DPF training dynamics, including convergence of δ and masks and IoU between consecutive masks.For IoU, smaller values indicate that a larger fraction of mask elements flip.
- Training dynamics: Figure 9 compares DPF with Incremental under the same gradual warmup schedule and hyperparameters, highlighting DPF’s drastic reparameterization over dense parameter space.The comparison concerns mask convergence at different target sparsity levels.
- Lottery-ticket effect: Figure 10 studies the lottery-ticket hypothesis under equal-epoch and equal-total-FLOP training budgets.The appendix uses these results to examine training-time structural exploration and implicit regularization effects.
- Lottery-ticket effect: DPF provides an alternative compression training scheme for extremely high compression ratios without sacrificing test accuracy, where several existing methods incur severe quality loss.The paper explicitly avoids questioning the importance of initialization or the existence of lottery tickets.
A.3.3 COMPUTATIONAL OVERHEAD AND THE IMPACT OF HYPER-PARAMETERS
This section examines DPF’s computational overhead and hyper-parameter choices, including reparameterization frequency, pruning scope, and mini-batch size. It also evaluates these choices across CIFAR-10 and ImageNet settings.
- Retraining comparison: Figure 10 compares lottery-ticket retraining with the same FLOPs budget for unstructured pruning of WideResNet28-2 on CIFAR-10.
- Hyper-parameter impact: Reparameterizing every p = 16 iterations is consistently optimal across tested sparsity ratios and mini-batch sizes.This setting was used in all experiments.
- Pruning scheme: Global-wise unstructured pruning enables more flexible dynamic parameter reallocation and performs better especially at aggressive sparsity ratios.For the same compressed parameter count, it slightly increases MACs compared with layer-wise pruning.
- Hyper-parameter impact: Figure 11 evaluates test top-1 accuracy under different reparameterization periods, sparsity ratios, mini-batch sizes, and layer-wise pruning choices.
- Computational overhead: Table 5 compares MACs between layer-wise and non-layer-wise reparameterization schemes at the same compressed parameter count.
- Computational overhead: DPF introduces trivial computational overhead when training a compressed ResNet-50 from scratch on ImageNet.The evaluation includes reparameterization cost, while practical inference speedup from high sparsity was not estimated.
A.3.4 IMPLICIT NEURAL ARCHITECTURE SEARCH
The experiments show that DPF can explore structural configurations during training, functioning as an implicit neural architecture search. Larger source networks can yield better generalization at the same compressed size, while structured pruning learns layerwise sparsity transitions.
- Implicit neural architecture search: At the same pruned model size, architecture search from WideResNet-28-8 produces better generalization than search from WideResNet-28-4.
- Implicit neural architecture search: Figure 13 compares test top-1 accuracy with compressed model size for WideResNet-28 models of different widths on CIFAR-10.
- Structured pruning: Figure 14 reports the relationship between parameter count and top-1 test accuracy for structured filter-wise pruning, averaging results over three runs.
- Structured pruning: Table 6 evaluates DPF and baseline methods on WideResNet variants using norm-based filter selection and gradual pruning warmup.
- Lottery-ticket analysis: For structured pruning, DPF does not find a lottery ticket, so its superior generalization is not explained by the mask or weight initialization scheme.
- Sparsity transitions: Under structured pruning, DPF gradually learns to prune entire residual blocks because of residual connections.
B MISSING PROOFS
This section supplies proofs for DPF’s convergence claims. The analysis combines smoothness inequalities with error-feedback bounds, weighted telescoping sums, and comparisons to standard convergence results.
- Proof strategy: The strongly convex proof combines standard SGD analysis with error-feedback estimates to control pruning errors.
- Theorem 4.1: DPF updates the iterate as w_t+1 = w_t − γ_t g(e_wt), which forms the starting point for the convergence analysis.
- Theorem 4.1: L-smoothness and an inner-product bound produce an inequality involving the pruning error norm and the objective gap.
- Theorem 4.1: Choosing γ_t = 4/[μ(t+2)] and λ_t = t+1 enables the proof’s weighted convergence bound.
- Theorem 4.1: Summing the λ_t-weighted inequalities from t = 0 to t = T yields a telescoping sum.
- Final-iterate pruning: The analysis justifies one-shot pruning of the final iterate and identifies a slight theoretical advantage of DPF over rounding at the end.
- Theorem 4.2: The dense model is also shown to converge to a neighborhood of an optimal solution using L-smoothness and the preceding estimate.
- Theorem 4.2: The proof of Theorem 4.2 follows the error-feedback analysis of Karimireddy et al. (2019).