[Paper Review] DSA: More Efficient Budgeted Pruning via Differentiable Sparsity Allocation

Xuefei Ning, Tianchen Zhao+4 more|arXiv (Cornell University)|Apr 5, 2020

Advanced Image and Video Retrieval Techniques25 references9 citations

TL;DR

DSA proposes a differentiable sparsity allocation method for efficient budgeted pruning, enabling end-to-end, pruning-from-scratch optimization via gradient-based layer-wise sparsity allocation. It achieves state-of-the-art performance with at least 1.5× speedup over iterative pruning methods on CIFAR-10 and ImageNet by using a novel differentiable pruning process and ADMM-based optimization under FLOPs constraints.

ABSTRACT

Budgeted pruning is the problem of pruning under resource constraints. In budgeted pruning, how to distribute the resources across layers (i.e., sparsity allocation) is the key problem. Traditional methods solve it by discretely searching for the layer-wise pruning ratios, which lacks efficiency. In this paper, we propose Differentiable Sparsity Allocation (DSA), an efficient end-to-end budgeted pruning flow. Utilizing a novel differentiable pruning process, DSA finds the layer-wise pruning ratios with gradient-based optimization. It allocates sparsity in continuous space, which is more efficient than methods based on discrete evaluation and search. Furthermore, DSA could work in a pruning-from-scratch manner, whereas traditional budgeted pruning methods are applied to pre-trained models. Experimental results on CIFAR-10 and ImageNet show that DSA could achieve superior performance than current iterative budgeted pruning methods, and shorten the time cost of the overall pruning process by at least 1.5x in the meantime.

Motivation & Objective

• Address the inefficiency of traditional iterative budgeted pruning methods that rely on discrete, search-based sparsity allocation.
• Eliminate the need for pre-training by enabling pruning-from-scratch optimization.
• Develop a differentiable pruning process that allows gradient computation w.r.t. pruning ratios for continuous optimization.
• Improve sample efficiency and computational efficiency in sparsity allocation under FLOPs or latency constraints.
• Ensure structural integrity by handling topological constraints from skip connections via a grouping procedure.

Proposed method

• Introduce a differentiable pruning process that softens hard pruning using probability distributions over channel masks, enabling gradient flow through pruning ratios.
• Use validation loss as a differentiable surrogate for accuracy to enable gradient-based optimization of sparsity allocation.
• Apply the Alternating Direction Method of Multipliers (ADMM) to jointly optimize weights and pruning ratios under budget constraints.
• Integrate ADMM updates into the main training loop, updating sparsity allocation every tens of weight optimization steps to reduce overhead.
• Implement a topological grouping strategy to preserve connectivity in residual networks by treating skip-connection groups as units during pruning.
• Normalize and align gradient magnitudes w.r.t. pruning ratios with layer-wise sensitivity to guide rational sparsity allocation.

Experimental results

Research questions

• RQ1 Can gradient-based optimization over continuous sparsity ratios outperform discrete search in budgeted pruning?
• RQ2 Does a differentiable pruning process enable effective and efficient sparsity allocation without pre-training?
• RQ3 How well do the gradients of the task loss w.r.t. pruning ratios align with actual layer-wise sensitivity?
• RQ4 Can DSA achieve superior accuracy and faster convergence compared to iterative pruning baselines on CIFAR-10 and ImageNet?
• RQ5 Does the pruning-from-scratch approach of DSA eliminate the need for costly pre-training while maintaining performance?

Key findings

• DSA achieves a 1.11% accuracy drop on ImageNet when keeping 60% of the FLOPs for ResNet-18, outperforming all iterative pruning baselines.
• DSA reduces the total pruning time by at least 1.5× compared to iterative methods, achieving 5/300 GPU hours on CIFAR-10 and 150/300 GPU hours on ImageNet.
• The gradient magnitudes w.r.t. pruning ratios align well with empirical layer-wise sensitivity, validating the method’s rationality in sparsity allocation.
• DSA successfully preserves model connectivity by grouping topologically related layers (e.g., residual blocks), maintaining structural integrity.
• The pruning-from-scratch approach in DSA eliminates the need for pre-training, reducing total computational cost by avoiding the 300-epoch pre-training stage.
• On CIFAR-10, DSA consistently outperforms baselines like AMC and SSL across multiple FLOPs budgets, demonstrating robustness and efficiency.