PaperlayerView exact revision on arXiv

Source-linked AI summary

Thinking at the Right Size: Amortized Distillation Across Post-Trained LLMs

Yan Zhou, Sara Kangaslahti, Jonathan Geuter, Nihal V. Nayak, Marco Fumero, Francesco Locatello, David Alvarez-Melis

arXiv:2608.22854v1cs.LG

TL;DR

LLM families need models spanning size and post-training variants, but producing every combination independently is infeasible. ADAPT amortizes distillation across both axes, and its interpolated families outperform matched-compute baselines while enabling adaptive model-size selection.

  • Problem

    Producing models across both size and post-training variant axes is computationally infeasible, leaving families with only a few coarse-grained options and motivating questions about zero-shot cross-variant size interpolation.

  • Method

    ADAPT combines two-phase distillation with weight-delta initialization to construct interpolated post-trained model families across sizes and variants from a single distillation run.

  • Results

    ADAPT substantially outperforms boomerang distillation and layer-pruning baselines at matched compute on reasoning and instruction-following benchmarks, while recovering smooth size-performance interpolation.

  • Takeaways & Limitations

    The interpolated model continuum enables adaptive inference that routes easier inputs to smaller models and traces a compute-accuracy Pareto frontier dominating any single fixed model.

  • Takeaways & Limitations

    Weight-delta transfer is not uniformly reliable across model families, especially in some out-of-distribution settings and intermediate sizes, so it requires validation beyond the studied settings.

Abstract

from arXiv · show

Practical deployment of large language models (LLMs) requires families of post-trained variants---instruction-tuned, reasoning-tuned, and chat-style models---each at multiple sizes to meet diverse latency and memory budgets. Producing each (variant, size) pair independently is prohibitive, so model families typically span only a handful of coarse-grained sizes per post-trained variant. Boomerang distillation (Kangaslahti et al., 2026) reduces this cost along the size axis for base models. Through model size interpolation, it constructs models of intermediate sizes from a single teacher-student pair without additional training. However, it still treats each post-trained variant as a separate object of optimization. We introduce ADAPT---Amortized Distillation Across Post-Trained LLMs---a framework for amortizing distillation across both axes of a model family: size and post-training variant, producing $L \times K$ models for $L$ interpolated sizes across $K$ post-trained variants with a single distillation run. ADAPT combines two components. First, a two-phase distillation procedure constructs post-trained students through pre-training alignment and supervised fine-tuning distillation, enabling smooth size--performance interpolation on generation and reasoning tasks. Second, weight-delta initialization approximates this construction across post-trained variants by transferring the distillation-induced weight change from the base model to students initialized from different post-trained variants. The resulting continuum of interpolated models also enables adaptive model-size selection at inference time, improving the compute--accuracy trade-off for long-form reasoning tasks.

1 Introduction

ADAPT amortizes distillation across model size and post-training variant, producing interpolated post-trained model families from a single distillation run. It combines two-phase distillation with weight-delta initialization and supports adaptive inference through smooth size–performance interpolation.

  • ADAPT framework: ADAPT constructs L interpolated model sizes across K post-trained variants using one two-phase distillation run and transferred weight deltas.Two-phase distillation uses pre-training alignment and supervised fine-tuning distillation; the resulting delta initializes students from multiple post-trained teachers.
  • Motivation: Direct boomerang distillation and cross-patching fail because generation and reasoning performance collapses across small to medium sizes on in-domain and out-of-domain tasks.The post-trained teacher itself achieves strong performance, suggesting instruction-following depends on properties absent from the original boomerang distillation recipe.
  • Results: ADAPT recovers smooth size–performance interpolation on generation and reasoning benchmarks and outperforms boomerang distillation and layer-pruning baselines at equal compute.The framework is designed for model families spanning both size and post-training variant axes.
  • Adaptive inference: Adaptive inference routes easier inputs to smaller models, tracing a compute–accuracy Pareto frontier that dominates any single fixed model.The interpolated family enables model-size selection at inference time rather than relying on one fixed-size model.
  • Weight-space analysis: Base and post-trained models, together with their distilled students, are empirically connected by low-loss linear interpolation paths in weight space.This smooth weight-space geometry helps explain why a single base-model distillation can transfer across multiple post-trained variants.

