PaperlayerView exact revision on arXiv

Source-linked AI summary

ReCal: Calibrating Structured Pruning for On-Policy Distillation Recovery

Houcheng Jiang, Mao Zheng, Mingyang Song, Qiyong Zhong, Jie Sun, Tianyu Zhang, Junfeng Fang

arXiv:2610.11332v1cs.CL

TL;DR

Pruning can damage reasoning, while OPD relies on student-generated trajectories that may preserve flawed prefixes. RECAL uses teacher–probe forward KL to reweight calibration before pruning; it improves recovered mathematical reasoning and benefits code generation in most comparisons. The results support accounting for recovery when selecting structures for pruning.

  • Problem

    Pruning can severely degrade reasoning, and OPD supervises student-generated prefixes where flawed generations can compromise feedback.

    Source: §1 · p. 1

  • Method

    RECAL measures teacher–probe forward KL on teacher trajectories and uses the scores to reweight calibration statistics before selecting pruned channels.

    Source: §1 · p. 2

  • Results

    RECAL improves mathematical reasoning after OPD on both models, with gains in seven of eight model–criterion comparisons and up to 16.7 percentage points, and benefits code generation in most comparisons.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The results suggest pruning should account for what the resulting student can recover, not only what it preserves immediately after compression.

    Source: §8 · p. 10

    Experiments focus on FFN-width pruning for two Qwen3 models and a fixed SFT–OPD pipeline; transfer to other structures and model families remains to be studied.

    Source: §7 · p. 9

Abstract

from arXiv · show

Structured pruning reduces the deployment cost of reasoning language models, but the resulting capability degradation can hinder subsequent on-policy distillation (OPD) recovery. Because OPD relies on student-generated trajectories, pruning damage that persists after offline distillation can limit its effectiveness. We propose RECAL, Recovery-Aware Calibration, a simple plug-and-play approach that improves OPD recovery by adjusting calibration before pruning. RECAL uses forward KL between an unpruned teacher and a pruned probe to identify teacher-supported predictions disrupted by pruning, then reweights calibration statistics to guide existing pruning criteria toward preserving these predictions. Across multiple models and pruning methods, RECAL consistently improves mathematical reasoning after OPD, achieving gains of up to 16.7 percentage points on AIME, alongside improvements in most code-generation comparisons. Further analysis shows that RECAL reduces residual damage at heavily affected tokens and establishes performance advantages that persist through recovery. These results demonstrate the value of recovery-aware calibration for improving on-policy distillation recovery of pruned reasoning models.

1 INTRODUCTION

Pruning can damage reasoning, and offline SFT only partly restores accuracy; OPD may then supervise student-generated prefixes where flawed reasoning persists. RECAL adjusts calibration before pruning to improve recovery, with reported gains in mathematical reasoning and code generation.

1 INTRODUCTION

  • Pruning can severely degrade reasoning, offline SFT only partly restores accuracy, and flawed student prefixes can compromise OPD feedback.A teacher may agree with next-token predictions at a flawed prefix, so low reverse KL on student rollouts need not indicate successful reasoning.Source: §1 · p. 1
  • RECAL measures token-level forward KL between a teacher and pruned probe on teacher trajectories, then reweights calibration statistics to recompute channel importance before applying a new mask.It preserves the base pruning criterion and target widths while changing the evidence used to select channels.Source: §1 · p. 2
  • RECAL improves mathematical reasoning after OPD on both Qwen3 models, with gains in seven of eight model–criterion comparisons and up to 16.7 percentage points, while also benefiting code generation in most comparisons.The comparisons are against representative structured-pruning baselines.Source: §1 · p. 2

2 PRELIMINARIES

The preliminaries introduce calibration-based structured pruning and the fixed SFT–OPD recovery pipeline that follows pruning. They define channel selection from calibration evidence and describe OPD training on student-generated responses.

2.1 STRUCTURED PRUNING

  • In FFN width pruning, selected channels and their associated input and output weights are removed, while a binary mask specifies which channels remain.Calibration-based criteria aggregate token contributions and retain the highest-scoring channels; RECAL reweights token evidence while preserving the criterion and target widths.Source: §2.1 · p. 3

