Source-linked AI summary

Comparing Rewinding and Fine-tuning in Neural Network Pruning

Alex Renda, Jonathan Frankle, Michael Carbin

arXiv:2003.02389v1cs.LGstat.ML

TL;DR

Neural network pruning must reduce resource demands without sacrificing accuracy, but the standard fine-tuning retraining method is only one possible choice. This paper compares fine-tuning with weight and learning-rate rewinding, finding that both outperform fine-tuning and support a network-agnostic pruning algorithm matching state-of-the-art tradeoffs.

  • Problem

    The paper examines how to retrain pruned networks while balancing accuracy, parameter-efficiency, and retraining search cost.

  • Method

    The paper compares fine-tuning, weight rewinding, and learning-rate rewinding as retraining techniques after pruning.

  • Results

    Both rewinding techniques outperform fine-tuning, while learning-rate rewinding matches or outperforms weight rewinding in all evaluated scenarios.

  • Takeaways & Limitations

    A learning-rate-rewinding pruning algorithm with network-agnostic hyperparameters matches state-of-the-art accuracy–parameter-efficiency tradeoffs across networks and datasets.

  • Takeaways & Limitations

    The study focuses on parameter-efficiency and does not study other efficiency forms such as inference cost in its main analysis.

Abstract

from arXiv · show

Many neural network pruning algorithms proceed in three steps: train the network to completion, remove unwanted structure to compress the network, and retrain the remaining structure to recover lost accuracy. The standard retraining technique, fine-tuning, trains the unpruned weights from their final trained values using a small fixed learning rate. In this paper, we compare fine-tuning to alternative retraining techniques. Weight rewinding (as proposed by Frankle et al., (2019)), rewinds unpruned weights to their values from earlier in training and retrains them from there using the original training schedule. Learning rate rewinding (which we propose) trains the unpruned weights from their final values using the same learning rate schedule as weight rewinding. Both rewinding techniques outperform fine-tuning, forming the basis of a network-agnostic pruning algorithm that matches the accuracy and compression ratios of several more network-specific state-of-the-art techniques.

1 INTRODUCTION

The paper compares fine-tuning with weight and learning-rate rewinding as retraining methods after pruning. Both rewinding methods outperform fine-tuning, and learning-rate rewinding supports a network-agnostic pruning algorithm matching state-of-the-art tradeoffs.

  • Motivation: Pruning removes network structures to reduce parameter count and inference resources while maintaining accuracy.The paper frames pruning around accuracy, efficiency, and search cost.
  • Retraining-based pruning: Retraining-based pruning trains the network, removes structures, and retrains the remainder to recover accuracy lost through pruning.Pruning and retraining may be repeated until a target sparsity or accuracy threshold is reached.
  • Retraining methods: Fine-tuning retrains unpruned weights from their final values for additional epochs using a fixed learning rate.This is described as the common standard retraining technique.
  • Retraining methods: Weight rewinding restores unpruned weights to earlier training values and resumes training with the original schedule, with the rewind point controlled by retraining time.Retraining for t epochs rewinds weights to t epochs before training ended.
  • Retraining methods: Learning-rate rewinding keeps final weights but restores the learning-rate schedule, matching or outperforming weight rewinding in all evaluated scenarios.The technique is proposed as a simplification of weight rewinding.
  • Results: Weight rewinding outperforms fine-tuning across networks and datasets, achieving higher accuracy at equivalent search cost across a wide rewind range.The result makes weight rewinding a drop-in replacement for fine-tuning.
  • Results: The proposed learning-rate-rewinding pruning algorithm uses network-agnostic hyperparameters and matches state-of-the-art accuracy–parameter-efficiency tradeoffs across networks and datasets.It repeatedly prunes the lowest-magnitude 20% of weights and retrains until the desired sparsity is reached.

2 METHODOLOGY