2 Related Work

Related work spans knowledge distillation, post-training efficiency, model arithmetic, adaptive compute, and layer pruning. ADAPT combines these directions by amortizing student distillation across model variants and sizes while adapting inference compute through model-size selection.

  • Knowledge distillation: Knowledge distillation trains smaller models from larger teachers, but existing LLM families require individual distillation runs for each smaller model.Examples include Qwen, DeepSeek, and Gemma.
  • Post-training LLMs: Post-training methods reduce compute through smaller datasets, yet still require training models at every size.ADAPT instead distills a single post-trained student and creates additional sizes.
  • Model arithmetic: Model arithmetic combines model skills by adding or subtracting weights; ADAPT transfers the distillation-induced weight delta to amortize student distillation.Prior applications include instruction-following, multitask models, and machine unlearning.
  • Adaptive compute: Adaptive-compute methods reduce inference latency by matching computation to input difficulty, while ADAPT adjusts model size for reasoning and generation tasks.Related approaches include test-time scaling, fine-grained models trained from scratch, and heuristics such as appending thinking tokens.
  • Layer pruning: Layer pruning can preserve classification accuracy but may collapse open-ended generation; ADAPT uses targeted two-phase distillation to recover generation capability.The passage states that ADAPT substantially outperforms approaches such as ShortGPT and LLM-…

3 ADAPT: Amortized Distillation Across Post-Trained LLMs

ADAPT amortizes post-trained model construction across both size and variant axes through two complementary methods. ADAPT (distilled) uses two-phase distillation for interpolated models, while ADAPT (weight-delta) transfers the resulting distillation change across variants without another distillation run.

  • Boomerang distillation: Boomerang distillation constructs intermediate models without additional training by patching student layers with corresponding teacher-layer blocks.Students are initialized by partitioning teacher layers, followed by knowledge distillation and front-to-back or back-to-front patching.
  • ADAPT (distilled): ADAPT replaces boomerang’s original distillation stage with pre-training followed by supervised fine-tuning to recover language modeling, instruction-following, and reasoning capabilities.The pre-training phase precedes restoration of post-training capabilities, while supervised fine-tuning adapts students to downstream generation tasks and preserves teacher alignment.
  • ADAPT (distilled): The pre-training phase uses the boomerang distillation objective on pre-training data, whereas supervised fine-tuning computes the objective only over response tokens.The two phases respectively recover general language modeling and enable instruction-following adaptation.
  • ADAPT (weight-delta): Weight-delta initialization transfers a base-model distillation run to other post-trained variants, approximately constructing distilled students without an additional distillation run.The method uses weight arithmetic to transfer the distillation-induced change and aims to match directly distilled post-trained students in downstream performance.

4 Experiments

Experiments show that ADAPT’s two-phase distillation enables strong interpolation across model sizes, while weight-delta initialization transfers this behavior across post-trained variants without additional distillation. Adaptive inference-time size selection further improves the performance–compute trade-off, consistent with preserved smooth weight-space structure after distillation.

  • 4.2 Size interpolation: ADAPT (distilled) achieves the strongest interpolation performance on both in-domain and out-of-domain tasks by combining pre-training alignment with SFT capability recovery.It also achieves the best teacher–student alignment by last-layer activation cosine similarity.
  • 4.2 Size interpolation: Base model BD collapses completely, while BD (pre-train) remains weak at small and medium sizes and BD (SFT) generalizes poorly to out-of-domain tasks.Cross-patch (pre-train) recovers nontrivial performance as more post-trained teacher layers are inserted, motivating weight-delta transfer.
  • 4.3 Post-training variant transfer: ADAPT (weight-delta) reuses one base-model distillation run across post-trained variants, matching ADAPT (distilled) on many in-domain settings without additional distillation.It remains competitive on out-of-domain tasks, with slight degradation concentrated in thinking-style models at smaller and intermediate sizes, and outperforms cross-patch (2-phase).
  • 4.4 Adaptive inference: Adaptive model-size selection outperforms the single-model baseline, tracing a superior Pareto frontier across average and minimum accuracy for both ADAPT variants.Smaller models are assigned to easier questions and larger models to harder ones, enabled by smoothly increasing performance across interpolated models.
  • 4.5 Weight-space analysis: Interpolation paths between base and post-trained models show no sharp loss barriers for teachers or students, and students retain similar interpolation behavior after distillation.This suggests distillation preserves much of the weight-space structure connecting base and post-trained models.

