Technical Summary: RAP (Runtime Adaptive Pruning for LLM Inference)

Problem Statement

Large Language Models (LLMs) face significant deployment barriers due to their massive computational and memory footprints. While model compression techniques (pruning, quantization, distillation) offer solutions, existing methods suffer from critical limitations:

1. Static Heuristics: Most current pruning strategies rely on fixed, one-shot heuristics that assume static workloads. They fail to adapt to runtime memory variations caused by fluctuating user requests (e.g., varying batch sizes and sequence lengths).
2. Neglect of KV Cache: Traditional pruning focuses primarily on reducing static model weights (parameters). However, in autoregressive inference, the Key-Value (KV) cache memory grows linearly with batch size and sequence length, often dwarfing the static parameter memory. Fixed pruning ratios cannot accommodate the dynamic shift from parameter-dominated to KV-cache-dominated memory bottlenecks.
3. System-Level Variance: Real-world edge and cloud environments experience stochastic memory availability due to co-running applications and interference, rendering fixed-budget assumptions invalid.

The core challenge addressed is how to select an optimal LLM pruning policy that dynamically adapts to heterogeneous, time-varying request workloads while satisfying fluctuating memory budgets without degrading output quality.

Methodology: The RAP Framework

The authors propose RAP (Runtime Adaptive Pruning), an elastic pruning framework driven by Reinforcement Learning (RL) that dynamically adjusts compression strategies based on real-time system states.

1. Greedy Sequential Importance (GSI)

To address the limitations of "one-shot" pruning (which ignores inter-layer dependencies and leads to performance degradation), RAP employs Greedy Sequential Importance (GSI).

• Mechanism: Instead of scoring all blocks once, GSI iteratively removes the block whose exclusion results in the minimal deterioration of model performance (measured by perplexity).
• Recalibration: After each removal, the importance scores of the remaining blocks are re-evaluated within the new, contracted architecture. This step-wise process controls error accumulation and accounts for the dynamic sensitivity of layers as the model structure changes.

2. Reinforcement Learning Formulation

RAP casts the pruning policy selection as a sequential decision-making problem (Markov Decision Process) where an RL agent learns to balance memory efficiency and model fidelity.

• State ($S_t$): The agent observes a composite state at each inference step $t$:
  • Request Characteristics: Batch size ($R_{bs}$) and sequence length ($R_{sql}$).
  • Model Configuration: Importance scores of all Multi-Head Attention (MHA) and Feed-Forward Network (FFN) blocks derived from GSI.
  • System State: Current available system memory ($Sys_{avail}$) and estimated memory overhead ($Sys_{req}$) for the proposed action.
• Action ($A_t$): The agent selects a pruning policy by deciding whether to retain or remove specific MHA and FFN blocks. To manage the combinatorial explosion of the action space ($2^{2N}$), the decision is decomposed into sequential single-block selections.
• Reward ($R_t$): A unified reward function balances two objectives:
  • Perplexity ($R_{ppl}$): Encourages preserving language modeling performance.
  • Memory Penalty ($R_{mem}$): Penalizes peak memory consumption during inference.
  • The reward is a weighted sum: $R_t = \sum (At)_i (\alpha R_{ppl, i} - \beta R_{mem, i})$, where hyperparameters $\alpha$ and $\beta$ tune the accuracy-memory trade-off.

3. Execution Flow

1. The RL agent observes the real-time environment (input length, batch, memory availability).
2. It selects a pruning action (which blocks to remove) based on the current state.
3. The base model executes inference on the pruned architecture.
4. The agent receives a reward based on the resulting memory usage and perplexity, updating its policy to optimize future decisions.

Key Contributions

1. Runtime-Adaptive Framework: RAP is the first framework to jointly consider model weights and KV cache memory "on the fly," dynamically adjusting sparsity levels based on instantaneous workload and device state.
2. RL-Driven Policy: The authors formulate pruning as an RL problem, training an agent to learn an optimal policy that balances memory savings with generation fidelity, outperforming static heuristic baselines.
3. GSI for Structural Pruning: The introduction of Greedy Sequential Importance addresses the inter-layer dependency issue in transformer pruning, providing a more robust method for determining block importance under aggressive compression.
4. Empirical Validation: Extensive experiments demonstrate that RAP consistently outperforms state-of-the-art static pruning baselines (e.g., LLM-Pruner, SliceGPT, ShortGPT) across diverse models (Llama2, Llama3, Qwen) and memory constraints.

Experimental Results

Experiments were conducted on Llama2-7B, Llama3-8B, Qwen1.5-7B, and Qwen2.5-7B under 80% and 60% unified memory budgets (covering both parameters and KV cache).

• Generation Ability: RAP exhibited the smallest perplexity drift. On Llama2-7B with an 80% budget, RAP increased perplexity by only +6.33 (WikiText2) and +21.47 (PTB), significantly outperforming the next-best method by margins of 22.95 and 146.93, respectively.
• Downstream Tasks: In zero-shot commonsense reasoning, RAP preserved 86% of dense model accuracy under 80% budget, surpassing leading baselines by +8.2% (Llama2-7B) and +11.5% (Llama3-8B).
• Resilience: Under severe 60% memory constraints, RAP was the only method to retain over 60% of the dense model's performance, whereas other methods suffered catastrophic degradation.
• Ablation Studies:
  • Removing GSI (using static one-shot scoring) caused a 13.17% drop in average commonsense accuracy.
  • Removing the RL agent (using random block dropping) resulted in a massive increase in perplexity (+301.71 on WikiText2) compared to the learned policy, highlighting the necessity of adaptive decision-making.
• Overhead: The RL controller adds negligible runtime overhead (<1% latency increase) due to its compact architecture (18K parameters), while the offline GSI calculation is a one-time cost amortized over deployment.

Significance and Claims

The paper claims that RAP bridges the gap between static compression techniques and dynamic real-world deployment scenarios. By moving away from "one-size-fits-all" heuristics, RAP enables:

• Scalable Edge Deployment: It allows LLMs to fit on-device under tight, fluctuating memory constraints by pruning substantial weights when necessary, while retaining full capacity when resources allow.
• Efficiency without Quality Loss: It achieves superior memory savings and faster inference with minimal impact on output quality, effectively navigating the trade-offs between parameter-heavy and KV-cache-heavy bottlenecks.
• Generalizability: The framework is shown to be architecture-agnostic, performing well across different model families (Llama, Qwen) and generations.

The authors position RAP as a flexible, context-aware compression mechanism that makes LLM inference feasible in heterogeneous environments where fixed strategies fail.