PaperlayerView exact revision on arXiv

Source-linked AI summary

Distillation Scaling Laws

Dan Busbridge, Amitis Shidani, Floris Weers, Jason Ramapuram, Etai Littwin, Russ Webb

arXiv:2502.08606v2cs.LGcs.AIcs.CLstat.ML

TL;DR

Distillation requires principled compute allocation because its efficiency depends on student budget and teacher reuse. The paper develops a distillation scaling law and finds that distillation can be more efficient than supervised learning under bounded student compute when an existing or reusable teacher is available, while identifying bias inheritance and misuse risks.

  • Problem

    The paper addresses how to allocate resources for large-scale distillation and whether distillation can outperform supervised pretraining under comparable resources.

  • Method

    The authors conduct a controlled study and develop a scaling law that estimates student performance from compute allocation among the teacher and student.

  • Results

    Distillation is more efficient than supervised learning below a student-size-dependent compute or token threshold when a teacher already exists or has uses beyond one distillation.

  • Takeaways & Limitations

    The scaling law identifies compute-optimal distillation recipes and expands the set of models with known compute-optimal configurations, reducing compute, hardware, and climate costs for a target capability.

  • Takeaways & Limitations

    Distillation can inherit bias from a teacher’s pretraining data even when the student is trained on unbiased data.

Abstract

from arXiv · show

We propose a distillation scaling law that estimates distilled model performance based on a compute budget and its allocation between the student and teacher. Our findings mitigate the risks associated with large-scale distillation by enabling compute-optimal allocation for both the teacher and student to maximize student performance. We provide compute-optimal distillation recipes for two key scenarios: when a teacher already exists, and when a teacher needs training. In settings involving many students or an existing teacher, distillation outperforms supervised learning up to a compute level that scales predictably with student size. Conversely, if only one student is to be distilled and a teacher also requires training, supervised learning is generally preferable. Additionally, our large-scale study of distillation increases our understanding of the process and helps inform experimental design.

1. Introduction

Distillation is studied as a way to obtain capable, efficient language models while allocating compute between teachers and students. The paper introduces a scaling law that identifies when distillation is more efficient than supervised learning and how teacher capability affects students.

  • Distillation targets small, capable models that reduce the inference and lifetime costs associated with compute-optimal models.Small models can be cheaper to serve and may reduce barriers to academic study and carbon costs.
  • The study addresses how to allocate substantial distillation resources and whether distillation can outperform supervised pretraining under equal resources.
  • The authors conduct a controlled study spanning transformer teachers and students from 143M to 12.6B parameters and training data from billions to 512B tokens.
  • The distillation scaling law predicts student performance from the teacher, student size, and amount of distillation data.
  • Distillation is more efficient than supervised learning only below a student-size-dependent compute or token threshold and when a teacher already exists or has uses beyond one distillation.
  • The resulting laws are intended to guide more capable models with lower inference cost and lower lifetime compute costs.

2. Background

The paper frames scaling laws as tools for predicting performance and allocating compute between model size and data. It adapts this approach to distillation by predicting student cross-entropy and identifying compute-efficient student configurations.

  • Scaling laws predict how compute should be distributed between model parameters and training data to achieve desired model properties.
  • Power-law scaling provides empirical and theoretical motivation for modeling cross-entropy as a function of parameters and data.
  • The experiments vary training tokens at fixed model size and vary model size and tokens under IsoFLOP constraints.
  • Data from both experimental strategies are combined to fit the scaling relationships.
  • The distillation objective is to predict a student’s cross-entropy and determine the lowest cross-entropy attainable for a given student size and compute budget.

3. Preliminaries

The preliminaries define autoregressive language modeling and distillation losses. In distillation, the teacher’s predicted next-token distribution replaces the one-hot target used in supervised training, with temperature-scaled divergence contributing to student optimization.

  • Autoregressive language models predict each next token from its preceding context and are trained with next-token prediction.
  • Distillation uses the teacher’s predicted next-token distribution as the target for the student distribution.
  • The distillation objective is equivalent to optimizing Kullback-Leibler divergence between teacher and student predictions.
  • The distillation temperature τ is positive and controls the temperature-scaled prediction comparison.

