Source-linked AI summary

From LLMs to LRMs: Rethinking Pruning for Reasoning-Centric Models

Longwei Ding, Anhao Zhao, Fanghua Ye, Ziyang Chen, Xiaoyu Shen

arXiv:2601.18091v1cs.LG

TL;DR

Pruning strategies have been studied mainly on instruction-following LLMs, leaving their transfer to reasoning-augmented models with long reasoning traces unclear. The paper compares static depth, static width, and dynamic pruning under training-distribution-aligned calibration and recovery across 17 tasks. It finds task- and paradigm-dependent trade-offs: dynamic pruning is stronger for instruction-following classification and generation, while static pruning better preserves reasoning performance.

  • Problem

    Existing pruning studies focus mainly on LLM-instruct, leaving unclear whether established strategies transfer to LLM-think models that generate long intermediate reasoning traces.

  • Method

    The study aligns calibration and post-pruning recovery data with each model’s original training distribution and evaluates static depth, static width, and dynamic pruning across 17 tasks.

  • Results

    Static depth pruning is more effective for classification, width pruning degrades more gracefully for generation and reasoning, and dynamic pruning benefits instruction-following models but underperforms static pruning on reasoning-augmented models.

  • Takeaways & Limitations

    Pruning strategies should account for the distinct characteristics of reasoning-augmented LLMs and their long-chain reasoning workloads.

  • Takeaways & Limitations

    The experiments use Llama-8B–scale models and do not explicitly test whether the observed pruning behaviors generalize across model scales.

Abstract

from arXiv · show

Large language models (LLMs) are increasingly costly to deploy, motivating extensive research on model pruning. However, most existing studies focus on instruction-following LLMs, leaving it unclear whether established pruning strategies transfer to reasoning-augmented models that explicitly generate long intermediate reasoning traces. In this work, we conduct a controlled study of pruning for both instruction-following ($\textbf{LLM-instruct}$) and reasoning-augmented ($\textbf{LLM-think}$) models. To isolate the effects of pruning, we align pruning calibration and post-pruning recovery data with each model's original training distribution, which we show yields more stable and reliable pruning behavior. We evaluate static depth pruning, static width pruning, and dynamic pruning across 17 tasks spanning classification, generation, and reasoning. Our results reveal clear paradigm-dependent differences: depth pruning outperforms width pruning on classification tasks, while width pruning is more robust for generation and reasoning. Moreover, static pruning better preserves reasoning performance, whereas dynamic pruning excels on classification and generation but remains challenging for long-chain reasoning. These findings underscore the need for pruning strategies that explicitly account for the distinct characteristics of reasoning-augmented LLMs. Our code is publicly available at https://github.com/EIT-NLP/LRM-Pruning.

1 Introduction

This study asks whether pruning methods developed for instruction-following models generalize to reasoning-augmented models that produce long intermediate traces. It uses aligned training-distribution data and evaluates three pruning families across 17 classification, generation, and reasoning tasks.

  • LLMs’ strong performance creates substantial computational costs for practical use in resource-constrained environments.
  • Reasoning-augmented LLM-think models externalize step-by-step reasoning and generate markedly longer outputs than LLM-instruct models.These generations often span thousands of tokens.
  • Existing pruning research focuses almost exclusively on LLM-instruct, leaving unclear whether its methods generalize to LLM-think.
  • The study aligns pruning calibration and post-pruning recovery data with each model’s original training distribution to isolate pruning effects.It constructs matched model–dataset pairs for instruction-following and reasoning-augmented models.
  • The aligned pipeline best preserves pre-pruning performance and produces more stable, effective pruning than mismatched datasets.
  • Across 17 tasks, static depth pruning is stronger for classification, while width pruning is more robust for generation and reasoning.The evaluation covers classification, generation, and reasoning tasks across both model paradigms.

2 Background

LLM-instruct models target instruction-aligned outputs, whereas LLM-think models generate long chains of thought before answering. This section introduces structured pruning formulations and the compression–performance-retention objective used to evaluate them.

  • LLM-instruct and LLM-thinking: LLM-instruct models optimize instruction-aligned outputs, while LLM-think models generate long chains of thought before final answers.The two model families share an underlying architecture but differ in internal representations and emergent behaviors.
  • Model pruning: Structured pruning removes entire coupled structures, unlike unstructured pruning, which removes individual weights by magnitude or importance.The background contrasts these two broad pruning approaches.
  • Pruning formulations: Static width pruning reduces hidden dimension d to d′ by removing less important neurons, whereas static depth pruning retains only a subset of layers.The retained-neuron formulation uses I_l, while retained layers form L′ ⊆ {1, . . . , L}.
  • Pruning criteria: Static depth pruning permanently removes layers selected using fixed importance scores, including Block Influence and perplexity criteria.Block Influence uses cosine similarity between layer inputs and outputs; high similarity indicates little newly contributed information.
  • Pruning formulations: Dynamic depth pruning uses a router and binary gate to execute or skip a block for each input token.The gate sets g = 1 for execution and g = 0 for skipping, with the block output updated accordingly.
  • Problem setting: The pruning objective maximizes compression ratio while minimizing degradation in performance retention relative to the dense baseline.Compression ratio is based on the reduction in average active parameters per token, while performance retention is Perf(M′)/Perf(M).

