One Shot vs. Iterative: Rethinking Pruning Strategies for Model Compression

Technical Summary: One Shot vs. Iterative: Rethinking Pruning Strategies for Model Compression

Problem Statement
As neural networks grow in scale to address complex tasks in computer vision and natural language processing, their deployment on resource-constrained devices becomes increasingly impractical. Pruning is a primary technique for reducing model size and computational cost without significantly sacrificing performance. However, the literature exhibits a significant bias toward proposing new pruning criteria (e.g., magnitude, sensitivity, game-theoretic approaches) while lacking a systematic comparison of pruning regimes—specifically, the choice between one-shot pruning (a single cycle of training, pruning, and retraining) and iterative pruning (multiple cycles of these steps). The relative efficacy of these strategies under varying conditions (different pruning ratios, model architectures, and criteria) remains under-explored, often leading to arbitrary selection of methods by practitioners.

Methodology
The authors conducted a comprehensive empirical study to isolate the pruning regime as a variable, controlling for pruning criteria, target sparsity, and model architectures. The study covers both structured and unstructured pruning across Convolutional Neural Networks (ResNet, EfficientNet), Vision Transformers (ViT), and a Transformer-based language model (TinyStories-33M).

Key methodological components include:

1. Regime Definitions:
   • One-Shot: A single pruning step followed by fine-tuning.
   • Iterative Constant: A fixed number of weights removed in each cycle.
   • Iterative Geometric: A fixed percentage of the remaining weights removed in each step, progressively reducing the number of pruned weights.
   • Hybrid (Few-Shot): A proposed combination where a large portion of weights is removed in an initial one-shot step, followed by fine-grained geometric iterative pruning.
2. Retraining Protocol: The study advocates for patience-based fine-tuning (early stopping based on validation performance) rather than fixed epoch counts, arguing that this dynamically adapts the retraining duration to the actual recovery needs of the model.
3. Evaluation Metrics: Performance was measured via Top-1 Accuracy for vision tasks and Perplexity for the NLP task, evaluated across various pruning ratios (from low to extreme sparsity).

Key Contributions

• Systematic Comparison: This work provides one of the first broad, controlled comparisons of one-shot versus iterative pruning regimes, challenging the assumption that iterative pruning is universally superior.
• Geometric Pruning Scheduler: The authors introduce and validate a geometric pruning ratio scheduler for iterative pruning, demonstrating that pruning a fixed percentage of remaining weights generally outperforms constant-rate pruning, particularly at high sparsity levels.
• Hybrid Approach: A novel hybrid pruning regime is proposed, which leverages the efficiency of one-shot pruning for initial coarse reduction and the precision of iterative geometric pruning for final refinement.
• Contextual Guidelines: The study establishes that the optimal pruning strategy is not universal but depends on specific constraints: pruning ratio, model architecture (CNN vs. Transformer), and the computational budget for retraining.

Key Results

• Pruning Ratio Dependency:
  • Low to Moderate Ratios (<80%): One-shot pruning with patience-based retraining often outperforms iterative methods, particularly for CNNs. It achieves higher accuracy with significantly less retraining time.
  • High Ratios (>80%): Iterative pruning (specifically the geometric variant) becomes superior. One-shot pruning struggles to recover performance at extreme sparsity levels.
• Architecture Sensitivity: Transformer-based models (both ViT and LLMs) generally benefit more from iterative pruning than CNNs. In NLP tasks, one-shot pruning showed sensitivity even at low pruning ratios (10–20%), whereas iterative pruning allowed the model to progressively adapt.
• Pruning Criteria Interaction:
  • For second-order criteria (e.g., Hessian-based), one-shot pruning is computationally efficient because the expensive ranking calculation is performed only once.
  • For high pruning ratios, the choice of pruning criterion becomes less significant; the iterative geometric regime yields similar performance across different criteria, making faster magnitude-based criteria preferable.
• Hybrid Performance: The hybrid approach consistently outperformed pure one-shot and pure iterative methods across nearly all pruning rates, effectively balancing the speed of one-shot removal with the refinement of iterative steps.
• Retraining Efficiency: Patience-based early stopping was shown to be critical. Fixed short retraining periods (common in literature) led to suboptimal results, while patience-based methods prevented both under-training and wasteful over-training.

Significance and Claims
The paper does not claim to establish a single "best" pruning method for all scenarios. Instead, it aims to provide a rigorous empirical basis for practitioners to select a pruning strategy tailored to their specific goals and constraints. The authors argue that the choice of pruning regime is as crucial as the choice of pruning criteria. By demonstrating that one-shot pruning is highly effective for lower sparsity and computationally constrained environments, while iterative (geometric) pruning is necessary for high sparsity and complex architectures, the study offers actionable guidelines for optimizing the trade-off between model accuracy, size, and computational cost. The proposed hybrid method serves as a practical solution that combines the strengths of both regimes, offering a more informed path for model compression.