[Paper Review] Effective Sparsification of Neural Networks with Global Sparsity Constraint

Xiao Zhou, Weizhong Zhang+2 more|arXiv (Cornell University)|May 3, 2021

Advanced Neural Network Applications41 references4 citations

TL;DR

This paper proposes ProbMask, a novel network sparsification method that uses probability as a global criterion to automatically determine weight importance and sparsity levels across all layers, avoiding manual pruning rate tuning. By formulating pruning as a continuous optimization problem over probability space with a global sparsity constraint and using the Gumbel-Softmax trick, ProbMask achieves state-of-the-art accuracy, especially at high sparsity rates—outperforming prior methods by up to 10% in Top-1 accuracy on ImageNet.

ABSTRACT

Weight pruning is an effective technique to reduce the model size and inference time for deep neural networks in real-world deployments. However, since magnitudes and relative importance of weights are very different for different layers of a neural network, existing methods rely on either manual tuning or handcrafted heuristic rules to find appropriate pruning rates individually for each layer. This approach generally leads to suboptimal performance. In this paper, by directly working on the probability space, we propose an effective network sparsification method called {\it probabilistic masking} (ProbMask), which solves a natural sparsification formulation under global sparsity constraint. The key idea is to use probability as a global criterion for all layers to measure the weight importance. An appealing feature of ProbMask is that the amounts of weight redundancy can be learned automatically via our constraint and thus we avoid the problem of tuning pruning rates individually for different layers in a network. Extensive experimental results on CIFAR-10/100 and ImageNet demonstrate that our method is highly effective, and can outperform previous state-of-the-art methods by a significant margin, especially in the high pruning rate situation. Notably, the gap of Top-1 accuracy between our ProbMask and existing methods can be up to 10\%. As a by-product, we show ProbMask is also highly effective in identifying supermasks, which are subnetworks with high performance in a randomly weighted dense neural network.

Motivation & Objective

• To address the challenge of selecting optimal, layer-specific pruning rates in neural network sparsification, which is difficult and suboptimal when done manually.
• To develop a global, comparable criterion for weight importance that is consistent across all layers, overcoming limitations of magnitude-based or normalized score-based methods.
• To eliminate the need for manual tuning of pruning rates per layer by learning redundancy budgets automatically through a global sparsity constraint in probability space.
• To resolve the training-testing performance discrepancy observed in prior pruning methods, especially on large-scale datasets like ImageNet.
• To effectively identify high-performing subnetworks (supermasks) in randomly initialized networks.

Proposed method

• Formulates network pruning as a continuous optimization problem over the probability space of binary masks, where each mask element is a Bernoulli random variable with learnable probability.
• Introduces a global sparsity constraint that limits the sum of all probabilities to control overall model sparsity, enabling automatic redundancy detection per layer.
• Uses the Gumbel-Softmax reparameterization trick to differentially train the probability parameters, allowing backpropagation through discrete mask sampling.
• Applies an annealing schedule for the temperature in Gumbel-Softmax to encourage convergence of probabilities to 0 or 1, forming a deterministic sparse mask.
• Leverages ℓ₁ regularization on the probability space to promote sparsity and improve convergence to low-variance, deterministic masks.
• Treats the final mask as a sparse, binary network, enabling efficient inference without retraining.

Experimental results

Research questions

• RQ1 Can probability serve as a globally comparable measure of weight importance across different layers in a deep neural network?
• RQ2 Does a global sparsity constraint in the probability space outperform layer-wise sparsity constraints in terms of accuracy, especially at high pruning rates?
• RQ3 Can the proposed method automatically learn appropriate sparsity budgets per layer without manual tuning?
• RQ4 Does ProbMask mitigate the training-testing performance discrepancy observed in prior pruning methods on large-scale datasets?
• RQ5 Can ProbMask effectively identify high-accuracy supermasks in randomly initialized dense networks?

Key findings

• ProbMask achieves a Top-1 accuracy of 74.68% on ImageNet with ResNet50 at a 99.9% pruning rate, outperforming previous state-of-the-art methods by up to 10 percentage points.
• On CIFAR-100 with ResNet32, ProbMask maintains high accuracy even at 99.9% pruning, while PBW and MLPrune collapse, indicating superior global comparability of probability-based importance scores.
• The global sparsity constraint in probability space leads to 57.75% higher accuracy than layer-wise uniform sparsity at 99.9% pruning on CIFAR-10, demonstrating the importance of per-layer budget adaptation.
• ProbMask successfully identifies a supermask with only 2% of the original weights on CIFAR-100, achieving nearly 50% accuracy when weights are fixed at initialization.
• The probability distribution of the final mask converges to 0 or 1 across all layers, confirming that the mask becomes deterministic and sparse after training.
• The method shows robustness across models (ResNet, VGG, DenseNet) and datasets (CIFAR-10/100, ImageNet), consistently outperforming magnitude-based and Hessian-based pruning baselines.