4. Distillation Scaling Laws

The paper fits and tests a distillation scaling law using complementary teacher–student experiments. The law summarizes teacher training through teacher cross-entropy, captures a capacity gap, and fits observations with low relative prediction error.

  • Experimental design: The study combines experimental protocols that vary student training, teacher training, and teacher–student sizes to identify scaling-law coefficients.
  • Experimental design: Fixed-token-ratio teachers are distilled into students across IsoFLOP profiles, while teacher training cost is varied separately for fixed students.
  • Capacity gap: A capacity gap occurs because increasing teacher capability can eventually worsen student performance.
  • Scaling-law structure: The capacity gap is modeled as a transition between power-law regimes determined by the relative learning capacities of teacher and student.
  • Validation: The supervised and distillation scaling laws fit observations with at most approximately 1% relative prediction error, including extrapolation from weaker to stronger models.

5. Distillation Scaling Law Applications

The distillation scaling law models student cross-entropy under compute and token constraints, enabling comparisons with supervised learning and compute-optimal teacher selection. Distillation is advantageous only in specific resource regimes and depends primarily on teacher quality rather than teacher size.

  • Compute accounting: The compute cost accounts for student training, teacher-logit inference, and optionally teacher pretraining.Teacher-logit and pretraining costs are controlled by indicators, while model FLOPs use a size- and context-sensitive approximation.
  • Fixed tokens or compute: Supervised learning outperforms distillation with enough student compute or tokens, whereas distillation is favorable under modest token budgets.In large-data regimes, supervised learning can reach the model-size-limited solution, while distillation may remain limited by the teacher.
  • Teacher inference: The optimal teacher cross-entropy decreases as a power law with student size until it reaches a matching-loss inflection point.This analysis treats teacher size indirectly: the relevant quantity for student performance is teacher cross-entropy rather than teacher parameter count.
  • Compute-optimal distillation: Supervised learning matches optimal distillation at sufficient total compute, and the compute transition favoring supervision increases with student size.Smaller models are more likely to benefit from supervised pretraining, while larger models are more likely to benefit from distillation.
  • Compute-optimal distillation: When teacher training is included, the best distilled student has higher cross-entropy than the supervised counterpart.Distillation may nevertheless be computationally preferable when the teacher serves multiple students or has uses beyond one distillation.
  • Compute-optimal distillation: In the teacher-pretraining-plus-inference scenario, student and teacher token counts scale as power laws, with student tokens growing faster.Teacher size initially increases, then plateaus because large-teacher inference is expensive and teacher overtraining becomes more efficient.

6. Conclusion

The paper proposes a distillation scaling law for allocating compute between teachers and students. It finds that distillation is more efficient than supervised learning only below a student-size-dependent resource threshold and when the teacher already exists or has broader uses.

  • Conclusion: The distillation scaling law estimates distilled model performance from a compute budget and its allocation between the student and teacher.The law is used to study practical distillation scenarios and identify optimal teachers.
  • Conclusion: Distillation is more efficient than supervised learning only when distillation resources remain below a student-size-dependent threshold and the teacher already exists or has additional applications.These conditions define the paper’s supported scope for compute-efficient distillation.
  • Conclusion: Compute-optimal recipes can produce smaller, more powerful models with lower inference costs and improve the feasibility of test-time scaling.The conclusion connects optimal allocation with reduced inference burden and carbon footprint.

Impact Statement