2.2 ON-POLICY DISTILLATION RECOVERY

  • Recovery first fits fixed teacher responses through offline SFT, then uses OPD to train on student-policy responses under frozen-teacher supervision.Fresh responses are collected as the student updates, so the teacher supervises prefixes reached by its current policy.Source: §2.2 · p. 3
  • The retained structure determines the starting student, while the recovery process is held fixed to study whether pruning damage can inform structure selection.Source: §2.2 · p. 3

3 UNDERSTANDING PRUNING DAMAGE AND RECOVERY

Pruning damage is concentrated in a small subset of tokens, and tokens most damaged initially remain disproportionately discrepant after SFT and OPD. This persistent concentration motivates using pre-recovery damage to guide structure selection.

3 UNDERSTANDING PRUNING DAMAGE AND RECOVERY

  • The highest-damage 20% of tokens contribute 81.5–82.9% of within-response damage across three pruning criteria.The second model shows the same pattern, with shares of 80.9–83.1%.Source: §3 · p. 3
  • In the fixed highest-damage group, mean KL falls from 2.110 after pruning to 0.705 after SFT and 0.627 after OPD, yet post-OPD KL remains nearly three times the next-highest group’s 0.214.The full pipeline removes 70.3% of the group’s initial discrepancy, with most reduction during SFT.Source: §3 · pp. 3–4, Figure 2 · p. 4
  • Without reselecting tokens, the initially highest-damage group carries 64.6% of within-response damage after SFT and 69.3% after OPD.Its absolute KL decreases from SFT to OPD even as its share of remaining discrepancy rises, making initial damage a pre-training signal for locating tokens that remain poorly matched.Source: §3 · pp. 3–4

4 RECOVERY-AWARE CALIBRATION FOR PRUNING

RECAL locates pruning damage on teacher-generated prefixes, then uses those token-level scores to reweight calibration for structure selection. It applies the resulting mask to the original model and leaves the recovery procedure unchanged.

Locate damage

  • A pruned probe and teacher evaluate identical teacher-generated prefixes to estimate forward-KL damage, without letting the probe restrict the final mask.Teacher trajectories expose damage independently of the probe’s own rollouts.Source: §4 · p. 4

Calibrate & prune

  • RECAL converts probe damage into calibration weights, assigning greater weight to more-damaged tokens while keeping each response’s total weight equal.When total damage is negligible, it uses uniform weights.Source: §4 · p. 4
  • Weighted statistics are computed on the original unpruned model, so probe-removed channels remain eligible for the final mask.The selected mask is applied to the original weights.Source: §4 · p. 4

OPD recovery

  • RECAL preserves the base criterion and target widths, then recovers the selected student with the baseline’s unchanged SFT and OPD procedure.Non-additive criteria incorporate the weights into activation statistics or token losses before aggregation.Source: §4 · p. 5, Figure 3 · p. 5

5 EXPERIMENTS

Across the experiments, RECAL is tested on two Qwen3 models, four structured-pruning criteria, recovery stages, compression ratios, and several capability benchmarks. It improves post-OPD mathematics across most comparisons and shows how its benefits and trade-offs vary by criterion, compression, and selected channels.

5.1 EXPERIMENTAL SETUP

  • The setup compares Qwen3-8B in thinking mode and Qwen3-4B-Instruct-2507 without thinking, using four pruning criteria and matched widths at each compression ratio.Students receive SFT followed by OPD; evaluation covers AIME, GPQA-Diamond, IFBench, and LiveCodeBench-v5.Source: §5.1 · p. 5

5.2 RECOVERY PERFORMANCE (RQ1)

  • 48.5% versus 41.7% AIME accuracy in 8B and 37.2% versus 32.4% in 4B after OPD: RECAL exceeds Wanda, the strongest baseline.The comparison is at approximately 25% parameter reduction; RECAL also leads in each AIME year for both models.Source: §5.2 · p. 6, Table 1 · p. 6

