Fast and Controllable Post-training Sparsity: Learning Optimal Sparsity Allocation with Global Constraint in Minutes

TL;DR

The paper tackles efficient post-training sparsity by learning a global sparsity allocation across all layers through a differentiable, KDE-based bridge between pruning thresholds and per-layer sparsity. It combines a reconstruction loss $L_{rec}$ with a global-sparsity control loss $L_c$, yielding a total objective $L = L_{rec} + L_c$ and enabling net-wide optimization via a learnable threshold $t_l$ for each layer. The method achieves substantial accuracy gains at high sparsity (e.g., over 30% improvement for ResNet-50 on ImageNet at 80% sparsity) and completes the sparse-model reconstruction in minutes on standard hardware, outperforming existing post-training sparsity and retraining-based approaches. FCPTS also demonstrates broad applicability to CNNs and ViTs across CIFAR-10/100, ImageNet, and PASCAL VOC, with practical benefits in inference latency and memory on specialized hardware.

Abstract

Neural network sparsity has attracted many research interests due to its similarity to biological schemes and high energy efficiency. However, existing methods depend on long-time training or fine-tuning, which prevents large-scale applications. Recently, some works focusing on post-training sparsity (PTS) have emerged. They get rid of the high training cost but usually suffer from distinct accuracy degradation due to neglect of the reasonable sparsity rate at each layer. Previous methods for finding sparsity rates mainly focus on the training-aware scenario, which usually fails to converge stably under the PTS setting with limited data and much less training cost. In this paper, we propose a fast and controllable post-training sparsity (FCPTS) framework. By incorporating a differentiable bridge function and a controllable optimization objective, our method allows for rapid and accurate sparsity allocation learning in minutes, with the added assurance of convergence to a predetermined global sparsity rate. Equipped with these techniques, we can surpass the state-of-the-art methods by a large margin, e.g., over 30\% improvement for ResNet-50 on ImageNet under the sparsity rate of 80\%. Our plug-and-play code and supplementary materials are open-sourced at https://github.com/ModelTC/FCPTS.

Figures (5)

• Figure 1: Comparison with existing post-training sparsity methods on ImageNet ResNet-50. Our FCPTS enjoys significant accuracy improvement, especially for the extremely high sparsity rates (e.g., 30% boost under 80% sparsity).
• Figure 2: An overview of our fast and controllable post-training sparsity (FCPTS) framework. The differentiable sparsity allocation transfers the learning of the sparsity rate to the threshold by a differentiable estimation, and the controllable sparsity reconstruction enables an optimized network with a specified global sparsity rate. Both components contribute to the final excellent performance.
• Figure 3: Explanation of bridge function $g^{-1}$ from threshold $t_l$ to sparsity $r_l$ (Equation \ref{['eq:bridge_function']}). The area of the shading region equals the sparsity rate $r_l$. The derivative of $r$ with respect to $t$ can be represented as $p(-t_l) + p(t_l)$.
• Figure 4: The effect of learnable sparsity. For each neural network, the left dash bars are results with fixed sparsity rate and the right bars are results with learned sparsity rate.
• Figure 5: Visualization of the optimized sparsity allocation at a sparsity rate of 90%. ResNet-18 and ResNet-50 are on ImageNet, ResNet-32 and ResNet-56 are on CIFAR-100.