5 Conclusion

ADAPT amortizes distillation across both model size and post-training-variant axes. It outperforms boomerang distillation baselines under matched compute and enables adaptive model-size selection through a continuum of interpolated models.

  • 5 Conclusion: ADAPT amortizes distillation across both the size and post-training-variant axes of a model family.This is the framework’s central contribution.
  • 5 Conclusion: ADAPT substantially outperforms boomerang distillation baselines on reasoning and instruction-following benchmarks under matched compute.The comparison is made under matched compute across both benchmark categories.
  • 5 Conclusion: The resulting continuum of interpolated models enables adaptive model-size selection at inference time.This capability is enabled by the continuum produced through ADAPT.
  • 5 Conclusion: Smooth linear weight interpolation between base and post-trained models is preserved across model sizes after distillation.The preservation of this interpolation property is a finding reported in the conclusion.

Limitations

ADAPT’s evaluation leaves open questions about interpolation failures, the reliability and prerequisites of weight-delta transfer, and the scope of its empirical validation. Adaptive inference also depends on task-specific difficulty estimation whose broader generalization and cost remain uncharacterized.

  • Limitations: Several aspects of ADAPT warrant further investigation, including its underlying mechanism, implicit procedural assumptions, and empirical scope.These open questions frame the specific limitations identified in the section.
  • Interpolation dips at specific sizes: Interpolation curves can show sharp drops at particular intermediate sizes, with the mechanism and prediction of these failure modes still unclear.A plausible explanation is local misalignment between inserted teacher layers and surrounding student layers.
  • Model and data scope: ADAPT’s experiments cover Qwen, Olmo, and Llama at 4B-14B scale using math and coding data for supervised fine-tuning, leaving larger scales and broader mixtures unverified.Multilingual instruction-tuning mixtures are given as an example of an untested broader scope.
  • Weight-delta assumptions: Weight-delta initialization requires access to the base checkpoint for each post-trained variant, limiting applicability when that checkpoint is unavailable or not exactly identifiable.This restriction affects some open-weight releases and most closed-weight releases without additional approximation.
  • Adaptive inference depends on a difficulty signal: Adaptive model-size selection relies on per-input difficulty estimates, but teacher-predicted difficulty in MATH is task-specific and broader generalization remains uncharacterized.The approach may be difficult to extend when difficulty is poorly defined or a separate predictor is expensive to obtain.

Ethical Considerations · A Training Implementation Details · B Hyperparameters

The study uses public datasets, benchmarks, and model checkpoints without collecting new human-subject data, while acknowledging inherited bias, harmful-content, and privacy risks. Training spans thinking and non-thinking model variants with scaled token budgets, reported hardware runtimes, Apache 2.0 release plans, and explicitly balanced distillation losses.

  • Ethical Considerations: Experiments rely on publicly available datasets, evaluation benchmarks, and model checkpoints, with no newly collected human-subject data.The datasets primarily concern mathematical reasoning, instruction following, and general model evaluation.
  • Ethical Considerations: Web-derived pre-training and post-training corpora may inherit biases, harmful content, and privacy concerns from their source material.The passage identifies these as limitations of the underlying data sources and models.
  • A Training Implementation Details: The experiments include both non-thinking and thinking models, including Qwen3, Olmo-3, and Llama-3.1 variants.Qwen3-4B and Qwen3-14B can toggle between thinking and non-thinking modes.
  • A Training Implementation Details: Training uses 1B tokens for Qwen 4B-sized models and scales to 2B and 4B tokens for larger models.The budgets are scaled proportionally following Hoffmann et al. (2022).
  • A Training Implementation Details: Non-thinking models use ground-truth responses without <think> tags, whereas thinking models use thinking ground truths version 1.1.The training data format is selected according to the model’s thinking configuration.
  • A Training Implementation Details: Each 4B run takes approximately 12 hours on 4 NVIDIA H100 GPUs, while larger or other-model runs use more GPUs and time.Olmo and Llama runs take approximately 24 hours on 4 NVIDIA H200 GPUs; Qwen3-14B takes approximately 48 hours on 16 H200 GPUs.
  • A Training Implementation Details: The authors plan to release the distilled models under an Apache 2.0 license.This licensing statement accompanies the reported training-resource requirements.
  • B Hyperparameters: Hyperparameters largely follow Kangaslahti et al. (2026a), with KL and cosine weights chosen to match early cross-entropy magnitude and one scheduler across distillation phases.M denotes the number of student layers.