The methodology evaluates pruning and retraining choices across standard vision and translation networks using accuracy, efficiency, and search-cost criteria. It compares magnitude-based one-shot and iterative pruning with fine-tuning and both rewinding methods.

  • Networks and datasets: The study evaluates pruning on ResNet-56, ResNet-34, ResNet-50, and GNMT across CIFAR-10, ImageNet, and WMT16 EN-DE.Implementations and most hyperparameters follow standard reference implementations.
  • Pruning structure: The paper studies both unstructured pruning of individual weights and structured pruning of groups such as filters, neurons, and channels.Structured pruning preserves dense computation, while unstructured pruning reduces parameter count.
  • Pruning heuristic: Magnitude pruning removes the lowest-magnitude weights globally for unstructured pruning and uses filter L1 norms with specified per-layer rates for structured pruning.The structured experiments use ResNet-56-B and ResNet-34-A configurations.
  • Pruning heuristic: The paper considers only magnitude-based pruning heuristics, excluding methods that learn pruning decisions during optimization or use other information.This is an explicit scope boundary of the pruning methodology.
  • Retraining methods: Fine-tuning starts from final unpruned weights with a fixed learning rate, whereas weight rewinding restores weights and schedule from t epochs earlier.The methods differ in both initialization and learning-rate schedule.
  • Retraining methods: Learning-rate rewinding starts from final weights but uses the learning-rate schedule from the last t training epochs.Formally, it runs TRAIN_t(W_T, m, T−t).
  • Experimental design: The experiments compare ten retraining times from zero to the original training duration for each network, repeating the choice at each iterative-pruning step.Both one-shot and iterative pruning are evaluated.
  • Metrics: Accuracy, parameter-efficiency, and search cost are the three evaluation criteria.Search cost is measured by retraining epochs, while parameter-efficiency is represented by compression ratio.

3 ACCURACY VERSUS PARAMETER-EFFICIENCY TRADEOFF

Across compression ratios, weight rewinding outperforms fine-tuning, while learning rate rewinding matches or outperforms weight rewinding. With iterative unstructured pruning, learning rate rewinding matches state-of-the-art accuracy–parameter-efficiency tradeoffs.

  • Weight rewinding achieves higher accuracy than fine-tuning across compression ratios on all studied networks and datasets.
  • Learning rate rewinding matches or outperforms weight rewinding in all tested scenarios.
  • One-shot pruning results: At higher compression ratios, fine-tuning loses more accuracy than either rewinding technique.
  • One-shot pruning results: Learning rate rewinding typically exceeds the unpruned network’s accuracy at low compression ratios, while other techniques only match it.
  • Iterative pruning results: Learning rate rewinding matches state-of-the-art accuracy–parameter-efficiency tradeoffs across all datasets with iterative unstructured pruning.
  • Iterative pruning results: 5.96× compression preserves the original ResNet-50 accuracy with iterative unstructured pruning and learning rate rewinding.

4 ACCURACY VERSUS SEARCH COST TRADEOFF

Both rewinding methods generally outperform fine-tuning across retraining times, with learning rate rewinding providing the highest accuracy when retrained for the original training duration. The rewinding safe zone spans 25% to 90% of the original training time for unstructured pruning.

  • Both rewinding techniques achieve higher accuracy than fine-tuning across many retraining times and can serve as drop-in replacements.
  • Unstructured pruning results: The rewinding safe zone spans 25% to 90% of the original training time for unstructured pruning.
  • Unstructured pruning results: Within the safe zone, both rewinding methods outperform fine-tuning at equivalent search costs.
  • Accuracy from rewinding usually saturates after roughly half of the original training time on most networks and compression ratios.
  • Structured pruning results: Structured pruning follows the same general trends as unstructured pruning, except weight rewinding does not drop in accuracy at full training duration.
  • Learning rate rewinding with full-duration retraining produces the highest accuracy among tested retraining techniques.
  • On CIFAR-10 ResNet-56 Structured-B at 1% Accuracy Drop, learning rate rewinding is lower than fine-tuning at 30 epochs, apparently because of random noise.

