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

2025-11-24

AI Summary

Existing approaches for layer-wise non-uniform sparsity allocation in large language model (LLM) pruning face a trade-off: heuristic methods yield suboptimal performance, while reinforcement learning (RL)-based search incurs prohibitive computational overhead. Method: This paper proposes an efficient one-shot decoupled RL framework. Its core innovation is the explicit decoupling of policy learning from resource constraint satisfaction, augmented by a curriculum learning mechanism to substantially reduce search complexity. Contribution/Results: The method achieves global sparsity allocation via a single policy update—eliminating iterative RL optimization. It consistently outperforms strong heuristic baselines on LLaMA, Mistral, and OPT models, delivering higher post-pruning accuracy. Compared to conventional multi-step RL methods, it reduces training cost by an order of magnitude, achieving both high efficiency and state-of-the-art performance in structured LLM pruning.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Search and Optimization: Learning to Search
• Natural Language Processing: (Large) Language Models

Application Category

• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• Search and Retrieval-Augmented AI: Large language models for search

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.

Problem

• Finding optimal non-uniform sparsity allocation in LLM pruning
• Overcoming computational cost barriers in reinforcement learning approaches
• Efficiently searching vast policy space while maintaining model performance

Innovation

• Single-step reinforcement learning for pruning
• Decoupled policy optimization from budget constraints
• Curriculum-based strategy reducing computational overhead