PaperlayerView exact revision on arXiv

Source-linked AI summary

Boomerang Distillation Enables Zero-Shot Model Size Interpolation

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

arXiv:2510.05064v2cs.LG

TL;DR

Building many LLM sizes is expensive because each student normally needs a separate training run. Boomerang distillation distills a teacher into a small, aligned student and then reinserts teacher-layer blocks without retraining, producing intermediate models whose performance interpolates between student and teacher. These models often match or surpass same-size distilled models, although the method requires compatible teacher–student construction and alignment.

  • Problem

    Each model size in conventional LLM families requires a full training run, making fine-grained size scaling prohibitively expensive.

  • Method

    Distill a layer-pruned, teacher-initialized student with alignment loss, then patch teacher-layer blocks back into it without additional training.

  • Results

    Interpolated models smoothly span student-to-teacher performance and match or sometimes surpass standard distilled models of the same size.

  • Takeaways & Limitations

    Boomerang distillation provides a simple zero-shot way to generate fine-grained model families from one aligned teacher–student pair.

  • Takeaways & Limitations

    The approach requires compatible layer dimensions and teacher-initialized students with alignment-compatible distillation; many existing distilled LLMs cannot be patched directly.

Abstract

from arXiv · show

Large language models (LLMs) are typically deployed under diverse memory and compute constraints. Existing approaches build model families by training each size independently, which is prohibitively expensive and provides only coarse-grained size options. In this work, we identify a novel phenomenon that we call boomerang distillation: starting from a large base model (the teacher), one first distills down to a small student and then progressively reconstructs intermediate-sized models by re-incorporating blocks of teacher layers into the student without any additional training. This process produces zero-shot interpolated models of many intermediate sizes whose performance scales smoothly between the student and teacher, often matching or surpassing pretrained or distilled models of the same size. We further analyze when this type of interpolation succeeds, showing that alignment between teacher and student through pruning and distillation is essential. Boomerang distillation thus provides a simple and efficient way to generate fine-grained model families, dramatically reducing training cost while enabling flexible adaptation across deployment environments. The code and models are available at https://github.com/dcml-lab/boomerang-distillation.

1 INTRODUCTION

Boomerang distillation addresses the cost of building fine-grained LLM size families by distilling a large teacher into a small student, then reinserting teacher-layer blocks without further training. The resulting interpolated models span student-to-teacher sizes and generally preserve smoothly improving performance while relying on teacher–student alignment.

  • Motivation: Fine-grained model families remain expensive because conventional approaches require a full training run for every student size.This cost limits the practicality of offering many deployment-specific model sizes.
  • Overview: Boomerang distillation reconstructs intermediate models by reinserting teacher-layer blocks into a distilled student without additional training.The procedure creates a spectrum of model sizes between the student and teacher.
  • Main finding: These interpolated models smoothly bridge student and teacher performance across downstream tasks.They use information from both models rather than only the teacher, unlike pruning-based approaches.
  • Conditions: The phenomenon depends on teacher-weight initialization, distillation, and an alignment loss such as cosine distance.Alignment loss is reported to support stable boomerang-distilled models and consistent advantages over pruning-based approaches.
  • Empirical scope: Interpolated models match or sometimes surpass standard distilled models of the same size.The reported phenomenon also appears in open-source model pairs including DistilBERT/BERT and DistilGPT2/GPT2.

2 BOOMERANG DISTILLATION: KNOWLEDGE DISTILLATION WITH STUDENT PATCHING

The method has three stages: initialize a smaller student from teacher layers, distill it while aligning hidden states to teacher blocks, and patch teacher layers back into the trained student. Repeated patching deterministically produces intermediate models with different layer counts.

  • Procedure: The pipeline consists of student initialization, knowledge distillation, and student patching.These stages convert a pretrained teacher and smaller student into interpolated models.
  • Assumptions: The construction assumes teacher and student layers share the same hidden-state dimension and supports deterministic models with M + K layers for nonnegative K where M + K < N.The student is trained without requiring access to the teacher’s pretraining data.
  • Student initialization: The student is initialized by partitioning the teacher’s layers into contiguous blocks and using those blocks to form the smaller model.The student has M layers while the teacher has N layers, with M < N.
  • Knowledge distillation: Distillation combines cross-entropy and KL-divergence losses with cosine-distance alignment across corresponding student and teacher hidden states.The alignment encourages each student layer to approximate the output of its associated teacher block, enabling later block replacement.
  • Student patching: After distillation, replacing student layers with their corresponding teacher blocks and repeating the substitution yields models with various intermediate sizes.The embedding layer and LM head are selected from the models contributing the first and last transformer layers.

