[Paper Review] Dynamic Model Pruning with Feedback

Tao Lin, Sebastian U. Stich+3 more|arXiv (Cornell University)|Jun 12, 2020

Advanced Neural Network ApplicationsComputer Science47 references87 citations

TL;DR

DPF introduces dynamic pruning with error feedback to train sparse networks in one pass, jointly maintaining a dense model to correct pruning errors, achieving state-of-the-art results on CIFAR-10 and ImageNet without retraining the sparse model.

ABSTRACT

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.

Motivation & Objective

• Motivate model compression to enable deployment on low-end devices without excessive retraining.
• Develop a dynamic pruning method that jointly maintains a dense and a pruned model to correct pruning errors.
• Obtain sparse networks with state-of-the-art accuracy on CIFAR-10 and ImageNet.
• Provide theoretical convergence analysis for convex and non-convex objectives under the proposed scheme.

Proposed method

• Propose Dynamic Pruning with Feedback (DPF) where the gradient is computed on the pruned weights but applied to the full dense weight vector.
• Use an error feedback mechanism to allow reactivation of previously pruned weights during training.
• Model pruning is performed with a mask m_t applied to w_t, yielding tilde{w}_t = m_t  w_t, and the update w_{t+1} = w_t - gamma_t g(m_t w_t).
• Show equivalence w_{t+1} = w_t - gamma_t g(w_t + e_t) with e_t = tilde{w}_t - w_t, linking to error-feedback SGD.
• Provide convergence guarantees for both convex (strongly convex) and non-convex objectives under assumptions of smoothness and bounded stochastic gradients.
• Extend the analysis to other compressors beyond masking, e.g., quantization, within the same delta_t framework.

Experimental results

Research questions

• RQ1 Can dynamic pruning with error feedback achieve high sparsity without significant loss in accuracy on large-scale datasets?
• RQ2 Does maintaining a simultaneous dense model and applying gradients computed on the pruned model improve generalization and enable recovery from premature pruning?
• RQ3 What are the convergence properties of DPF on convex and non-convex objectives under standard optimization assumptions?
• RQ4 How does DPF compare to existing pruning schemes (one-shot, incremental, dynamic) across datasets like CIFAR-10 and ImageNet?

Key findings

• DPF achieves state-of-the-art performance in sparsity-aware pruning, outperforming prior dynamic and static pruning methods on CIFAR-10 and ImageNet.
• DPF enables high sparsity (up to around 99% in reported experiments) while maintaining viable accuracy on large models.
• The method shows that pruning masks converge during training, with only a small fraction of weights still changing near the end.
• Compared to lottery ticket-style pruning, DPF finds effective sparse networks without requiring training from scratch or extensive fine-tuning, and performs especially well at very high sparsity levels.
• Theoretical results provide non-asymptotic convergence guarantees for both convex and non-convex objectives under the proposed error-feedback pruning scheme.
• Experiments demonstrate consistent gains over SNIP, Incremental pruning, DSR, and Sparse Momentum across CIFAR-10 and ImageNet.