5.3 PLUG-AND-PLAY IMPROVEMENT (RQ2)

  • RECAL improves mathematics in seven of eight model–criterion pairs, with gains up to 16.7 points; LiveCodeBench improves in seven pairs, while IFBench declines in all eight.Figure 4 compares Base + RECAL with Base across models and criteria; the intervention changes calibration evidence while holding scoring rules, widths, and recovery fixed.Source: §5.3 · p. 7, Figure 4 · p. 7

5.4 SENSITIVITY TO COMPRESSION RATIO (RQ3)

  • RECAL exceeds Base in all 18 measured SFT comparisons across compression ratios, while all measured SFT curves decline as compression increases.At 35% reduction, 8B Wanda + RECAL retains 27.1% accuracy versus 12.5% for Base; calibrated Minitron reaches 3.6%.Source: §5.4 · p. 7, Figure 5 · p. 8

5.5 STRUCTURE AND RECOVERY ANALYSIS (RQ4)

  • At fixed widths, activation-based RECAL pairs retain 75.8–79.3% of baseline deletions with ranking correlations of 0.784–0.839, while reducing Q5 post-OPD KL to 0.470 versus 0.627 for Base.Across eight pairs, the mean AIME advantage is 7.9 points after SFT and 7.7 after OPD; Q1/Q2 intervals include zero.Source: §5.5 · pp. 7–8, Figure 5 · p. 8

6 RELATED WORK

Prior work covers structured pruning, recovery through training or distillation, and approaches that adapt structure selection or recovery supervision. RECAL differs by using token-level teacher–probe damage to reweight calibration while retaining the base criterion and widths.

Structured pruning and capability recovery

  • Structured pruning uses gradient, activation, or reconstruction criteria, and continued training or distillation can recover pruned models.DarwinLM evaluates structures after brief training, while NIRVANA models fine-tuning dynamics and uses KL to select calibration data (Tang et al., 2026; Ai et al., 2026).Source: §6 · p. 8
  • Reasoning-oriented pruning methods use generated traces, preserve sampling diversity, or improve student initialization; RECAL instead reweights token contributions within teacher trajectories.RECAL preserves the base criterion and widths while changing calibration evidence for student selection.Source: §6 · p. 8

Offline and on-policy distillation

  • Distillation methods modify supervision or sampling, whereas RECAL uses forward KL on teacher prefixes to select the starting structure and leaves the OPD loss unchanged.Related approaches include skew divergences, adaptive reuse, mixed teacher–student responses, and short-to-long rollouts (Ko et al., 2024; 2025; Xu et al., 2025; Zhang et al., 2026).Source: §6 · pp. 8–9

7 LIMITATIONS

The experiments cover FFN-width pruning on two Qwen3 models with a fixed SFT–OPD pipeline, and RECAL’s benefits are not uniform across criteria or tasks. The paper also identifies unresolved questions about how its calibration choices relate to downstream recovery.

Generalization beyond the current pruning setting

  • Experiments are limited to FFN-width pruning on two Qwen3 models followed by a fixed SFT–OPD recovery pipeline.The authors identify attention-head and layer pruning, other model families, and interactive agent workloads as settings for future study.Source: §7 · p. 9
  • Mixed instruction-following results and the 8B LLM-Pruner case show that RECAL does not benefit every criterion or task uniformly.Its effectiveness may depend on the structure removed and the capability evaluated.Source: §7 · p. 9

Understanding what makes recovery-aware calibration effective

  • The complete procedure is evaluated without independently isolating trajectory choice, divergence direction, or token weighting.The authors say the connection between reduced residual discrepancy on heavily pruned tokens and downstream recovery remains only partially understood.Source: §7 · p. 9
  • Current diagnostics use fixed teacher prefixes rather than following student states encountered during OPD.The authors propose studying student-state diagnostics to clarify when recovery-aware calibration is useful.Source: §7 · p. 9

8 CONCLUSION

The conclusion presents RECAL as a way to reweight calibration evidence using teacher–probe disagreement before structure selection. It reports improved recovered mathematical reasoning and reduced residual damage at vulnerable tokens.