C Evaluation Implementation Details … E Additional Ablations for ADAPT (Distilled)

The appendices define evaluation protocols and provide additional ADAPT results, showing that two-phase distillation best aligns students with teachers while classification is mostly strategy-insensitive except for BD (SFT).

  • C Evaluation Implementation Details: Generation evaluation averages accuracy across five benchmarks, separating AIME, GSM8K, and MATH500 as in-domain from IFEval and MMLU-Redux as out-of-domain.Coding uses a separate in-domain benchmark set, while Section 4.4 uses the MATH dataset.
  • C Evaluation Implementation Details: Classification evaluation uses lm-evaluation-harness and averages accuracy across 10 benchmarks, with distinct sampling parameters for non-thinking and thinking models.The benchmark set includes ARC, BoolQ, HellaSwag, MMLU, OpenBookQA, PIQA, RACE, RTE, and WinoGrande tasks.
  • D Additional Evaluation Results for ADAPT (Distilled): Additional ADAPT (distilled) analysis examines final-layer cosine similarity to teachers and classification performance.The cosine-similarity analysis averages activations over all downstream tasks.
  • D.1 Cosine Similarity Analysis: ADAPT (distilled) achieves the highest student-teacher cosine similarity across nearly the entire interpolation curve, with BD (SFT) consistently below it.Cross-patch (pre-train) has the weakest alignment, especially at smaller model sizes.
  • D.1 Cosine Similarity Analysis: The weakest alignment of cross-patch (pre-train), particularly at smaller sizes, indicates that patching alone does not sufficiently align students with post-trained teachers.These results further motivate the two-phase distillation approach.
  • D.2 Classification Performance: Classification accuracy is relatively insensitive to distillation strategy across Qwen, Olmo, and Llama, with most setups showing similar interpolation behavior.Classification tasks therefore do not fully expose the alignment challenges arising in post-trained generation.
  • E Additional Ablations for ADAPT (Distilled): A 50:50 pre-training-to-SFT distillation split balances in-domain and out-of-domain performance, whereas increasing either phase favors its associated task type.The single-phase endpoints are consistently dominated by two-phase settings.
  • D.2 Classification Performance: BD (SFT) is the main classification exception, underperforming across much of the interpolation curve for all three model families.The results suggest that task-specific distillation alone can degrade broader alignment with the teacher.

E.1 Pre-training vs. SFT Phase Ratios · E.2 Distillation with Cosine Loss Only · F Additional Ablations for ADAPT (Weight-Delta)

ADAPT’s two-phase training shows a trade-off between ID and OOD accuracy while remaining relatively insensitive to the precise phase split, and both two-phase settings outperform single-phase endpoints. Cosine-only distillation achieves performance comparable to the full objective, while pre-training weight-delta initialization is additionally examined for post-trained Qwen models.

  • E.1 Pre-training vs. SFT Phase Ratios: Increasing the SFT proportion improves ID accuracy but degrades OOD accuracy, whereas increasing the pre-training proportion has the opposite effect.
  • E.1 Pre-training vs. SFT Phase Ratios: The 50:50 phase split achieves strong ID accuracy without sacrificing OOD performance.
  • E.1 Pre-training vs. SFT Phase Ratios: The three two-phase ratios perform comparably at larger model sizes, indicating limited sensitivity to the exact split.
  • E.1 Pre-training vs. SFT Phase Ratios: Both single-phase endpoints are dominated by the two-phase settings.
  • E.2 Distillation with Cosine Loss Only: Cosine-only distillation yields performance comparable to the full objective combining cross-entropy, KL, and cosine losses.This result suggests intermediate representation alignment may drive interpolation behavior and could simplify the training objective.
  • F Additional Ablations for ADAPT (Weight-Delta): Pre-training weight-delta initialization is evaluated for post-trained Qwen models.