3 Experimental setup

The study compares pruning methods for instruction-following and reasoning models under a shared Llama-3.1-8B backbone and evaluates them across 17 task benchmarks. It covers static depth, static width, and dynamic pruning strategies, including representative methods from each category.

  • Models and evaluation: Both LLM-instruct and LLM-think share the same Llama-3.1-8B backbone for controlled pruning comparisons.Calibration and post-pruning recovery use each model’s accessible training data rather than downstream task data.
  • Models and evaluation: The evaluation spans 17 tasks covering instruction following, classification, generation, and reasoning capabilities.LLM-instruct is assessed on instruction-following benchmarks, while LLM-think is assessed on reasoning benchmarks.
  • Pruning methods: Static depth pruning removes entire transformer layers, with SLEB, ShortGPT, Shortened-PPL, and Shortened-Taylor evaluated as representative methods.These methods use layer-importance, representation similarity, perplexity, or Taylor-based criteria.
  • Pruning methods: Static width pruning reduces model width by pruning attention heads or shrinking hidden-state dimensions, using LLM-Pruner and SliceGPT.LLM-Pruner prunes coupled structures with gradient-based criteria, while SliceGPT removes low-variance components using principal component analysis.
  • Pruning methods: Dynamic pruning methods adapt computation per input through token or layer routing, including MOD, D-LLM, and SkipGPT.Their routing mechanisms selectively process tokens, skip transformer layers, or prune MLP and self-attention modules.

4 Effects of Calibration/Recovery Data

The study first tests whether calibration and post-pruning recovery data should match each model’s original training distribution. Across both model paradigms, aligned data best preserves capabilities, especially on complex tasks.

  • Pipeline design: Calibration estimates structural importance, while post-pruning recovery adapts the pruned model to restore performance.The study treats both stages as critical components of the pruning pipeline.
  • Experimental comparison: The experiments compare general-domain, instruction-tuning, and original training corpora for ShortGPT and SliceGPT at a fixed 20% pruning ratio.Original corpora include tulu-3-sft-mixture for LLM-instruct and openthoughts for LLM-think.
  • Results: Using the original training corpus for both calibration and recovery yields the strongest performance for LLM-instruct and LLM-think.This result is reported across IFEval, HumanEval+, classification tasks, and GSM8K with CoT prompting.
  • Results: Mismatched datasets preserve basic linguistic competence but can produce near-zero performance on challenging tasks, whereas original-corpus alignment preserves strong performance across tasks.The contrast appears for both instruction-following and reasoning models.
  • Conclusion: Aligning calibration and recovery with the original training corpus consistently yields superior preservation of pre-pruning capabilities.The finding supports using matched model–dataset pairs when comparing pruning strategies.

5 Effects of Pruning

Pruning effects differ sharply between LLM-instruct and LLM-think. Static depth pruning is stronger for classification, width pruning is more robust for generation and reasoning, and static methods generally preserve reasoning performance better than dynamic methods.

  • Controlled pruning setting: Aligning calibration and recovery data with the original training distribution is used as the default because it better recovers pruned-model performance.The study treats dataset alignment as a control for isolating pruning-strategy effects.
  • Static depth pruning: Shortened-Taylor consistently outperforms competing static depth methods for LLM-instruct at 20% and 40% pruning.Its average generation-task decline remains below 7.04% at 20% pruning.
  • Static depth pruning: For LLM-think, all static depth methods exceed 42% performance decline even at 20% pruning, indicating severe sensitivity to layer removal.The reported result suggests static depth pruning is not ideal for LLM-think.
  • Static width pruning: SliceGPT consistently outperforms LLM-Pruner across pruning ratios on both model paradigms, especially on generation and reasoning tasks at higher sparsity.The paper attributes this robustness to subspace-based compression rather than structured component deletion.
  • Static versus dynamic pruning: For LLM-think, dynamic methods often converge to higher training loss, whereas static methods remain lower and more stable; MOD is a notable exception.The contrasting optimization behavior is shown at a 20% pruning ratio.
  • Depth versus width: At 40% pruning, depth pruning declines by 55.57% for LLM-instruct and 85.38% for LLM-think, versus 41.83% and 77.69% for width pruning.Width pruning is therefore more robust at higher sparsity, particularly for generation and reasoning.
  • Depth versus width: Preserving model depth becomes more important than preserving width at higher sparsity levels, especially for generation and reasoning tasks.This structural interpretation summarizes the reported robustness gap between depth and width pruning.
  • Static versus dynamic pruning: Dynamic pruning performs best for LLM-instruct, but static pruning retains higher performance for LLM-think.The broader comparison reverses across paradigms, with dynamic methods benefiting classification and generation but struggling on reasoning.

