[Paper Review] Good Subnetworks Provably Exist: Pruning via Greedy Forward Selection

Mao Ye, Chengyue Gong+4 more|arXiv (Cornell University)|Mar 3, 2020

Advanced Neural Network ApplicationsComputer Science52 references37 citations

TL;DR

The paper proves that greedy forward selection can find small subnetworks with provably lower loss than training from scratch, and demonstrates strong empirical pruning results on ImageNet models.

ABSTRACT

Recent empirical works show that large deep neural networks are often highly redundant and one can find much smaller subnetworks without a significant drop of accuracy. However, most existing methods of network pruning are empirical and heuristic, leaving it open whether good subnetworks provably exist, how to find them efficiently, and if network pruning can be provably better than direct training using gradient descent. We answer these problems positively by proposing a simple greedy selection approach for finding good subnetworks, which starts from an empty network and greedily adds important neurons from the large network. This differs from the existing methods based on backward elimination, which remove redundant neurons from the large network. Theoretically, applying the greedy selection strategy on sufficiently large {pre-trained} networks guarantees to find small subnetworks with lower loss than networks directly trained with gradient descent. Our results also apply to pruning randomly weighted networks. Practically, we improve prior arts of network pruning on learning compact neural architectures on ImageNet, including ResNet, MobilenetV2/V3, and ProxylessNet. Our theory and empirical results on MobileNet suggest that we should fine-tune the pruned subnetworks to leverage the information from the large model, instead of re-training from new random initialization as suggested in \citet{liu2018rethinking}.

Motivation & Objective

• Motivate why pruning via forward selection can yield better small subnetworks than training from scratch.
• Propose a simple greedy forward selection algorithm to build subnetworks from a pre-trained large network.
• Provide theoretical guarantees on loss decay for the pruned subnetworks under mild assumptions.
• Demonstrate strong empirical pruning performance on ImageNet architectures (ResNet, MobileNet variants, ProxylessNet).
• Discuss practical implications for finetuning pruned subnetworks versus retraining from random initialization.

Proposed method

• Define the subnetwork selection problem as choosing a subset S of at most n neurons from a large network to minimize loss.
• Propose a forward greedy rule: at each step add the neuron i that minimizes L[f_{S ∪ {i}}].
• Set f_S as the average of activations from neurons in S, preserving original weights.
• Show that the greedy procedure achieves L[f_{S_n}] = O(1/n) + L_N^* under mild assumptions (Assumption 1).
• Prove a faster rate L[f_{S_n}] = O(1/n^2) when the large network is sufficiently over-parameterized (Assumption 2).
• Relate Assumption 2 to over-parameterization and the Frank-Wolfe interior-point condition.
• Extend results to randomly weighted networks (Lottery Ticket context) and compare forward vs backward elimination.
• Provide layer-wise practical Algorithm 1 for deep networks and discuss finetuning vs retraining.

Experimental results

Research questions

• RQ1 Do good small subnetworks exist inside a large pre-trained network that outperform networks of the same size trained from scratch?
• RQ2 Can a simple greedy forward selection procedure provably find such subnetworks with favorable loss decay rates?
• RQ3 How do the theoretical guarantees compare to gradient-descent training of small networks and to backward-elimination pruning?
• RQ4 Are the theoretical results borne out empirically on modern architectures and datasets (e.g., ImageNet)?

Key findings

• Greedy forward selection yields subnetworks of size n with loss L[f_{S_n}] = O(1/n) + L_N^* generally.
• If the large network is sufficiently over-parameterized, L[f_{S_n}] = O(1/n^2).
• The O(1/n^2) rate extends to randomly initialized large networks (i.i.d. weights).
• Pruned subnetworks often outperform similarly-sized networks trained from scratch under ImageNet FLOPs constraints, especially for mobile architectures.
• Finetuning pruned MobileNet variants after pruning consistently improves over retraining from scratch, aligning with the theory.
• Forward selection empirically outperforms backward elimination on ResNet34 and MobileNetV2 in pruning efficiency and accuracy.