The framework identifies when distillation can reduce the costs of producing capable language models and narrows the search for compute-optimal configurations. It also highlights risks from inherited teacher bias and the easier misuse of efficient small models.

  • Benefits: Compute-optimal distillation recipes create more opportunities to produce powerful models with lower inference costs and carbon footprints.Lower inference costs address a major component of language-model lifecycle emissions.
  • Benefits: Identifying optimal recipes reduces compute, hardware, and climate costs for models targeting a given capability.The paper frames this as expanding the set of models with known compute-optimal configurations.
  • Benefits: The distillation scaling law reduces experimentation by treating teacher cross-entropy as the primary driver of student cross-entropy.Teacher size and teacher-token count can therefore be removed as independent search dimensions.
  • Benefits: Small powerful models broaden participation in studying highly capable models and their safety properties.Their lower resource requirements make capability and safety research more accessible.
  • Potential negative consequences: Distillation can transmit teacher-pretraining bias to students even when the student is trained on unbiased data.The stated risk concerns bias inherited through the teacher’s targets.
  • Potential negative consequences: More efficient small models can reduce the resources needed by malicious actors to generate targeted misinformation at scale.The paper identifies this as a potential misuse consequence of improved inference efficiency.

A. Limitations

The study is limited by its language-modeling setting, English C4 data, transformer architectures, and same-distribution training design. Its scaling law coefficients are not guaranteed to generalize, and several distillation variants and downstream capabilities remain insufficiently tested.

  • The experiments study language modeling only, so the described distillation behavior may not generalize across domains.The authors note that cross-domain applicability is supported by some scaling-law evidence but remains uncertain for distillation.
  • Using repeated English C4 data limits certainty about whether repeated-data effects transfer from supervised training to distillation.The larger runs repeat data, while the cited evidence for negligible loss impact concerns supervised learning.
  • C4 limits downstream analysis, especially for reasoning-related capabilities that the dataset is not well suited to probe.Standard English downstream performance closely follows cross-entropy, but this does not establish coverage of reasoning performance.
  • The study focuses on full-probability-distribution distillation and does not establish that its findings apply to sequence-level, synthetic-data, intermediate-layer, or modified-objective methods.The authors describe these approaches as common variations but state that broader applicability is uncertain.
  • The analysis uses transformer teacher-student pairs trained on the same data distribution, omitting domain-shifted teacher-to-student workflows.This design isolates algorithmic effects but leaves teacher-student domain differences for future study.
  • The scaling law’s coefficients are specific to the studied architectures and dataset, and its form is not derived from first principles.The authors therefore cannot fully guarantee the law’s correctness or generalization to other architectures and datasets.

C.1.3. U-SHAPE IN THE STUDENT ERROR

Student error follows a U-shaped relationship with teacher capacity: increasing teacher capacity helps below the student’s learning capacity, but eventually worsens performance because of the capacity gap. Across compute settings, distillation is most useful with an appropriately trained teacher and limited resources, while supervised learning catches up or becomes preferable as student compute grows.

  • U-shaped student error: For teachers below the student capacity, increasing teacher capacity does not increase error and typically improves the student.This monotonic improvement reflects the regime where the teacher’s error decreases as teacher capacity increases.
  • U-shaped student error: When teacher capacity reaches or exceeds the student’s, student error increases monotonically, so a stronger teacher can produce a worse student.The capacity gap reflects differences in learning capacity, including hypothesis space and optimization ability, not only model size.
  • U-shaped student error: Student error first decreases and then increases as teacher capacity rises, establishing a U-shaped relationship caused by the capacity gap.The decreasing regime holds below the student capacity; beyond that point, additional teacher capacity increases error.
  • Distillation versus supervised learning: Distillation can outperform supervised learning with modest resources, but both methods produce the same student at sufficiently large compute.The compute level where supervised learning catches up increases with student size.
  • Distillation versus supervised learning: If teacher training is included in the budget and only one target model is needed, supervised learning generally yields lower student cross-entropy than training a teacher and distilling.Distillation may become more favorable when the teacher is reused for multiple students or served as a model.
  • Compute-optimal teacher allocation: The teacher should be overtrained: optimal teacher tokens continue increasing, while teacher size eventually plateaus or decreases because larger teachers impose greater compute and inference costs.The optimal teacher cross-entropy generally decreases with total compute, with an inflection at low compute for larger students.

