Weight Concentration Regularization for Improving Pruning Robustness Under High Sparsity

TL;DR

This paper introduces a Weight Concentration Regularizer (WCR) that enhances pruning robustness by concentrating weight energy on a few parameters, improving model performance under high sparsity.

Contribution

The paper proposes WCR, a novel regularizer that amplifies key weights during training, leading to better pruning robustness and compatibility with existing optimizers.

Findings

WCR improves pruning robustness across tasks and architectures.

WCR is compatible with existing pruning-robust optimizers.

WCR demonstrates consistent performance gains in various experiments.

Abstract

Deep neural networks achieve outstanding performance across vision and language tasks, yet their large parameter counts limit deployment in resource-constrained settings. One-shot pruning reduces model size without retraining, but models trained with standard objectives often suffer substantial accuracy drops under aggressive sparsity. Prior work mitigates this drop along two directions: regularizers such as $\ell_1$ and DeepHoyer that shape the weight distribution during training, and pruning-robust optimizers such as SAM, CrAM, and S$^2$SAM that flatten the loss landscape. However, existing regularizers either shrink all weights uniformly ($\ell_1$) or induce scale-invariant sparsity (DeepHoyer), without concentrating weight energy onto a small set of informative parameters. We propose a Weight Concentration Regularizer (WCR), a training-time regularizer that amplifies the magnitude of a small subset of parameters while driving the remainder toward zero, so that magnitude pruning predominantly removes parameters with negligible functional contribution. We provide a convergence analysis and evaluate WCR on LLM fine-tuning, image classification, and medical segmentation, demonstrating consistent improvements in pruning robustness across architectures and compatibility with existing pruning-robust optimizers.