8 CONCLUSION

  • RECAL uses teacher–probe disagreement to reweight calibration evidence before structure selection.Source: §8 · p. 10
  • RECAL improves recovered mathematical reasoning and reduces residual damage at vulnerable tokens.Source: §8 · p. 10
  • The results suggest pruning should account for what a student can recover after compression, not only what it preserves immediately.Source: §8 · p. 10

AI USE STATEMENT

The authors report using generative AI tools for coding and writing assistance. They state that the tools were not used for research design, interpretation, or generation of results, and that outputs and results were reviewed by the authors.

AI USE STATEMENT

  • Generative AI assisted with implementation, debugging, refactoring, documentation, data processing, visualization, and writing edits.Source: AI USE STATEMENT · p. 10
  • The tools were not used for research ideation, methodological or experimental design, result interpretation, or generating experimental results.Source: AI USE STATEMENT · p. 10
  • Authors reviewed AI-assisted outputs and code, tested the code, independently verified reported results, and take responsibility for the manuscript and artifacts.Source: AI USE STATEMENT · p. 10

ETHICS STATEMENT

The study evaluates compression and recovery of existing language models using public benchmarks and teacher-generated responses. The authors caution that recovery gains do not establish deployment safety, fairness, or reliability.

ETHICS STATEMENT

  • The study examines existing language models using public benchmarks and teacher-generated responses.Source: ETHICS STATEMENT · p. 10
  • Reported recovery gains do not establish safety, fairness, or reliability in deployment.Source: ETHICS STATEMENT · p. 10
  • Instruction-following and scientific reasoning do not improve uniformly, and compressed students retain failure modes of pruning and their teachers.Source: ETHICS STATEMENT · p. 10
  • The authors call for separate capability and safety evaluation of intended applications and use under applicable licenses.Source: ETHICS STATEMENT · p. 10

REPRODUCIBILITY STATEMENT

The reproducibility materials document the models, pruning adapters, calibration and recovery data, training settings, and evaluation procedures. RECAL’s calibration uses teacher responses and damage-weighted statistics, while comparisons hold widths and recovery settings fixed.

REPRODUCIBILITY STATEMENT

  • The code and source bundle provide aggregate measurements and scripts, while the appendices document experimental settings and theoretical assumptions.Sections 4 and 5.1 describe the method and evaluation comparisons; Appendix A covers experiments, and Appendix C gives assumptions and proof.Source: REPRODUCIBILITY STATEMENT · p. 10

A EXPERIMENTAL DETAILS

  • Appendix A specifies model and adapter details, training data, recovery settings, benchmark evaluation, and fixed-prefix analysis.The experimental-details section maps these topics to Sections A.1–A.6.Source: A EXPERIMENTAL DETAILS · p. 15

A.1 MODELS AND PRUNED ARCHITECTURE

  • Pruning removes gated-FFN channels while retaining attention, embeddings, normalization, and the language-model head, with equal retained widths across layers.Removing a channel deletes its gate and up rows and down-projection column; widths align to multiples of 128.Source: §A.1 · p. 15

A.2 PRUNING CRITERIA AND CALIBRATION

  • Compared adapters differ in channel scores, while RECAL uses teacher responses and damage weights to reweight calibration statistics under shared target widths.The baseline uses prompt cross-entropy; RECAL weights response-token loss, and each response receives equal total weight.Source: §A.2 · pp. 15–16

A.3 TRAINING DATA

  • SFT and OPD use separate, initially balanced prompt pools, and OPD consumes prompts rather than source answers or external task rewards.Teacher-generated SFT targets are shared across methods; deduplication and benchmark-overlap screening are applied to prompts.Source: §A.3 · p. 16

A.4 RECOVERY HYPERPARAMETERS

  • Recovery uses SFT followed by 200 OPD updates, with fresh student responses evaluated by the teacher at student-generated prefixes.OPD applies a student-Top-16, restricted-support reverse-KL-style signal as an advantage in clipped policy updates, without GAE or external reward; paired runs share data and optimization settings.Source: §A.4 · p. 17

A.5 BENCHMARKS AND EVALUATION