3 EXPERIMENTS

Experiments show that boomerang distillation creates zero-shot intermediate models whose performance interpolates smoothly between distilled students and teachers across model families and baselines. Success depends primarily on teacher-weight initialization, while alignment losses improve stability and quality.

  • 3.1 THE BOOMERANG DISTILLATION PHENOMENON: Teacher-weight initialization plus distillation is necessary for strong interpolation, whereas naive pruning and randomly initialized distillation fail to produce comparable gains.Naive pruning sharply degrades smaller models, while patching randomly initialized students provides almost no improvement.
  • 3.1 THE BOOMERANG DISTILLATION PHENOMENON: Across Qwen3, Pythia, and Llama families, boomerang distillation yields smooth size–performance interpolation across different model sizes.The phenomenon therefore extends beyond the primary Qwen3-4B-Base setup.
  • 3.2 HOW GOOD IS BOOMERANG DISTILLATION?: Boomerang distillation produces intermediate models that match or exceed standard distilled models and pretrained models, while requiring training only one small student.At larger sizes, interpolated models outperform standard distilled models, likely because distillation on lower-quality corpora causes catastrophic forgetting.
  • 3.3 EFFECT OF KNOWLEDGE DISTILLATION: Per-layer cosine alignment produces the lowest perplexity and smoother interpolation, especially near the first and last teacher layers.Boomerang distillation still occurs without the per-layer loss, indicating that teacher-layer initialization supplies substantial alignment information.
  • 3.4 ZERO-SHOT MODEL SIZE INTERPOLATION WITH EXISTING OFF-THE-SHELF MODELS: Off-the-shelf DistilBERT/BERT and DistilGPT2/GPT2 pairs also support zero-shot interpolation without additional training, outperforming naive teacher-layer pruning.BERT interpolation is cleaner than GPT2 interpolation, but both improve over pruning.
  • 3.4 ZERO-SHOT MODEL SIZE INTERPOLATION WITH EXISTING OFF-THE-SHELF MODELS: Existing distilled LLMs are often incompatible with boomerang distillation because pruning can mismatch hidden-state dimensions and many frameworks omit cosine alignment loss.The authors identify adapting such models with layer pruning, cosine loss, and large token budgets as future work.
  • 3.3 EFFECT OF KNOWLEDGE DISTILLATION: Against LaCo and ShortGPT, boomerang distillation maintains higher accuracy across intermediate sizes, particularly when pruning methods’ generation performance collapses.The advantage is reported for both classification and generation tasks.
  • 3.6 ABLATIONS: The method performs well with smaller students when the distilled student retains non-trivial target-task performance, and more student training tokens improve interpolated models.These ablations identify student quality and training budget as practical conditions for successful interpolation.

4 RELATED WORK

Prior work covers model interpolation, knowledge distillation, pruning, and dynamic compute allocation, while boomerang distillation interpolates model sizes by patching a distilled student with teacher layers without a specialized router.

  • Model Interpolation: Model interpolation traditionally combines weights of same-size models, whereas boomerang distillation creates different-size models by patching a distilled student with teacher layers.Unlike elastic-transformer approaches, it uses a standard knowledge-distillation pipeline and does not require training a specialized router.
  • Knowledge Distillation: Knowledge distillation trains a smaller student to mimic a larger teacher and can accommodate architectures that do not match.It has been applied across vision models, language models, and proprietary API-based systems.
  • Pruning: Pruning reduces computational requirements by removing redundant parameters while aiming to preserve full-model performance.Representative strategies include layer dropping, neuron pruning, and SVD-based pruning.
  • Dynamic Compute Allocation: Dynamic compute allocation adjusts inference-time computation to task complexity using methods such as early exiting, test-time scaling, and compute-adaptive embeddings.This work focuses on creating model-size choices through boomerang distillation rather than dynamically allocating inference computation.

5 CONCLUSION

The paper identifies boomerang distillation as a way to construct fine-grained model families between a student and teacher without additional training. It reports that the method works across existing open-source models and can match or outperform comparable alternatives.

  • 5 CONCLUSION: Boomerang distillation creates models that smoothly interpolate in size and performance between a given student and teacher without additional training.The student must be initialized from the teacher through layer pruning and trained with knowledge distillation.
  • 5 CONCLUSION: Interpolated models consistently match or outperform same-size models directly trained with knowledge distillation and outperform existing pruning approaches downstream.
  • 5 CONCLUSION: The phenomenon also occurs in existing open-source models such as DistilBERT and DistilGPT2.
  • 5 CONCLUSION: A single student-teacher pair can produce fine-grained model families while significantly reducing training time and cost.

ETHICS STATEMENT

