Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

FastForward Pruning: Efficient LLM Pruning via Single-Step Reinforcement Learning

Xin Yuan, Siqi Li, Jiateng Wei, Chengrui Zhu, Yanming Wu, Qingpeng Li, Jiajun Lv, Xiaoke Lan, Jun Chen, Yong Liu

TL;DR

FastForward Pruning tackles the challenge of finding an optimal non-uniform sparsity allocation for large language models under a global budget. It introduces a decoupled, single-step RL framework with Progressive Scheduling to efficiently search for layer-wise retention policies, followed by a retraining-free Ridge regression calibration to recover performance. The method demonstrates competitive or superior results versus strong heuristics and other search-based approaches while substantially reducing search cost, validated on OPT and LLaMA variants. This work provides a practical path to high-performance LLM pruning, enabling deployment with reduced compute and storage demands.

Abstract

Pruning is an effective method for compressing Large Language Models, but finding an optimal, non-uniform layer-wise sparsity allocation remains a key challenge. While heuristic methods are fast but yield suboptimal performance, more powerful search-based approaches like Reinforcement Learning are often hindered by prohibitive computational costs on large-scale models. To overcome this efficiency barrier, we propose FastForward Pruning. Its core is a decoupled, single-step RL framework that separates policy optimization from the complex budget satisfaction problem. Such a decoupling is crucial for efficiently searching the vast policy space of LLMs. This curriculum-based strategy begins with low-cost, simple tasks and gradually increases in complexity, significantly reducing the search's computational overhead. Evaluated on the LLaMA, Mistral, and OPT model families, our framework discovers pruning policies that achieve superior performance over strong heuristic baselines. Crucially, when compared to other search-based algorithms, our method achieves competitive or superior results at a fraction of the computational cost, demonstrating a clear advantage in search efficiency.

FastForward Pruning: Efficient LLM Pruning via Single-Step Reinforcement Learning

TL;DR

FastForward Pruning tackles the challenge of finding an optimal non-uniform sparsity allocation for large language models under a global budget. It introduces a decoupled, single-step RL framework with Progressive Scheduling to efficiently search for layer-wise retention policies, followed by a retraining-free Ridge regression calibration to recover performance. The method demonstrates competitive or superior results versus strong heuristics and other search-based approaches while substantially reducing search cost, validated on OPT and LLaMA variants. This work provides a practical path to high-performance LLM pruning, enabling deployment with reduced compute and storage demands.

Abstract

Pruning is an effective method for compressing Large Language Models, but finding an optimal, non-uniform layer-wise sparsity allocation remains a key challenge. While heuristic methods are fast but yield suboptimal performance, more powerful search-based approaches like Reinforcement Learning are often hindered by prohibitive computational costs on large-scale models. To overcome this efficiency barrier, we propose FastForward Pruning. Its core is a decoupled, single-step RL framework that separates policy optimization from the complex budget satisfaction problem. Such a decoupling is crucial for efficiently searching the vast policy space of LLMs. This curriculum-based strategy begins with low-cost, simple tasks and gradually increases in complexity, significantly reducing the search's computational overhead. Evaluated on the LLaMA, Mistral, and OPT model families, our framework discovers pruning policies that achieve superior performance over strong heuristic baselines. Crucially, when compared to other search-based algorithms, our method achieves competitive or superior results at a fraction of the computational cost, demonstrating a clear advantage in search efficiency.

Paper Structure

This paper contains 11 sections, 7 equations, 2 figures, 4 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Methodology
  3. Reinforcement Learning Formulation for Pruning
  4. Search Acceleration via Progressive Scheduling
  5. Efficient Post-Pruning Calibration
  6. Experiments
  7. Experiment Setup
  8. Main Performance
  9. Analysis of Search Efficacy
  10. Conclusion
  11. acknowledgements

Figures (2)

  • Figure 1: Our single-step RL paradigm (d) avoids the key drawbacks of legacy approaches for LLMs: (a) sparse rewards in sequential decision-making he2018amc, (b) complex state representations yu2021autograph, and (c) multi-agent coordination challenges alwani2022decore.
  • Figure 2: An overview of our FastForward Pruning framework, consisting of two main stages. Stage 1: Efficient Search, where the search process is accelerated by our Progressive Scheduling (ProgSched) mechanism. It uses a unified schedule $\alpha(t)$ to govern a Sparsity Curriculum (SC) for task difficulty and a Fidelity Curriculum (FC) for evaluation cost. An RL Agent outputs raw retention scores, which are then converted into a budget-compliant policy by the Policy Mapping module. Stage 2: Calibration, where the final pruned weights undergo a retraining-free calibration via Ridge Regression to recover performance.