The supplementary material specifies benchmark and fixed-prefix diagnostic protocols, then reports additional adapter comparisons, damage and generation analyses, and limits on interpreting KL-based preservation. It also distinguishes RECAL’s calibration signal from other recovery-aware pruning and distillation approaches.

A.5 BENCHMARKS AND EVALUATION

  • AIME uses eight sampled responses per question averaged across three years, while GPQA, IFBench, and LiveCodeBench report multiple-choice accuracy, instruction-level strict accuracy, and pass@1, respectively.Evaluation uses a common 20,480-token response budget; AIME uses temperature 0.6 and top-p 0.95, while other benchmarks use greedy decoding.Source: §A.5 · pp. 17–18

A.6 FIXED-PREFIX DIAGNOSTIC PROTOCOL

  • Fixed-prefix diagnostics rank raw-pruned KL within each response into five equal-count groups and track them through recovery, using 128 unfiltered responses per model and paired trajectory-bootstrap intervals.Concentration statistics average per-response ratios; the diagnostic pool is evenly split across four domains, and tokens within a response are not treated as independent samples.Source: §A.6 · p. 18

B.4 INTRODUCTION FIGURE MEASUREMENT; B.5 ADDITIONAL BENCHMARK RESULTS

  • The introduction diagnostics measure AIME accuracy and response-budget truncation separately; truncation cannot identify failure type, and Tables 5–6 add calibrated Pruned/SFT and post-OPD science and instruction scores.The supplementary tables report scores at 25% parameter reduction and complement the main-paper comparisons.Source: §B.4 · p. 19, §B.5 · p. 20, Table 5 · p. 21

B.1 LLM-PRUNER INTEGRATION

  • 8B LLM-Pruner AIME falls from 12.4% to 9.2%, while GPQA rises from 34.8% to 44.9% and code from 17.5% to 28.3%.The mathematics decline occurs in each AIME year; OPD narrows the post-SFT mathematics deficit but does not reverse it. The 4B comparison improves on AIME, GPQA, and code while IFBench declines slightly.Source: §B.1 · p. 19, Figure 8 · p. 20

B.2 ADDITIONAL DAMAGE ANALYSIS

  • The 4B concentration pattern extends the earlier analysis, while 8B Wanda shows increased discrepancy in some lower-damage groups after pruning and SFT despite reduced Q5 damage.Figure 9 supplies the additional concentration curves and paired KL changes before and after SFT.Source: §B.2 · p. 19, Figure 9 · p. 20

B.3 GENERATION AND ALIGNMENT DIAGNOSTICS

  • All seven post-OPD completion-diagnostic pairs show positive net rescue, but RECAL’s aggregate truncation rate rises from 76.8% to 79.4%.The 8B comparisons enter OPD with 1.9–3.6 percentage points higher teacher–student Top-16 overlap and finish 1.0–2.9 points higher; these own-rollout statistics do not isolate agreement at matched states. Across probes, response length correlates with summed KL (0.783–0.886), while high instruction-response damage does not consistently predict IFBench improvement.Source: §B.3 · p. 19, Figure 10 · p. 21

C FORWARD KL AND TEACHER-TRAJECTORY PRESERVATION; C.2 INTERPRETATION AND SCOPE

  • The teacher-trajectory connection motivates teacher-generated prefixes and teacher-weighted disagreement, but does not establish that KL weighting minimizes D(m), orders task accuracy, or certifies recovery-stage mask ordering.Tail-token coarsening can only decrease KL, so the implemented coarsened KL is not an upper bound on full-vocabulary damage and cannot certify sequence preservation.Source: §C.2 · p. 22

D ADDITIONAL RELATED-WORK COMPARISONS

  • RECAL’s distinction from prior work is its pre-recovery teacher–probe token signal for calibration with unchanged SFT–OPD, rather than recovery-aware pruning in general or forward KL alone.Experiments retain layers and attention modules and compare FFN channels at identical widths, so they do not establish the best compression axis or deployment speedups.Source: D ADDITIONAL RELATED-WORK COMPARISONS · p. 23
Loading 2610.11332v1…