[Paper Review] ChipNet: Budget-Aware Pruning with Heaviside Continuous Approximations

Rishabh Tiwari, Udbhav Bamba+2 more|arXiv (Cornell University)|Feb 14, 2021

Advanced Neural Network ApplicationsComputer Science35 references31 citations

TL;DR

ChipNet presents a deterministic, budget-aware structured pruning method that uses continuous Heaviside projections and a crispness loss to obtain highly sparse pruned networks from pretrained dense models, outperforming state-of-the-art baselines across multiple budgets and datasets.

ABSTRACT

Structured pruning methods are among the effective strategies for extracting small resource-efficient convolutional neural networks from their dense counterparts with minimal loss in accuracy. However, most existing methods still suffer from one or more limitations, that include 1) the need for training the dense model from scratch with pruning-related parameters embedded in the architecture, 2) requiring model-specific hyperparameter settings, 3) inability to include budget-related constraint in the training process, and 4) instability under scenarios of extreme pruning. In this paper, we present ChipNet, a deterministic pruning strategy that employs continuous Heaviside function and a novel crispness loss to identify a highly sparse network out of an existing dense network. Our choice of continuous Heaviside function is inspired by the field of design optimization, where the material distribution task is posed as a continuous optimization problem, but only discrete values (0 or 1) are practically feasible and expected as final outcomes. Our approach's flexible design facilitates its use with different choices of budget constraints while maintaining stability for very low target budgets. Experimental results show that ChipNet outperforms state-of-the-art structured pruning methods by remarkable margins of up to 16.1% in terms of accuracy. Further, we show that the masks obtained with ChipNet are transferable across datasets. For certain cases, it was observed that masks transferred from a model trained on feature-rich teacher dataset provide better performance on the student dataset than those obtained by directly pruning on the student data itself.

Motivation & Objective

• Motivate robust structured pruning that can enforce arbitrary budget constraints without retraining from scratch.
• Develop a pruning mechanism that yields near-discrete channel masks from dense networks using gradient-based optimization.
• Enable stability under very low budget regimes and demonstrate transferability of learned masks across datasets.
• Show that ChipNet achieves superior accuracy under various budgets compared to existing pruning methods.

Proposed method

• Learn sparsity masks for channels in a pretrained dense CNN.
• Use a continuous Heaviside projection combined with a logistic proxy to drive masks toward 0 or 1.
• Introduce a crispness loss that penalizes intermediate mask values to encourage discrete masks.
• Impose budget constraints via a budget loss that can accommodate channel, activation-volume, parameter, or FLOPs budgets.
• Train with a joint loss that includes cross-entropy, crispness, and budget terms; perform soft pruning followed by hard pruning via binary masking.
• Allow budget function V to be any of several options (channel, volume, parameter, FLOPs) and compute estimates on a logistic-projected version of masks during budgeting.

Experimental results

Research questions

• RQ1 Can a deterministic, budget-aware pruning method produce highly sparse networks without retraining from scratch?
• RQ2 Do continuous Heaviside projections plus a crispness loss yield near-discrete masks that satisfy diverse budget constraints?
• RQ3 Are learned pruning masks transferable across datasets and task domains?
• RQ4 How does ChipNet perform relative to state-of-the-art budget-aware pruning methods under multiple budget formulations?
• RQ5 Is ChipNet stable and effective under very low resource budgets (e.g., extreme pruning)?

Key findings

• ChipNet consistently outperforms state-of-the-art pruning baselines across different budget types and datasets.
• The method remains stable and effective even at very low budgets, including extreme pruning scenarios.
• Masks learned by ChipNet can be transferred across datasets, sometimes outperforming masks learned directly on the target data.
• In several experiments, ChipNet achieved accuracy gains of up to a notable margin over baselines under similar budgets.
• The transfer of masks from a feature-rich teacher dataset can provide better student performance than pruning on the student data alone.