One-sentence Summary

Researchers at Beijing Institute of Technology and collaborators establish a law of capacity gap in language model distillation showing that the optimal teacher size scales linearly with the student size across different model and data scales, thereby eliminating exhaustive teacher search and yielding $7\text{B}$7B LLMs that outperform a wide array of competitors.

Key Contributions

• The paper reframes the curse of capacity gap in large language model distillation as an impossible triangle, clarifying why task-agnostic distillation with very large teachers becomes computationally costly and motivating a law-based approach to teacher selection.

• This work discovers and formalizes the law of capacity gap from experiments on models below 3B parameters, showing that the optimal teacher scale increases linearly with the student scale across different model and data scales, which supports extrapolation to 7B distillation.

• Experiments show that using this law enables compute-efficient 7B language model distillation, producing versatile models that improve fundamental language understanding and instruction following benchmarks while strengthening the compute-performance Pareto frontier.

Introduction

Language model distillation transfers knowledge from a large teacher model to a smaller student model, making large language models easier to deploy while preserving useful behavior. Prior work has revealed a key limitation called the curse of capacity gap: a larger teacher does not always produce a better student, so finding the best teacher often requires costly experiments across model sizes. This creates a difficult trade off between the desired student scale, the optimal teacher scale, and limited compute budget. The authors address this bottleneck by discovering a linear scaling law between the intended student scale and its optimal teacher scale. They leverage this law to choose teachers efficiently, reducing distillation overhead and producing small language models with strong performance on language understanding and instruction following tasks.

Dataset

• The authors use two main data groups: a Chinese-enhanced general-text mixture for model adaptation, pruning, and distillation, and an instruction-following collection for supervised finetuning.

• General mixture:

  • Sources: subsampled Pile, GitHub, and WuDao.
  • Purpose: used to adapt LLaMA2-7B with a Chinese vocabulary and to support pruning and distillation for creating MINIMA models at 0.4B, 7B, and 3B scales.
  • Scale and mixture ratios: not specified in the excerpt; the paper refers to Table 1 for statistics.

• Instruction-following data:

  • Sources: a broad collection of single-turn question-answering pairs and multi-turn conversation sessions.
  • Size: 1.1M examples total.
  • Purpose: used to finetune MINIMA into MINICHAT.
  • Statistics: detailed in Appendix B of the paper.

• Data processing and sequence construction:

  • For distillation, examples are packed into sequences of 4,096 tokens, with a batch of 1,024 sequences, corresponding to about 4M tokens.
  • For finetuning, examples are padded into sequences of 4,096 tokens, with a batch of 256 sequences, corresponding to about 1M tokens.
  • The excerpt does not mention additional cropping, metadata construction, or train, validation, and test split details.

Method

The authors leverage a capacity gap exploration pipeline that treats teacher scale as a tunable variable rather than an unbounded resource. For a student model of expected scale, the goal is to identify the teacher scale that yields the best distilled student performance with limited computational overhead. The pipeline begins by choosing a teacher language model of a specific scale, constructing a smaller student model by pruning the teacher, performing distillation, and then recording the resulting performance. By repeating this procedure across many teacher and student scale pairs, the method obtains evidence for whether a consistent optimal teacher scale exists for each target student scale.

The student construction stage uses structured pruning guided by parameter expressive scores. The student is pruned from the teacher by adding trainable scaling variables to architectural units. Attention heads are weighted by variables $\xi_j$ξj​, feedforward neurons are weighted by variables $\nu_k$νk​, and hidden states are weighted by a shared variable $\mu$μ across layers. The modified computation is expressed as