Interpolated models may inherit or amplify biases from the pretrained teacher, so deployment should include target-task evaluation and additional mitigation training.

  • ETHICS STATEMENT: Interpolated models may inherit or amplify biases from the pretrained teacher model.The paper recommends comprehensive evaluation on target tasks before deployment.
  • ETHICS STATEMENT: Instruction-following training and additional safety training are suggested to mitigate residual biases.

A LIMITATIONS AND FUTURE WORK

The method depends on an expensive distilled student and layer-compatible pruning, while scaling, patching order, and transfer beyond language remain open questions. More similarity-guided initialization and patching may improve stability, but broader settings require further work.

  • A LIMITATIONS AND FUTURE WORK: Training the distilled student remains computationally expensive even though intermediate models require no additional training.The distilled student is crucial for boomerang distillation and requires significant compute.
  • A LIMITATIONS AND FUTURE WORK: The experiments are limited in model size and distillation-token budget, leaving scaling to larger models as future work.
  • A LIMITATIONS AND FUTURE WORK: Naive patching orders can destabilize performance, while similarity-guided initialization and patching may mitigate this instability.The paper reports smoother interpolation when patching with layers having low activation cosine similarity to their teacher counterparts.
  • A LIMITATIONS AND FUTURE WORK: Layer-pruned students are required, and width or attention-head pruning can create hidden-dimension mismatches that prevent patching out of the box.
  • A LIMITATIONS AND FUTURE WORK: Whether boomerang distillation extends beyond language models to transformer-based vision and audio models remains unresolved.

B BOOMERANG DISTILLATION IMPLEMENTATION

The implementation distills pruned teacher models into students, then patches student layers with teacher blocks to create intermediate models without retraining. Smooth interpolation depends on non-trivial student performance and alignment-preserving initialization, with behavior extending across several model families.

  • Implementation: Student models are initialized by pruning teacher layers, trained with distillation and alignment losses, then progressively patched with corresponding teacher-layer blocks.The default setup keeps every other teacher layer, trains on 2.1B Pile tokens, and replaces student layers progressively during inference.
  • Student size: Boomerang distillation interpolates smoothly when the distilled student retains non-trivial performance, whereas the smallest naively initialized students do not.Every-2nd and every-3rd-layer students behave similarly, but every-4th and every-5th-layer students show unsmooth interpolation.
  • Alignment: Cosine similarity between student layers and their corresponding teacher blocks correlates with interpolation quality and motivates patching high-similarity layers first.Low similarity explains accuracy drops in smaller students and provides a heuristic for avoiding substantial performance loss.
  • Further compression: Cosine-similarity-informed initialization compresses Qwen3-4B-Base students up to 8.7x while preserving interpolation behavior.Such aggressive compression yields fewer intermediate sizes because each student layer corresponds to an entire teacher block.
  • Generalization: Boomerang distillation appears across Qwen3-8B, Qwen3-14B, Pythia, and Llama-3.2-3B, but task behavior depends on the teacher and training setup.Pythia interpolation improves classification while generation remains near zero, whereas Llama-3.2-3B produces smooth interpolation across classification and generation tasks.

G PYTHIA-2.8B FULL RESULTS

For Pythia-2.8B, boomerang-distilled models are comparable to standard distilled models and layer-pruning methods on classification, while generation remains trivial because the teacher performs poorly there.

  • Standard knowledge distillation: Boomerang-distilled Pythia-2.8B models have comparable performance to intermediate models trained with standard knowledge distillation.Standardly distilled models perform better across all sizes in this comparison.
  • Effect of knowledge distillation: Cross-entropy plus alignment loss produces lower perplexity and meaningful classification differences, especially at smaller model sizes.Interpolated models trained without these losses perform worse on the reported measures.
  • Task scope: Generation performance is nearly 0% for interpolated Pythia models and pretrained models, so boomerang distillation does not improve beyond the pretrained baseline on those tasks.The reported improvement for Pythia is confined to classification accuracy.
  • Layer-pruning methods: Pythia boomerang distillation and layer pruning show similar classification performance, with boomerang distillation stronger at smaller sizes and LaCo or ShortGPT stronger at larger sizes.The comparison includes LaCo and ShortGPT as layer-pruning baselines.

H LLAMA-3.2-3B FULL RESULTS