D.4.7. OPTIMAL TEACHER TRAINING AND STUDENT DISTILLATION TOKENS

Compute-optimal distillation allocates student and teacher resources differently across compute scenarios. Student tokens increase with compute, while teacher-token and teacher-size optima interpolate according to how teacher-training costs are counted.

  • Student allocation: Student tokens increase polynomially with compute across distillation scenarios, matching supervised scaling in the best-case setting.Pretraining distillation approaches the same allocation asymptotically; other scenarios differ mainly by offsets.
  • Teacher allocation: Optimal teacher tokens interpolate between scenarios: at low compute, teacher-pretraining strategies coincide; at high compute, they approach teacher-inference or best-case allocation.As student tokens grow, teacher-pretraining costs become less important.
  • Teacher allocation: Optimal teacher size shows the same compute-allocation interpolation behavior as optimal teacher tokens.The shared mechanism is the changing relative importance of teacher-pretraining costs.
  • Distillation efficiency: When teacher training is discounted, distillation can require less compute and fewer tokens than supervised learning for cross-entropies slightly above the lowest achievable value.This advantage applies to the base-case and teacher-inference scenarios.
  • Distillation efficiency: As the target approaches the strongest possible student performance, supervised learning becomes more efficient than distillation.The authors suggest a possible distillation-to-supervised transition, but leave its threshold for future work.
  • Distillation efficiency: Including teacher pretraining in the computational cost makes supervised learning more efficient than distillation for all student sizes.Both teacher-pretraining and teacher-pretraining-plus-inference scenarios require more compute than supervised learning.

E.6. Distillation with infinite data

With effectively unlimited data, the scaling analysis examines distillation’s asymptotic lower bounds, weak-to-strong behavior, and calibration. Optimal teachers can match the supervised lower bound, while finite-data distillation can let larger students outperform weaker teachers.

  • Infinite-data limit: Only the optimal teacher for each student size reaches a cross-entropy lower bound consistent with supervised learning; weaker or mismatched teachers remain worse.The consistency is numerical evidence for the scaling law’s behavior beyond the observed models.
  • Weak-to-strong generalization: Weak-to-strong generalization occurs only with finite distillation data; with sufficiently many tokens, student cross-entropy rises toward the teacher’s cross-entropy.The asymptotic behavior is consistent with a larger student eventually matching the teacher across the data manifold.
  • Calibration: Distillation from the full teacher distribution produces a calibrated student against ground truth, whereas teacher-top-1 training produces an overconfident student.Teacher-top-1 is not a proper scoring metric or an unbiased estimator of the data distribution.
  • Calibration: Against teacher top-1, full-distribution students are underconfident, while top-1-trained students are calibrated to that evaluation target.The calibration outcome depends on matching the training objective to the evaluation metric.
  • Calibration: Increasing distillation tokens improves ground-truth calibration for full-distribution training but worsens it for teacher-top-1 training.When evaluated against teacher top-1, more tokens improve calibration only for students trained on that same metric.

F.3. Scaling law coefficients parameteric fit

The fitted scaling coefficients are estimated by bootstrapping and align with supervised scaling estimates. The resulting scaling law predicts student cross-entropy with low relative error, but its coefficients depend on the experimental setup.

  • Coefficient fitting: The supervised scaling estimates for model size and training data are consistent with Hoffmann et al. (2022).The comparison concerns the fitted supervised-case coefficients.
  • Coefficient fitting: Scaling coefficients are estimated with 90% confidence intervals using 4096 bootstrap resamples.The fitting procedure follows the cited resampling methodology.
  • Coefficient fitting: The fitted irreducible error is lower than in Hoffmann et al. (2022), which the authors suspect is related to µP.This is presented as a suspicion rather than a demonstrated causal explanation.