F.1 Weight-Delta Transfer from Pre-training-only and SFT-only Deltas … F.4 Ablating SFT Phase Loss Components

Across ablations, pre-training-derived weight deltas transfer more reliably than SFT-only deltas, while two-phase distillation remains strongest overall. Post-trained deltas and SFT-loss ablations clarify how transferability, continued training, and ID–OOD trade-offs depend on the distillation procedure.

  • F.1 Weight-Delta Transfer from Pre-training-only and SFT-only Deltas: Pre-training weight-delta initialization transfers comparably to or better than BD (pre-train) across model families, whereas SFT-only deltas transfer less smoothly.This indicates that alignment learned during base-model pre-training distillation transfers across post-trained variants more effectively than SFT-only alignment.
  • F.1 Weight-Delta Transfer from Pre-training-only and SFT-only Deltas: Single-phase weight-delta methods remain weaker than ADAPT (weight-delta) and ADAPT (distilled), showing that both distillation phases are needed to recover downstream capabilities.SFT-only transfer closely tracks BD (SFT) for instruct models but underperforms on the other three variants and even loses to cross-patch (SFT) for Qwen3-4B.
  • F.2 Continued Training from Weight-Delta Initialization: Continued training for 0.5B or 1B tokens does not meaningfully improve interpolation beyond ADAPT (weight-delta), regardless of whether initialization uses pre-training, SFT, or two-phase deltas.The results support zero-training transferability as the main benefit of weight-delta initialization rather than improved initialization for further training.
  • F.3 Weight-Delta Initialization with Post-trained Model Deltas: Post-trained instruct and thinking model deltas both substantially outperform the cross-patch (2-phase) baseline when transferred to other post-trained variants.The instruct delta is competitive with ADAPT (distilled) on ID tasks but shows slight OOD degradation similar to the base-model delta.
  • F.4 Ablating SFT Phase Loss Components: Removing the KL term leaves interpolation behavior largely unchanged, indicating that KL contributes little during the SFT phase.The ablation retains cross-entropy while removing KL, cosine, or both terms from SFT-phase distillation.
  • F.4 Ablating SFT Phase Loss Components: Removing the cosine term improves ID performance across all four variants to match ADAPT (distilled), but degrades OOD performance overall.The ID gains arise mainly in smaller interpolated models, while OOD performance declines at intermediate and larger sizes.
  • F.4 Ablating SFT Phase Loss Components: 3.3B parameters marks the approximate point until which the full objective leaves weight-delta students near total collapse in smaller interpolated models.Removing both KL and cosine terms produces similar behavior but with even weaker OOD performance, revealing a cosine-mediated ID–OOD trade-off.

G Weight-Delta Alignment Experiments … I.3 Discussion of Total Token Count

