[Paper Review] Winning the Lottery Ahead of Time: Efficient Early Network Pruning

John Rachwan, Daniel Zügner+4 more|arXiv (Cornell University)|Jun 21, 2022

Advanced Neural Network Applications4 citations

TL;DR

This paper proposes EarlyCroP, a method for efficient early network pruning that preserves gradient flow to maintain model performance while enabling structured pruning. It achieves state-of-the-art accuracy with up to 98.5% weight sparsity, reduces training time by 7x, GPU memory by 4.9x, and carbon emissions by 70%, making large models trainable on commodity hardware without sacrificing performance.

ABSTRACT

Pruning, the task of sparsifying deep neural networks, received increasing attention recently. Although state-of-the-art pruning methods extract highly sparse models, they neglect two main challenges: (1) the process of finding these sparse models is often very expensive; (2) unstructured pruning does not provide benefits in terms of GPU memory, training time, or carbon emissions. We propose Early Compression via Gradient Flow Preservation (EarlyCroP), which efficiently extracts state-of-the-art sparse models before or early in training addressing challenge (1), and can be applied in a structured manner addressing challenge (2). This enables us to train sparse networks on commodity GPUs whose dense versions would be too large, thereby saving costs and reducing hardware requirements. We empirically show that EarlyCroP outperforms a rich set of baselines for many tasks (incl. classification, regression) and domains (incl. computer vision, natural language processing, and reinforcment learning). EarlyCroP leads to accuracy comparable to dense training while outperforming pruning baselines.

Motivation & Objective

• To address the high computational cost of finding sparse winning tickets in deep neural networks.
• To enable structured pruning that reduces training time, GPU memory usage, and carbon emissions without sacrificing model accuracy.
• To identify the optimal early training point for pruning by tracking gradient flow and Neural Tangent Kernel stability.
• To make training of large, sparse models feasible on commodity GPUs by preserving learning dynamics through gradient flow preservation.
• To outperform existing unstructured and structured pruning baselines across vision, NLP, and reinforcement learning tasks.

Proposed method

• Proposes a gradient flow (GF)-based pruning criterion that minimizes disruption to the model's learning dynamics.
• Leverages the connection between gradient flow and the Neural Tangent Kernel (NTK) to identify stable training phases suitable for early pruning.
• Applies pruning during the lazy kernel regime, where NTK remains approximately constant, ensuring minimal performance drop.
• Extends the method to structured pruning by removing entire neurons or filters while preserving gradient flow and model performance.
• Uses a single training run to identify and prune sparse subnetworks, eliminating the need for multiple dense training cycles.
• Employs a principled pruning schedule based on NTK stability and GF preservation to determine the earliest viable pruning time.

Experimental results

Research questions

• RQ1 Can we identify an early training phase where pruning has minimal impact on model performance, enabling efficient subnetwork discovery?
• RQ2 Can gradient flow preservation serve as a principled criterion for structured pruning that maintains model accuracy and efficiency?
• RQ3 Does early structured pruning outperform existing structured and unstructured baselines in terms of accuracy, inference speed, and carbon footprint?
• RQ4 Can sparse models trained via early pruning achieve better performance than similarly sized dense models?
• RQ5 To what extent does early pruning reduce training time, GPU memory usage, and carbon emissions compared to dense training?

Key findings

• EarlyCroP-U achieves test accuracy comparable to or better than the dense model and outperforms all unstructured baselines, including LTH and GraSP, at high sparsity levels.
• EarlyCroP-S reduces training time by 7x, GPU memory usage by 4.9x, and carbon emissions by up to 70% compared to dense training.
• On CIFAR-100 with VGG16, EarlyCroP-S achieves 92.5% test accuracy with 98.5% weight sparsity and 89.9% node sparsity, outperforming the dense model.
• For the ResNeXt101_32x48d model, EarlyCroP-S achieves 93.2% accuracy after 80 epochs with 98.5% sparsity, training 6.2x faster and emitting 9.5x less carbon than a similar-sized dense model.
• In NLP tasks, EarlyCroP-S maintains performance at 89% sparsity and outperforms all baselines beyond that point, while LTR fails due to layer-specific weight decay.
• In reinforcement learning, EarlyCroP-S outperforms LTR and unstructured baselines, with unstructured models catching up only after longer training due to inefficient gradient computation.