Unlocking Structured Pruning for Reasoning Language Models via Attention-Preserving MLP Trimming

1 Introduction

• Pruning methods developed for non-reasoning LLMs can cause performance collapse on long-chain tasks such as AIME.The paper reports that existing structured-pruning methods produce longer reasoning traces and lower accuracy.Source: §1 · p. 1
• MLP trimming produces less representation error than attention pruning or module removal, and concentrating trimming in redundant layers produces less error than distributing the same budget across layers.The analysis also finds MLP feature importance concentrated in a subset of features.Source: §1 · p. 1, Figure 2 · p. 2
• LayerTrim uses dense-model reasoning traces and Block Influence to select redundant layers, retaining MLP features with the largest expected squared output contributions.It preserves attention and decoder depth while selectively trimming MLP features.Source: §1 · pp. 1–2
• Across three reasoning models and five benchmarks, LayerTrim outperforms existing structured-pruning methods in most settings, improving average accuracy by up to 21.4 percentage points over RESP at matched parameter budgets.On Qwen3-14B, it also achieves higher accuracy while generating fewer tokens on average than other evaluated pruning methods.Source: Figure 1 · p. 1, §1 · p. 2
• LayerTrim is training-free and designed to preserve long-form reasoning without gradients or recovery training.The paper describes it as the first training-free structured-pruning method designed for this goal.Source: §1 · pp. 1–2

2 Related Work

• Structured-pruning methods use gradient- or Taylor-based importance, error compensation, additional training, interaction-aware pruning, or layer removal, merging, and replacement.ShortGPT uses Block Influence to identify redundant layers, while LoRP uses inter-layer redundancy and architectural locality for training-free depth pruning.Source: §2 · p. 2
• Existing structured-pruning approaches do not explicitly account for long-form chain-of-thought reasoning, and removing entire layers can degrade reasoning performance.LayerTrim addresses this challenge by selectively trimming MLP features while preserving attention and decoder depth.Source: §2 · pp. 2–3
• Reasoning-focused methods include RAC and OBC-Prune, which use unstructured pruning and require sparse-kernel support for practical acceleration; structured RESP uses reasoning traces, gradient-based importance, and progressive recalibration.RESP retains lower reasoning accuracy than dense models across the evaluated benchmarks.Source: §2 · p. 2

3.1 Motivation

• Across tested budgets, MLP-only trimming yields lower final-hidden representation error than module removal, attention trimming, or joint attention-and-MLP trimming.This motivates preserving attention and decoder depth while reducing MLP width.Source: §3.1 · p. 3
• The top 38.2% of MLP features account for 89.4% of total feature importance, supporting their retention during width reduction.Feature importance is defined as the expected squared norm of a feature’s MLP output contribution.Source: §3.1 · p. 3
• Under matched parameter budgets, trimming 14 low-Block-Influence MLPs yields lower final-hidden representation error than trimming all 40 MLPs or 14 randomly selected MLPs.The random-selection comparison indicates that layer choice matters beyond the number of trimmed MLPs.Source: §3.1 · p. 3

3.2 LayerTrim

• LayerTrim gathers dense-model reasoning-trace statistics, selects low-Block-Influence layers, ranks MLP features, and trims selected MLPs while preserving attention and decoder depth.Its five stages run from reasoning-trace calibration through layer selection, feature ranking, and structured trimming to a compressed model.Source: §3.2 · pp. 3–4, Figure 3 · p. 4
• For a budget of N layer-equivalents, LayerTrim selects K = 2N layers and retains the top m features by score in each selected MLP.The feature count is set by equating removed parameters, 3CK(I − m), to the budget NPlayer.Source: §3.2 · pp. 3–4
• LayerTrim physically removes I − m features to create smaller dense MLPs, with no inference-time masks, fine-tuning, distillation, or recovery training.Attention modules, unselected MLPs, and decoder depth remain unchanged.Source: §3.2 · p. 4

4.2 Main Results

• At N = 9, LayerTrim exceeds the strongest baseline average by 6.8, 14.3, and 21.4 percentage points on Qwen3-14B, Phi-4-reasoning-plus, and Magistral-Small-2506, respectively.It has the highest average benchmark score among pruning methods for all three models at both matched budgets; its lead grows from N = 7 to N = 9.Source: Table 1 · p. 5, §4.2 · p. 5

4.1 Experimental Settings