ADAPT’s weight-delta transfer reliability follows directional alignment, while both ADAPT variants outperform layer-pruning baselines across sizes. Adaptive selection uses teacher-estimated difficulty, but smaller models may generate more total tokens because unresolved problems trigger the token limit.

  • G Weight-Delta Alignment Experiments: The authors identify when weight-delta transfer succeeds as an important open question, noting its difficulty in related model-merging and task-vector settings.The preceding appendices show ADAPT (weight-delta) is more successful for some models and setups than others, matching ADAPT (distilled) at its best.
  • G Weight-Delta Alignment Experiments: Weight-delta transfer reliability tracks alignment: instruct and hybrid variants reach cosine 0.73-0.75, thinking variants 0.57-0.64, and Llama 0.30.Norm ratios remain near 1, suggesting differences in direction rather than magnitude; the analysis presents directional alignment as a candidate diagnostic, not a proven causal rule.
  • H.1 ShortGPT / H.2 LLM-Streamline: ShortGPT prunes layers with the lowest Block Influence score, whereas LLM-Streamline replaces the lowest-scoring layer block with a lightweight network trained to imitate it.LLM-Streamline additionally post-trains the replacement layer on language-modeling loss for a fairer comparison.
  • H Comparison to Layer Pruning Methods: Across all model sizes, ADAPT (distilled) and ADAPT (weight-delta) consistently outperform ShortGPT and LLM-Streamline on generation tasks.The layer-pruning baselines struggle to recover meaningful generation capability, especially at smaller sizes.
  • I.1 Input Difficulty Classification: Teacher difficulty estimates achieve MAE 0.970 versus student MAE 1.995, making the teacher more suitable for adaptive model selection.The teacher’s predictions are typically within one difficulty level of ground truth, while student predictions deviate by approximately two levels on average.
  • I.1 Input Difficulty Classification: Teacher-defined difficulty labels produce well-separated, consistently ordered interpolation curves resembling ground-truth trends, whereas student labels are noisy.This supports using the teacher for difficulty estimation in adaptive model selection.
  • I.2 Adaptive Model-Size Selection Results for Olmo and Llama: Qwen3-4B-Instruct-2507 supplies difficulty labels for Olmo-3-7B-Instruct and Llama-3.1-8B-Instruct because it is smaller than both teacher models and therefore cheaper.The adaptive model-size selection experiments for these variants are reported in Figures 24 and 25.
  • I.3 Discussion of Total Token Count: Smaller models often generate more total tokens despite favorable FLOPs per token because unresolved problems continue generating until reaching max_tokens.The authors believe further training could mitigate this issue but leave it for future work.

J Patching Order Experiment … O Use of Large Language Models

The supplementary experiments show that ADAPT’s interpolation behavior depends on model architecture and training phase, while extending across additional models, coding tasks, and benchmark-level evaluations. The authors also disclose limited generative-AI use for manuscript preparation, with substantive research conducted by the authors.

  • J Patching Order Experiment: Patching-order effects vary by model: KLPatch improves Qwen3-4B-Instruct-2507 and Olmo-3-7B-Instruct, whereas front-to-back patching is best for Llama-3.1-8B-Instruct.The experiments use back-to-front patching for Qwen and front-to-back patching for Olmo and Llama; KLPatch uses a 64-example Nemotron non-thinking calibration set.
  • K Additional Smooth Weight Interpolation Results: Base and post-trained students exhibit smooth interpolation in generation accuracy and SFT loss across all three setups.The loss curves differ partly because calibration data produce different losses for base, instruct, and thinking models.
  • L Training Dynamics: Two-phase distillation converges to distinct pre-training and SFT interpolation solutions, combining general-purpose alignment with improved in-domain student performance.The pre-training phase resembles the BD (pre-train) curve, while SFT initially improves students, dips at larger sizes, and then recovers smooth interpolation.
  • M ADAPT for Additional Models and Tasks; M.1 Results for Qwen3-14B; M.2 Results for Olmo; M.3 Results for Llama: Additional experiments extend two-phase distillation and weight-delta initialization to Qwen3-14B thinking and non-thinking modes, Olmo models, and Llama-3.1-8B-Instruct.Results are reported in Figures 31–37 for these model families and modes.
  • M.4 Results for Coding Tasks: On coding tasks, two-phase distillation outperforms BD (pre-train) and BD (SFT), while ADAPT (weight-delta) trades lower cost for performance between ADAPT (distilled) and cross-patch (2-phase).Coding evaluation replaces math SFT data with coding data and uses MBPP+, HumanEval+, and Live-CodeBench for ID tasks; Qwen3-4B results are evaluated in thinking mode.
  • N Individual Benchmark Results; N.1 Two-phase Post-trained Distillation Results; N.2 Weight-delta Initialization Results: The supplementary benchmark analyses report Qwen3-4B-Instruct-2507 interpolation results and Qwen-family weight-delta initialization results on individual benchmarks.These evaluations appear in Figures 40 and 41.
  • O Use of Large Language Models: Generative AI tools were used for code completion, debugging, and minor grammatical corrections, while authors performed the substantive research, analyses, and interpretations.The disclosure appears in the paper’s use-of-large-language-models section.
Loading 2608.22854v1…