$$\mathrm{MHA}(\mathbf{x}_i)=\sum_{j=1}^{A}\xi_j\cdot \mathrm{Attn}(\mu \mathbf{x}_{\le i};\mathbf{W}_j^{\mathrm{Q}},\mathbf{W}_j^{\mathrm{K}},\mathbf{W}_j^{\mathrm{V}},\mathbf{W}_j^{\mathrm{O}}),$$MHA(xi​)=j=1∑A​ξj​⋅Attn(μx≤i​;WjQ​,WjK​,WjV​,WjO​), $$\mathrm{FFN}(\mathbf{x}_i)=\sum_{k=1}^{I}\nu_k\cdot g(\mathbf{x}_i\mathbf{W}_k^{\mathrm{I}})\mathbf{W}_k^{\mathrm{O}}.$$FFN(xi​)=k=1∑I​νk​⋅g(xi​WkI​)WkO​.

These variables expose the importance of each unit to the pruning objective. The expressive score of an attention head, neuron, or hidden state is computed as the expected absolute gradient of the pruning loss with respect to its corresponding variable:

$$I_j^{\mathrm{head}}=\mathbb{E}_{\mathbf{x}_i}\left|\frac{\partial \mathcal{L}^{\mathrm{pruning}}}{\partial \xi_j}\right|, \qquad I_k^{\mathrm{neuron}}=\mathbb{E}_{\mathbf{x}_i}\left|\frac{\partial \mathcal{L}^{\mathrm{pruning}}}{\partial \nu_k}\right|, \qquad I^{\mathrm{hidden}}=\mathbb{E}_{\mathbf{x}_i}\left|\frac{\partial \mathcal{L}^{\mathrm{pruning}}}{\partial \mu}\right|.$$Ijhead​=Exi​​​∂ξj​∂Lpruning​​,Ikneuron​=Exi​​​∂νk​∂Lpruning​​,Ihidden=Exi​​​∂μ∂Lpruning​​.

The pruning loss uses the teacher outputs and ground truth labels:

$$\mathcal{L}^{\mathrm{pruning}}=\sum_{i}^{|\mathcal{V}|}-\mathbf{y}_i^{\mathrm{truth}}\log \mathbf{y}_i^{\mathrm{teacher}}.$$Lpruning=i∑∣V∣​−yitruth​logyiteacher​.

Units with lower expressive scores are considered less important and are pruned first. The method applies $\ell_2$ℓ2​ normalization to the scores to reduce imbalanced pruning across layers. Since hidden state pruning interacts multiplicatively with attention or neuron pruning, the individual component sparsity is set to

$$p'=1-\sqrt{1-p},$$p′=1−1−p​,

where $p$p is the desired overall model sparsity. This adjustment helps achieve the target global sparsity while keeping embeddings relatively less compressed than transformer internal parameters.

The distillation stage uses a simple token level objective combining teacher soft targets and ground truth hard targets:

$$\mathcal{L}^{\mathrm{distillation}} = \frac{1}{2}\sum_{i}^{|\mathcal{V}|}-\mathbf{y}_i^{\mathrm{teacher}}\log \mathbf{y}_i^{\mathrm{student}} + \frac{1}{2}\sum_{i}^{|\mathcal{V}|}-\mathbf{y}_i^{\mathrm{truth}}\log \mathbf{y}_i^{\mathrm{student}}.$$Ldistillation=21​i∑∣V∣​−yiteacher​logyistudent​+21​i∑∣V∣​−yitruth​logyistudent​.

This objective preserves teacher signal from the output distribution while still anchoring the student with ground truth supervision. After training, the performance of each pruned student is compared to select the teacher scale that produces the best result for that student scale.

The collected pairs of student scale $\mathbb{S}$S and optimal teacher scale $\mathbb{T}^{\star}$T⋆ are then used to estimate the scaling relation. The method fits the observed pairs to a linear law:

$$\mathbb{T}^{\star}\approx \alpha \cdot \mathbb{S}+\beta.$$T⋆≈α⋅S+β.

As shown in the figure below:

The fit gives $\alpha=2.498$α=2.498, $\beta=-11.498$β=−11.498, and $R^2=0.9957$R2=0.9957, supporting an almost linear relationship between student scale and optimal teacher scale. The resulting estimated law is