• Using 24 dense-model AIME 2024 traces per model for calibration, LayerTrim gathers statistics without updating weights; evaluation uses native reasoning templates and three generation seeds.Models are evaluated against their original dense checkpoints for performance retention, on AIME, GPQA-Diamond, HumanEval, and MBPP+ benchmarks.Source: §4.1 · pp. 4–5

4.3 Ablation Studies

• MLP-only trimming yields the lowest final-hidden NMSE and higher average accuracy than module removal, attention pruning, or joint attention–MLP trimming.The comparisons motivate preserving attention and decoder depth while reducing MLP width.Source: §4.3 · p. 6
• LayerTrim’s feature-importance score achieves the highest average accuracy among evaluated criteria, outperforming activation-only ranking and FLAP’s variance-based score.Activation-based scores outperform weight-only ranking.Source: §4.3 · p. 6
• Selecting low-BI layers beats random selection and all-layer trimming in average accuracy; Block BI also beats MLP-only BI.This supports measuring redundancy across the full decoder block rather than only the MLP update.Source: §4.3 · p. 6
• Under a fixed budget, K = 1.5N yields the highest AIME25 accuracy, K = 3N the highest GPQA-Diamond accuracy, and default K = 2N the highest average accuracy by a small margin.The default provides a practical balance between the two tasks.Source: §4.3 · p. 6

4.4 Efficiency Measurement

• At the higher pruning budget, LayerTrim achieves 1.24×, 1.22×, and 1.27× prefill speedups on Qwen3-14B, Phi-4-reasoning-plus, and Magistral-Small-2506, respectively.Peak GPU memory falls by 19.7%, 20.3%, and 20.8% on those models, respectively, while LayerTrim retains the highest average accuracy across all three models at both budgets. On Qwen3-14B, pruning preparation takes 27.9 s, 94× and 140× faster than RESP at N = 7 and N = 9, respectively.Source: §4.4 · p. 6

5.1 Analysis of LayerTrim

• The feature-importance scores rank measured squared MLP output changes across held-out conditions with Spearman ρ = 0.999.The low-overhead approximation gives comparable average accuracy to full-error greedy selection with 13.2× faster calibration and selection.Source: §5.1 · p. 7
• MLP trimming yields the lowest attention KL divergence across tested budgets and the highest overlap with dense-model top-8 attended tokens, including at 16K–32K query–key distances.The passage interprets this as better preservation of attention to earlier tokens during long reasoning.Source: §5.1 · pp. 6–7, Figure 4 · p. 8

5.2 Quantitative Analysis of Reasoning Traces

• LayerTrim records 24.4% token-limit hits and 21.9% strong repetition, versus FLAP’s 70.2% and 49.2%, alongside higher accuracy.Figure 5 examines these output diagnostics for GPQA-Diamond responses from Phi-4-reasoning-plus at Lp/Lt = 9/40.Source: §5.2 · p. 7, Figure 5 · p. 8
• On HumanEval, LayerTrim achieves 99.6% compilable-code rate versus Dense’s 99.4%, and 88.4% correctness among qualifying-code responses versus RESP’s 72.7%.The passage also reports overall HumanEval accuracy of 88.0% for LayerTrim versus 62.0% for RESP.Source: §5.2 · p. 8
• On AIME26, LayerTrim reaches 94.8% arithmetic consistency versus Dense’s 96.4%; with RESP at 94.1%, its accuracy is 60.0% versus RESP’s 50.0%.The consistency comparison also reports 83.8% for FLAP and 73.6% for CoCurve.Source: §5.2 · p. 8

• LayerTrim selectively trims MLP features while preserving attention and decoder depth, using dense-model reasoning traces to identify redundant layers and retain high-importance features without gradients or recovery training.Source: §6 · p. 8
• Across three reasoning language models and five benchmarks, LayerTrim achieves higher average accuracy than evaluated structured-pruning baselines at both pruning budgets.Source: §6 · p. 8
• LayerTrim offers a step toward more efficient deployment of reasoning language models through training-free structured pruning.Source: §6 · p. 8

• Evaluation is limited to three reasoning language models and five benchmarks, and uses fixed pruning budgets and layer allocations.Additional architectures, scales, and reasoning tasks would help assess generalizability; adapting budgets and layer allocation to models and tasks may improve the compression–accuracy trade-off.Source: Limitations · p. 9

Ethical considerations

• Pruned models may still produce incorrect answers, incomplete reasoning traces, or unexpected outputs, a risk to consider where reasoning accuracy is critical.Source: Ethical considerations · p. 9