6 Conclusion

The study finds that pruning strategies transfer unevenly from LLM-instruct to LLM-think, with task- and paradigm-dependent trade-offs. Dynamic methods perform better on classification and generation, while static pruning better preserves reasoning performance.

  • Static depth pruning is more effective for classification, whereas static width pruning degrades more gracefully for generation and reasoning.
  • Dynamic pruning benefits classification and generation but is substantially harder to optimize for long-chain reasoning.
  • Aggressive compute skipping is risky for reasoning-intensive workloads.
  • The findings highlight an important gap in current dynamic pruning for LLM-think.

Limitations

The experiments are limited to Llama-8B–scale language models, so the reported pruning behaviors are not explicitly evaluated across model scales.

  • The study does not explicitly evaluate whether its pruning behaviors generalize across model scales.
  • Extending the analysis to smaller and larger models remains an important direction for future work.

A The Use of Large Language Models

This section describes structured pruning and its main depth, width, and dynamic variants, alongside representative methods and their computational trade-offs. It also contrasts static component removal with input-dependent computation paths.

  • Structured pruning: Structured pruning removes entire coupled components rather than individual weights, making it more suitable for practical hardware acceleration.
  • Static depth pruning: Static depth pruning removes fixed Transformer layers or internal submodules, reducing sequential computation through a shorter model.
  • Static width pruning: Static width pruning removes redundant attention heads or MLP neurons while typically preserving the residual-stream dimensionality.
  • Dynamic pruning: Dynamic pruning selects the computational path separately for each input by executing only essential components at inference time.
  • Dynamic pruning: Early exit can accelerate inference by terminating simple inputs early, but may compromise deep semantic reasoning.
  • Dynamic pruning: Layer skipping uses routers to decide whether layers execute, with vertical dynamics varying computation across depth and horizontal dynamics across sequence positions.

C Experiments Details

The experiments use matched model and dataset setups for instruction-following and reasoning-augmented models, with pruning and recovery procedures evaluated under specified training configurations. Additional experiments examine calibration and post-training dataset effects for ShortGPT and SliceGPT.

  • LLM-instruct: For LLM-instruct, experiments use Llama 3.1-Tulu-3-8B-SFT with its SFT corpus for calibration and subsequent LoRA fine-tuning.
  • LLM-instruct: The LLM-instruct baselines are trained for 10,000 steps with batch size 16, cosine learning-rate scheduling, and maximum sequence length 4,096 tokens.Learning rates are selected by grid search for each baseline.
  • LLM-think: For LLM-think, the authors fine-tune Llama-3.1-8B-Instruct on OpenThoughts to create a reasoning-augmented model.
  • LLM-think: LLM-think pruning experiments use a single A800 GPU with DeepSpeed Stage-2 offloading and gradient checkpointing.
  • Dataset validation: Varying ShortGPT calibration datasets produces only minor differences in estimated layer importance, while the pruned-layer set remains unchanged.
  • Dataset validation: Figure 7 reports experiments examining the effect of post-training datasets for ShortGPT on LLM-instruct and LLM-think.

D Additional Empirical Findings on Calibration and Post-training Datasets

Calibration behavior depends on both calibration-set size and dataset choice, while post-pruning recovery is most effective on the original training corpus for both model paradigms.

  • Calibration-set size: SliceGPT performance rises sharply from 128 to 512 calibration examples and plateaus beyond 1024.This forms a saturation trend in the usable calibration signal.
  • Calibration datasets: SliceGPT is highly sensitive to calibration-dataset choice, whereas ShortGPT shows substantially more stable behavior.
  • Post-training datasets: Post-pruning recovery is most effective on the original training corpus for both LLM-instruct and LLM-think.The pattern holds consistently for both SliceGPT and ShortGPT.

E Repetitive Generation in LLM-Think

Dynamic depth pruning methods can produce severe repetitive and incoherent generation even at a 20% pruning ratio, with failures appearing across mathematical and reasoning examples.

  • Failure cases: At a 20% pruning ratio, D-LLM and SkipGPT exhibit severe semantic degradation with repetitive and incoherent outputs.The failures persist until the maximum generation length is reached.
  • MATH500 example: On the MATH500 example, D-LLM repeats an incorrect divisor count while SkipGPT produces repetitive, incoherent problem text.
  • GPQA example: D-LLM repeatedly outputs the same 0 eV conclusion instead of completing the energy-difference reasoning task.
  • GPQA example: SkipGPT produces fragmented repetition and incoherent number sequences on the GPQA example.
Loading 2601.18091v1…