For Llama-3.2-3B, boomerang distillation yields smooth performance interpolation and compares favorably with standard distillation and layer pruning, especially at larger or smaller sizes respectively.

  • Standard knowledge distillation: Boomerang-distilled Llama-3.2-3B models have comparable performance to standard distilled models and outperform them at larger sizes.Distilling on The Pile can reduce performance relative to the pretrained model, while patching retains some pretraining benefits.
  • Effect of knowledge distillation: Alignment loss lowers perplexity across most intermediate sizes and slightly improves classification accuracy, especially around 2.5-3B inference parameters.Removing alignment loss also degrades generation performance at high parameter counts.
  • Layer-pruning methods: Boomerang distillation outperforms layer-pruning approaches at all sizes, although LaCo becomes competitive around 3.2B parameters.The initial classification gap is large, but LaCo recovers performance as model size increases.

I LLAMA-3.2-3B COSINE SIMILARITY ANALYSIS

Cosine similarity analysis explains why naive Llama-3.2-3B initialization and patching can disrupt interpolation. Preserving first-layer information and patching low-similarity blocks first restores smoother, stronger interpolation.

  • Similarity analysis: Per-layer activation cosine similarity explains when boomerang distillation becomes noisy or interpolates poorly.The analysis uses 128 calibration samples and recommends patching low-similarity student layers first while avoiding pruning consecutive low-similarity layers.
  • Standard initialization: Naive Llama initialization aligns the first student layer with the first teacher layer but leaves most corresponding layers mismatched until the final student layers.Patching from the last layers therefore fails to recover smooth interpolation, whereas patching from the first layers mitigates the issue.
  • Best configuration: The best Llama-3.2-3B configuration initializes with the first two teacher layers plus alternating later layers, then patches from the first layer to the last.This configuration creates the best interpolated models among the tested initialization and patching orders.
  • Preserving first-layer information: Keeping the first two teacher layers during initialization substantially raises student-teacher cosine similarity and improves student and interpolation performance.The gain is strongest when combined with patching from the first layer to the last.
  • Takeaways: For Llama-3.2-3B, naive initialization and patching are insufficient; alignment can instead be improved by prioritizing low-similarity blocks or preserving similarity within each teacher block.The paper links this alignment condition to improved interpolation performance.

J STUDENT MODEL SIZE ABLATION COSINE SIMILARITY ANALYSIS

Cosine alignment between student layers and corresponding teacher blocks predicts whether patching improves interpolation, while progressively patching layers makes hybrids better approximate the teacher. Boomerang distillation also amortizes training cost by producing intermediate models after one initial distillation run.

  • Cosine similarity and interpolation: Student–teacher cosine similarity predicts interpolation behavior: high-similarity layers improve patched models, whereas low-similarity layers reduce performance.Patching layers 7, 6, 5, and 4 improved performance, patching layers 3 and 2 reduced it, and later patching increased it again.
  • Mechanism: Aligned initialization and distillation make each student layer approximate its corresponding teacher block, enabling useful layer replacement.The proposed intuition is that patching a student layer with its corresponding teacher block yields a better approximation of the teacher model.
  • Teacher approximation: Patching more student layers increases last-layer cosine similarity between interpolated and teacher models across model families.The proof-of-concept evaluates this similarity on 128 held-out texts from The Pile.
  • Computational cost: Boomerang distillation amortizes training because intermediate models require no additional training after the initial distillation.The computational comparison measures training cost over 2.1B tokens for individually distilled intermediates versus boomerang distillation.

M.1 PERPLEXITY

Across perplexity, classification, and generation evaluations, boomerang distillation produces smoothly interpolated models with stronger stability than layer-pruning approaches. Its interpolated models remain comparable to pretrained and naively distilled models, while pruning causes larger perplexity degradation as more layers are removed.

  • Perplexity: Interpolated models have perplexity comparable to standard distilled models, with smaller distilled models slightly better and larger distilled models slightly worse.Pretrained Pythia-2.8B and Llama-3.2-3B models have lower perplexity but similar classification and generation performance to interpolated models.
  • Perplexity: Boomerang distillation interpolates perplexity smoothly between student and distilled models, while layer-pruning approaches degrade substantially after more than six layers are dropped.Randomly initialized distilled interpolations have perplexity above 104.
  • Baselines: The pruning comparisons include Layer Collapse and ShortGPT alongside naive layer-pruning baselines.LaCo merges layer chunks using activation cosine similarity, while ShortGPT removes layers sequentially using Block Importance scores.
  • Classification accuracy: Boomerang distillation produces smoothly interpolated classification accuracy and comparable per-task performance to pretrained and naively distilled models.Per-layer loss yields stable and smoother classification accuracy, while boomerang distillation outperforms depth-pruning methods.
  • Generation accuracy: Boomerang distillation produces smoothly interpolated generation accuracy and comparable per-task performance to pretrained and naively distilled models.Per-layer loss yields stable and smoother generation accuracy, while boomerang distillation outperforms depth-pruning methods.
Loading 2510.05064v2…