$$\mathbb{T}^{\star}\approx 2.498\cdot \mathbb{S}-11.498\approx 2.5\cdot \mathbb{S}.$$T⋆≈2.498⋅S−11.498≈2.5⋅S.

Equivalently, if a teacher has scale $\mathbb{T}^{\star}$T⋆, then the corresponding student size for which this teacher is predicted to be optimal is approximately $0.4\cdot \mathbb{T}^{\star}$0.4⋅T⋆. This characterization allows teacher selection to be guided by a direct scaling estimate rather than by exhaustive trial of many teacher sizes.

Experiment

The experiments first distill pruned GPT2 and Pythia students across multiple model and data scales, then extend the capacity-gap analysis to larger LLaMA-based MINIMA and instruction-tuned MINICHAT models using broad language, reasoning, coding, and dialogue evaluations. The small-scale results show that student performance improves with scale, while increasing teacher scale follows an increasing-steady-decreasing pattern and reveals a stable optimal teacher scale, and the large-scale models perform competitively against efficient and instruction-tuned baselines. Ablations further indicate that symmetric pruning and simple distillation remain effective, supporting the robustness and practical value of the capacity-gap law.

The distilled 3B models are trained on a mixed corpus combining Pile, GitHub, and WuDao, totaling about 126B tokens over one epoch. Pile provides the largest proportion, while GitHub contributes the smallest share, supporting both adapted LLaMA2 and LLaMA3.1 distillation. The mixture is also used for pruning, helping shape compact student models intended for Chinese and general tasks. The training mixture combines general web text, code, and Chinese web data. Pile is the dominant source by proportion, while GitHub is the smallest source. The same mixture supports pruning and distillation for the 3B MINIMA student models.

MINIMA is described as a 3B distilled model that outperforms comparable pruning and layer-dropping baselines across knowledge, reasoning, and coding benchmarks. It remains stronger than a compute-comparable ShearedLLaMA setting, suggesting that distillation can compete with other efficiency approaches. The reported patterns also support the paper's capacity-gap explanation and the value of stronger teacher models. MINIMA outperforms layer-dropping and global-pruning baselines at similar model size and token budget. Increasing ShearedLLaMA's token budget does not close the gap to MINIMA. Stronger teacher language models can improve the distilled model's standard benchmark performance.

MINICHAT achieves clearly stronger macro-average pairwise GPT4 assessment scores than the scale-matched 3B baseline OpenBuddy-3B on both Vicuna-Bench and BELLE-Bench. It also remains broadly competitive with larger or differently scaled instruction models, although it is slightly behind BiLLa-7B on Vicuna-Bench and behind ChatGLM2-6B on BELLE-Bench. MINICHAT shows a large margin over the scale-matched OpenBuddy-3B baseline on both assessed benchmarks. MINICHAT remains competitive with scale-mismatched baselines, with only limited cases where the baseline scores higher. The assessment results support MINICHAT as a strong small instruction-following model despite its reduced scale.

The ablation results show that MINIMA remains broadly competitive when its pruning priority is changed. Global pruning can provide small advantages on some knowledge benchmarks, while layer dropping and teacher assistant are generally less competitive. Naive distillation is close to patient knowledge, suggesting that the pruning design is more central to the observed behavior than the distillation objective. MINIMA performs comparably to variants using other pruning strategies. Global pruning gives mixed effects, with better scores on some tasks and worse scores on others. Layer dropping and teacher assistant typically underperform MINIMA across the reported tasks.

Experiments evaluated 3B MINIMA and MINICHAT models trained on a mixed English, code, and Chinese corpus, using general knowledge, reasoning, coding, and instruction-following benchmarks alongside GPT-4-based pairwise assessments. MINIMA consistently outperforms comparable pruning, layer-dropping, and compute-comparable efficiency baselines, with stronger teachers improving performance and ablations showing pruning design matters more than the distillation objective. MINICHAT substantially surpasses the scale-matched OpenBuddy-3B baseline and remains broadly competitive with larger instruction models, indicating that distillation can yield strong small language models despite reduced scale.