5 OUR PRUNING ALGORITHM

The proposed pruning algorithm uses network-agnostic iterative pruning with learning rate rewinding. It matches the accuracy of more hyperparameter-intensive pruning methods without per-compression-ratio hyperparameter search.

  • The proposed algorithm uses network-agnostic hyperparameters and lies on the state-of-the-art accuracy–parameter-efficiency Pareto frontier.
  • Each iteration globally prunes the 20% lowest-magnitude weights, then retrains with learning rate rewinding for the original training time.
  • The algorithm repeats pruning and retraining until the desired compression ratio is reached.
  • Its accuracy–parameter-efficiency tradeoff matches state-of-the-art behavior without per-compression-ratio hyperparameter search.
  • Shared hyperparameters apply across all evaluated networks and tasks, with no layer-wise pruning rates or pruning schedule to select.
  • The algorithm matches the accuracy of pruning methods requiring more hyperparameters or additional methods.

6 DISCUSSION

The discussion finds that rewinding improves pruning retraining, while highlighting open questions about optimization, retraining design, search cost, and efficiency metrics.

  • Weight rewinding outperforms fine-tuning and remains competitive with learning rate rewinding despite reverting most of training.The authors argue that the pruning mask itself is a valuable output, rather than merely a way to identify weights requiring extensive further training.
  • Learning rate rewinding achieves state-of-the-art accuracy versus parameter-efficiency tradeoffs, but the paper does not investigate why its schedule performs better than fine-tuning.The authors identify optimization of sparse neural networks as a direction for explaining and extending this result.
  • The retraining techniques reuse original-training hyperparameters, narrowing the design space and leaving possible gains from decoupled retraining schedules.Those gains would require additional hyperparameter search.
  • SEARCH COST: T · (1 + k) training epochs are required to reach compression ratio 1 / 0.8^k, making search cost a substantial consideration.Here T is the original training time and k is the number of pruning iterations; the paper compares this cost with several other methods.
  • EFFICIENCY: The study focuses on parameter-efficiency, although both rewinding techniques outperform fine-tuning after structured pruning and can reduce FLOPs after iterative unstructured pruning.Storage-, communication-, and energy-efficiency are identified as other efficiency measures not studied here.
  • The Lottery Ticket Hypothesis: The paper presents a detailed comparison between lottery ticket networks and pruned networks produced by standard fine-tuning.Weight rewinding originated in work on sparse subnetworks that can train in isolation to full accuracy from near initialization.

7 CONCLUSION

The paper concludes that both rewinding methods outperform fine-tuning, and that learning rate rewinding can match more complex network-specific pruning techniques.

  • Both weight rewinding and learning rate rewinding outperform fine-tuning for retraining after pruning.
  • Iterative unstructured pruning with learning rate rewinding matches the accuracy versus parameter-efficiency tradeoffs of more complex network-specific techniques.The authors propose this algorithm as a valuable baseline and compelling practical default.

Appendix D. We evaluate other retraining methods using the algorithm described in Section 5.

Appendix E extends the paper’s evaluation across a broader set of experimental comparisons.

  • The extended evaluation includes more networks, pruning techniques, baselines, and ablations.

C STATE-OF-THE-ART BASELINES