A LayerTrim Procedure

• Algorithm 1 calibrates on generated reasoning traces by teacher-forcing the dense model and accumulating block similarities and feature second moments.Calibration averages are taken over token positions.Source: A LayerTrim Procedure · p. 11
• The procedure selects layers by Block Influence, scores MLP features, installs selected projections at reduced width, and returns a model with unchanged attention and depth.The algorithm applies the width reduction to selected layers and keeps the highest-scoring features.Source: A LayerTrim Procedure · p. 11

B Main Evaluation Implementation Details

• All methods use the same target budget, defined as N times the prunable parameters in one original decoder layer, with N ∈ {7, 9}.Actual removal counts can differ slightly because heads and channels are removed in whole units.Source: B Main Evaluation Implementation Details · p. 11
• LayerTrim selects K = 2N MLPs and removes 3C parameters per removed gated-MLP feature, yielding a realized budget of 3CK(I − m).The selected MLPs share retained width m, set by Equation 3.Source: B Main Evaluation Implementation Details · p. 11
• Calibration uses 24 dense-model AIME 2024 reasoning traces; generation uses native reasoning-enabled templates, three seeds, and a 32K-token output limit.Scoring uses exact-match answers for AIME and GPQA-Diamond, and pass@1 for HumanEval and MBPP+.Source: B Main Evaluation Implementation Details · p. 11

B.1 Calibration Overhead

• LayerTrim requires substantially less preparation time than RESP at both pruning budgets (Table 4).Calibration time is measured from model loading until pruning masks are computed, including RESP trace regeneration and intermediate checkpoint preparation.Source: §B.1 · p. 11, Table 4 · p. 12

C Ablation Experimental Details

• The ablations compare trimming targets at approximately 1.494B removed parameters and fix the design-ablation budget at N = 7, approximately 2.312B parameters.Other variations test feature-ranking scores, layer selection, and the number of selected MLPs while retaining the specified budget.Source: C Ablation Experimental Details · p. 12

C.1 Additional Accuracy Ablations

• Accuracy is stable beyond 1,024 calibration tokens, and 2,048 tokens are used by default.The ablation varies generated reasoning-token length on the same 24 AIME 2024 traces, with the budget and K = 14 fixed.Source: §C.1 · p. 12
• Reasoning-trace calibration yields higher average accuracy than token-matched WikiText-2 or prompts alone.The comparison fixes Qwen3-14B at N = 7 and K = 14; prompt-only calibration uses fewer tokens.Source: §C.1 · p. 12, Table 6 · p. 12
• 68.1% average accuracy is achieved by both feature-importance ranking and greedy selection using squared MLP output change, while the latter takes 13.2× as long.The comparison holds Qwen3-14B at N = 7, with the same calibration traces, selected layers, and retained widths.Source: §C.1 · p. 13, Table 7 · p. 13

D Representation error and Attention Diagnostics Experimental Details

• Attention overlap measures the fraction of the dense model’s top-eight attended earlier tokens retained among the pruned model’s top eight, grouped by token distance.Higher overlap means pruning better preserves which earlier tokens the model attends to.Source: D Representation error and Attention… · p. 13
• Feature-importance scores strongly rank-agree with measured squared MLP output changes across 70 layer-and-width combinations.Feature importance omits cross-feature terms, while measured output change includes them.Source: D Representation error and Attention… · p. 13

E Qualitative Analysis

• The examples are not prevalence estimates, and lexical collapse may evade the exact-token repetition criterion despite frequent token-limit hits.The appendix examples illustrate repetition, incoherent fragments, code-format failures, and arithmetic errors.Source: E Qualitative Analysis · p. 13
• A ShortGPT coding response terminates after 337 tokens but claims correctness without supplying comparison logic.The example is a has_close_elements task whose prompt requires returning False for the given input.Source: E Qualitative Analysis · p. 14
• A ShortGPT AIME excerpt reaches 32,256 tokens and hits the limit while repeating an incorrect arithmetic calculation.The trace repeatedly states a division by 17017.Source: E Qualitative Analysis · pp. 15–16

F Licenses and Terms of Use

• The study uses publicly available models and datasets under their respective licenses and usage conditions, solely for non-commercial research.The listed models include Apache 2.0- and MIT-licensed releases; the listed datasets and software have varied licensing terms.Source: F Licenses and Terms of Use · p. 17