cs.LGAug 24, 2026
SaveThinking at the Right Size: Amortized Distillation Across Post-Trained LLMs
Organizations: Harvard University · Kempner Institute · IST Austria
Abstract
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 models for interpolated sizes across 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.
Figures & tables
Figure 1: ADAPT: Amortized Distillation Across Post-Trained LLMs. We perform two-phase distillation (pre-training + SFT) on a base student once, yielding interpolated model sizes from a single teacher–student pair. We then compute the distillation-induced weight delta and add it to the initialized students from multiple post-trained teachers. This amortizes distillation for post-trained LLMs: one distillation run can now produce models across post-trained variants.
Figure 2: ADAPT improves generation performance over boomerang distillation. For generation tasks that are both in-domain (math reasoning) and out-of-domain (instruction-following and general knowledge) relative to our SFT dataset, SFT training is necessary to recover generation performance, especially for smaller models. For OOD tasks, two-phase distillation stabilizes performance compared to SFT-only training. We observe similar results for Qwen3-14B, Olmo-3-7B-Instruct, and Llama-3.1-8B-Instruct in Appendix M .
Figure 3: ADAPT (weight-delta) enables reusable base model distillation across post-trained variants. ADAPT (weight-delta) matches ADAPT (distilled) on ID tasks and remains competitive on OOD tasks, showing that one base model distillation run can transfer to multiple post-trained variants without additional training. ADAPT (weight-delta) outperforms baseline cross-patch (2-phase) across all variants. Degradation in thinking-style settings is mainly due to malformed <think> formatting, rather than broad loss of generation quality.
Figure 4: ADAPT enables adaptive model-size selection based on input difficulty for more compute-efficient inference. By producing a continuum of interpolated models with smoothly varying performance, ADAPT allows for more fine-grained matching between model capacity and task difficulty. This approach achieves a Pareto-optimal trade-off between average and worst-case accuracy and computational efficiency.
Figure 5: Smooth weight interpolation is preserved after distillation. Across post-trained variants, both teacher and distilled student models exhibit smooth interpolation paths between base and post-trained endpoints, suggesting that the relevant weight-space structure is preserved after distillation.
Figure 6: Weight-delta initialization produces a student in a connected region of high generation performance. We report AUC for generation accuracy between student and teacher models on the plane intersecting base, instruct, and weight-delta students. The weight-delta student has improved AUC compared to the base student and lies in a high-performance region connected to both the base and instruct students.
Appendix figures & tables42 assets
Supplementary material from the paper’s appendix.
Appendix
| Hyperparameters | Values |
|---|---|
| Learning rate | 3e-4 |
| LR scheduler | cosine |
| Warmup ratio | 0.01 |
| Optimizer | AdamW |
| Adam betas | (0.9, 0.95) |
| Adam epsilon | 1e-8 |
Table 1: Hyperparameters for distillation training.
| Models | Non-thinking | Thinking |
|---|---|---|
| Max output len. | 16384 | 32768 |
| Temperature | 0.7 | 0.6 |
| Top p | 0.8 | 0.95 |
| Top k | 20 | 20 |
| Presence penalty | 0.5 | 0.5 |
| Min. tokens | 32 | 32 |
Table 2: Sampling parameters for non-thinking and thinking models.
| Pre-training Phase | SFT Phase | ||||||
| Model Type | Setup | Steps | Effective BS | Seq. len. | Steps | Effective BS | Seq. len. |
| 4B (Non-think) | Pile only | 480 | 2048 | 1024 | – | – | – |
| Nemotron only | – | – | – | 585 | 4096 | 1024 | |
| Pile+Nemotron | 240 | 2048 | 1024 | 293 | 4096 | 1024 | |
| 4B (Think) | Pile only | 500 | 512 | 4096 | – | – | – |
| Nemotron only | – | – | – | 400 | 1024 | 4096 | |
Table 3: Training hyperparameters across model families. We report the number of steps, effective batch size, and sequence length for the pre-training and SFT phases across all distillation setups.
| Benchmarks | System Prompts |
|---|---|
| AIME, GSM8K, MATH500, MATH | Please reason step by step, and put your final answer within \boxed{} . |
| IFEval | None |
| MMLU-Redux | Please reason step by step, and select the answer from the given choices A, B, C, or D. Respond only with the letter of the correct answer. Put the index within \boxed{} . |
Table 4: System prompts for generation tasks.
Figure 7: Final layer activation cosine similarity between student and teacher models. The two-phase distillation setup yields the best student-teacher alignment, further motivating our approach.
Figure 8: Classification accuracy is relatively insensitive to the distillation strategy. Across Qwen, Olmo, and Llama, most setups exhibit similar interpolation behavior on classification benchmarks. The main exception is BD (SFT), which underperforms across much of the interpolation curve, suggesting that task-specific distillation alone can degrade broader alignment with the teacher.
Figure 9: Distilling with equal pre-training and SFT phase offers an appropriate tradeoff between ID and OOD performance. Increasing the SFT proportion improves ID accuracy but degrades OOD accuracy, while increasing the pre-training proportion has the opposite effect; the 50:50 split balances the two. The single-phase endpoints are consistently dominated by the two-phase settings.
Figure 10: Distilling with cosine loss only yields interpolation performance comparable to the full objective. We compare training with cosine alignment loss only to training with all of the loss terms in Equation 1 and find that cosine only training has similar performance, indicating that the efficiency of interpolation training can potentially be improved by using only alignment loss.
Figure 11: Pre-training weight-delta initialization for post-trained Qwen models.
Figure 12: SFT weight-delta initialization for post-trained Qwen models.
Figure 13: Pre-training weight-delta initialization performance for post-trained Olmo models.
Figure 14: Pre-training weight-delta initialization performance for post-trained Llama models.
Figure 15: Continued training from weight-delta initialization does not improve interpolation performance. We initialize the instruct student using a pre-training, SFT, or two-phase weight delta, computed from base distillation on pre-training data, SFT data, or both respectively. We then continue training for 0.5B or 1B tokens. Continued training does not meaningfully improve over weight-delta initialization.
Figure 16: Results for ADAPT (weight-delta) using delta from Qwen3-4B-Instruct-2507.
Figure 17: Results for ADAPT (weight-delta) using delta from Qwen3-4B-Thinking-2507.
Figure 18: Results for ablating KL loss term in SFT phase.
Figure 19: Results for ablating cosine loss term in SFT phase.
Figure 20: Results for ablating both KL and cosine loss terms in SFT phase (cross-entropy loss only).
| Variant | Cosine | Norm ratio |
|---|---|---|
| Qwen3-4B-Instruct-2507 | 0.73 | 1.00 |
| Qwen3-4B-Thinking-2507 | 0.57 | 0.92 |
| Qwen3-4B | 0.75 | 1.00 |
| Qwen3-14B | 0.75 | 1.00 |
| Olmo-3-7B-Instruct | 0.73 | 1.02 |
| Olmo-3-7B-Think | 0.64 | 0.94 |
Table 5: Cosine similarity and norm ratio between base and post-trained deltas for each variant.
Figure 21: ADAPT significantly outperforms layer pruning baselines on generation tasks. We compare ADAPT (distilled) and ADAPT (weight-delta) to two popular layer pruning methods for generation tasks, ShortGPT ( 32 ) and LLM-Streamline ( 7 ) , and show that our method performs substantially better across all sizes.
| Hyperparameters | Values |
|---|---|
| LR (MSE) | 2e-4 |
| LR (LM-loss) | 5e-5 |
| LR scheduler | cosine |
| Warmup ratio | 0.01 |
| Optimizer | AdamW |
| Adam betas | (0.9, 0.999) |
Table 6: Hyperparameters for LLM-Streamline training.
Figure 22: Prompt for difficulty classification. Adapted from 46 .
Figure 23: Interpolation performance of Qwen3-4B-Instruct-2507 for different difficulty levels. We find that the teacher-defined difficulty score produces similar interpolation curves to the ground truth difficulty scores, but student-defined difficulty is unreliable. This shows that the teacher-defined labels can act as a reasonable approximation for groundtruth difficulties.
| Levels | Teacher | Student |
|---|---|---|
| 1 | 184 | 1035 |
| 2 | 282 | 9 |
| 3 | 149 | 2 |
| 4 | 195 | 0 |
| 5 | 240 | 4 |
Table 7: Breakdown of input-difficulty classification results from teacher and student models. The student is unreliable for difficulty classification, as it chooses the lowest difficulty level for the vast majority of examples.
Figure 24: Adaptive model size selection results for Olmo-3-7B-Instruct.
Figure 25: Adaptive model size selection results for Llama-3.1-8B-Instruct.
Figure 26: Comparison of interpolation performance for different patching orders. We test patching from the front, patching from the back, and KLPatch ( 24 ) . We find that KLPatch can further improve performance, indicating that more sophisticated patching strategies are orthogonal to our approach and can be used in conjunction with our student models.
Figure 27: Smooth weight interpolation results for generation accuracy. Base and post-trained students display smooth interpolation behavior in downstream generation accuracy.
Figure 28: Smooth weight interpolation results for loss on non-thinking Nemotron data. Base and post-trained students display smooth interpolation behavior in loss evaluated on non-thinking Nemotron data.
Figure 29: Smooth weight interpolation results for loss on thinking Nemotron data. Base and post-trained students display smooth interpolation behavior in loss evaluated on thinking Nemotron data.
Figure 30: Training dynamics study of two-phase distillation. We find that two-phase distillation converges to two distinct interpolation curves during the pre-training and SFT phases, which allows the model to retain general-purpose and in-domain performance.
Figure 31: ADAPT (distilled) results for Qwen3-14B (Thinking). Two-phase distillation produces the highest and most stable performance over both ID and OOD tasks.
Figure 32: ADAPT (distilled) results for Qwen3-14B (Non-Thinking). Two-phase distillation produces the highest and most stable performance over both ID and OOD tasks.
Figure 33: ADAPT (weight-delta) results for Qwen3-14B with both thinking and non-thinking modes. ADAPT (weight-delta) remains competitive with ADAPT (distilled) across both ID and OOD settings, while consistently outperforming cross-patch (2-phase).
Figure 34: ADAPT (distilled) results for Olmo models. Two-phase distillation produces the highest and most stable performance over both ID and OOD tasks.
Figure 35: ADAPT (weight-delta) results for Olmo models. ADAPT (weight-delta) remains competitive with ADAPT (distilled) across both ID and OOD settings, while consistently outperforming cross-patch (2-phase) (with the exception of OOD performance for Olmo-3-7B-Think).
Figure 36: ADAPT (distilled) results for Llama models. Two-phase distillation produces the highest and most stable performance over both ID and OOD tasks.
Figure 37: ADAPT (weight-delta) results for Llama models. ADAPT (weight-delta) remains comparable with ADAPT (distilled) across both ID and OOD settings.
Figure 38: ADAPT (distilled) results for Qwen3-4B on coding tasks. Similar to the math setup, two-phase distillation produces the highest and most stable performance over both ID and OOD tasks. The evaluation is done on thinking mode only.
Figure 39: ADAPT (weight-delta) results for Qwen3-4B on coding tasks. ADAPT (weight-delta) maintains a trade-off between the higher-compute upper bound ADAPT (distilled) and compute-matched baseline cross-patch (2-phase). It is cheaper than ADAPT (distilled) while performing better than cross-patch (2-phase).
Figure 40: Interpolation results for Qwen3-4B-Instruct-2507 on individual benchmarks.
Figure 41: Weight-delta initialization experiment results for Qwen 4B model family on individual benchmarks.
More in this topic
Language Model Distillation
Language model distillation transfers a larger model’s capabilities to smaller or specialized models using outputs, probabilities, or internal representations.
See all 514 papers