G.1. Mixing coefficient (λ) sensitivity analysis

The sensitivity analyses identify robust defaults for distillation while showing where objective and hyperparameter choices matter. Pure distillation, unit temperature, and the supervised-optimal learning rate perform robustly, whereas truncation and reverse KL reduce performance.

  • Mixing coefficient: Pure distillation with λ = 1 provides robust performance across model scales, although optimal λ varies by teacher-student pair.Student cross-entropy remains mostly flat for λ > 0.5; lower λ is preferred mainly for weak teachers.
  • Learning rate: η = 0.01 is optimal across the tested student and teacher sizes, with less than 0.1 nats variation over one order of magnitude.This transfers the supervised setting’s optimal learning rate to distillation under µP.
  • Distribution truncation: At k = 128, truncation raises validation loss by 0.11 nats, while λ = 0.7 keeps it within 0.01 nats of full distillation.The mixed truncated objective mitigates the degradation of Top-k truncation.
  • KL direction: Forward KL achieves a 0.28-nat average improvement over reverse KL and avoids computing student entropy.The authors recommend standard forward KL for distillation.

H.1. Alternative approximation for FLOPs per token as a function of N

This section derives forward-FLOP approximations from non-embedding parameter count N while retaining context- and vocabulary-dependent terms important for small models. It recommends fixed-aspect-ratio models because they improve scaling-law fit and simplify compute estimation.

  • Parameter accounting: The paper uses non-embedding parameters in its scaling laws and accounts for embedding and normalization contributions in approximate total compute.This treatment is intended to reduce estimation bias, especially for small models.
  • FLOP approximation: For small capable models, context-dependent and vocabulary-projection terms cannot generally be ignored without systematic error.The authors specifically identify 2n_layersn_ctxd_model and embedding-related terms as important in the relevant configuration range.
  • Design recommendation: Fixed-aspect-ratio models are recommended because they provide tighter scaling-law fits, simpler compute expressions, and more predictable hyperparameters at larger scales.The recommendation also reflects improved intuition that model size more directly corresponds to work performed.
  • Comparisons: The parameter-count and FLOP-count tables compare the paper’s expressions with commonly used alternatives and provide simplified substitutions for analysis.The attention accounting differs across methods, including causal versus bidirectional attention FLOP conventions.
  • FLOP approximation: The approximation expresses forward FLOPs per token as a function of non-embedding parameters N, including context and vocabulary contributions.The resulting expression supports estimating forward FLOPs from model size while dropping subleading terms.

I. Model architecture

The experiments use transformer models trained with a consistent optimization setup on English C4 data, while varying model sizes, training regimes, and distillation configurations. The study also includes dedicated analyses of scaling laws, capacity gaps, downstream metrics, and practical distillation choices.

  • Model implementation: All models use AXLearn implementations based on Gunter et al. and apply decoupled weight decay, simplified μP, RMSNorm, and multi-headed attention.The simplified μP setup fixes the learning rate at 1e−2 across model sizes.
  • Optimization: Training uses gradient rescaling, cosine learning-rate decay with warmup, and Z-loss regularization for stability.The gradient global norm is capped at 1.0, warmup lasts 2000 steps, and Z-loss is 10−4.
  • Data: Experiments use the English-only C4 subset, with separate teacher and student data splits for distillation training.C4 provides roughly 180B total tokens, divided into approximately 90B unique tokens for teacher and student training.
  • Experimental scope: The experimental program covers fixed-aspect-ratio models, FLOP counting, IsoFLOP design, teacher training, distillation, and longer 512B-token training runs.These components support the paper’s scaling-law analysis across teacher–student configurations.
  • Supplementary analyses: Additional analyses examine downstream metrics, capacity-gap phenomena, and practical distillation sensitivities such as mixing coefficients, temperature, learning rate, truncation, and reverse KL.These analyses extend beyond the core scaling-law fits to characterize distillation behavior and design choices.
Loading 2502.08606v2…