The paper compares its pruning algorithm with methods selected around accuracy versus parameter-efficiency tradeoffs across several network and task settings.

  • Accuracy is measured as relative loss from the original network, while parameter-efficiency is measured by compression ratio.
  • CIFAR-10 ResNet-56: Learning Compression is used as the CIFAR-10 ResNet-56 baseline because it is identified as most accurate at high sparsities.
  • CIFAR-10 ResNet-56: The CIFAR-10 baseline uses unstructured gradual global magnitude pruning with a pruning schedule followed by fine-tuning.
  • ImageNet ResNet-50: AMC is used for ImageNet ResNet-50 and reports 0.02% top-1 accuracy loss at 5.13× compression.
  • ImageNet ResNet-50: AMC iteratively prunes and fine-tunes ResNet-50, using layerwise pruning rates determined with reinforcement learning.The reported pruning sequence is 50%, 35%, 25%, then 20%, producing 80.5% sparsity.
  • WMT16 EN-DE GNMT: Zhu and Gupta’s WMT16 EN-DE GNMT method is included after a search found no cited work claiming a better accuracy versus parameter-efficiency curve.
  • WMT16 EN-DE GNMT: Their method gradually prunes during training using a polynomial pruning schedule and excludes attention layers from equal-rate pruning.
  • Figure 6 evaluates the pruning algorithm and finds weight rewinding generally similar to learning rate rewinding, except at high GNMT sparsities.

D OTHER INSTANTIATIONS OF OUR PRUNING ALGORITHMS

The appendix compares alternative retraining techniques and extends evaluation across additional networks, pruning methods, baselines, and ablations. The algorithm repeatedly trains, prunes low-magnitude weights, retrains, and stops at the desired compression.

  • Retraining techniques: The appendix compares weight rewinding for 90% of the original training time, learning rate rewinding for the original schedule, and fine-tuning for the original schedule.
  • Algorithm: The pruning algorithm trains to completion, globally removes the 20% lowest-magnitude weights, retrains, and repeats until reaching the desired compression ratio.Retraining can use weight rewinding, learning rate rewinding, or fine-tuning.
  • Results: Weight rewinding performs similarly to learning rate rewinding except at high sparsities on GNMT.
  • Evaluation scope: The evaluation adds more networks, structured pruning techniques, a reinitialization baseline, and an ablation combining rewound weights with fine-tuning’s learning rate.
  • Evaluation scope: The discovered pruning structure is determined after training, using magnitude pruning for unstructured pruning and layerwise rates from Li et al. for structured pruning.

RESULTS

Across accuracy, parameter-efficiency, and search-cost evaluations, rewinding generally outperforms fine-tuning, while reinitialization behaves differently for unstructured and structured pruning. Low-learning-rate weight rewinding is a notable weaker ablation.

  • Accuracy versus parameter-efficiency: Low-LR weight rewinding causes a large accuracy drop relative to the best achievable accuracy, while reinitialization performs poorly for unstructured pruning.
  • Accuracy versus parameter-efficiency: With structured pruning, reinitialization roughly matches weight and learning rate rewinding, consistent with prior findings that it performs comparatively well for structured techniques.
  • Accuracy versus search cost: Low-LR weight rewinding peaks at different training points across networks and compression ratios, often near the middle of training rather than after longer training.
  • Accuracy versus search cost: Reinitialization typically saturates with the original training schedule, whereas sufficient retraining can let learning rate rewinding achieve higher accuracy under structured pruning.
  • Figure comparisons: Figures 11 and 12 report speedup over the original network for different retraining techniques and networks, including iterative-pruning FLOP speedup.

F COMPRESSION RATIO VS FLOPS

Compression ratio alone does not capture inference cost because pruning different layers removes different numbers of FLOPs. The appendix therefore compares FLOP speedups for the most accurate networks at each compression ratio.

  • Motivation: Compression ratio does not fully measure efficiency because pruning an early convolutional-layer weight can remove more FLOPs than pruning a later-layer weight.
  • Methodology: The analysis compares FLOPs for iteratively pruned vision networks, where earlier-layer pruning produces larger FLOP reductions than later-layer pruning.
  • Methodology: Reported FLOPs correspond to the most accurate network at each compression ratio, not necessarily the minimum-FLOP network at that ratio.
  • Results: Both rewinding techniques find lower-FLOP networks than fine-tuning, and a ResNet-50 with equal original accuracy achieves a 4.8× theoretical speedup using learning rate rewinding.
Loading 2003.02389v1…