MiniPLM: Knowledge Distillation for Pre-training Language Models

## Abstract

Abstract Knowledge distillation (KD) is widely used to train small, high-performing student language models (LMs) using large teacher LMs.
While effective in fine-tuning, KD during pre-training faces challenges in efficiency, flexibility, and effectiveness.
Existing methods either incur high computational costs due to online teacher inference, require tokenization matching between teacher and student LMs, or risk losing the difficulty and diversity of the teacher-generated training data.
To address these issues, we propose MiniPLM , a KD framework for pre-training LMs by refining the training data distribution with the teacher’s knowledge.
For efficiency, MiniPLM performs offline teacher LM inference, allowing KD for multiple student LMs without adding training-time costs.
For flexibility, MiniPLM operates solely on the training corpus, enabling KD across model families.
For effectiveness, MiniPLM leverages the differences between large and small LMs to enhance the difficulty and diversity of the training data, helping student LMs acquire versatile and sophisticated knowledge.
Extensive experiments demonstrate that MiniPLM boosts the student LMs’ performance on 9 widely used downstream tasks, improves the language modeling capabilities, and reduces pre-training computation.
The benefit of MiniPLM extends to large pre-training scales, evidenced by the extrapolation of the scaling curves.
Further analysis reveals that MiniPLM supports KD across model families and enhances the utilization of pre-training data. Our model, code, and data are available at https://github.com/thu-coai/MiniPLM .

## 1 Introduction

Recent advances in language models (LMs; , , , ) have largely been driven by scaling up model sizes, but this comes with high inference costs for deployment.
At the same time, training small, deployment-friendly LMs faces training computation challenges, as small models are typically far from compute-optimal configurations according to Scaling Laws []. This has spurred a growing interest in exploring the limits of small LMs under the constraint of pre-training computation [].

Figure: Figure 2: Results of applying KD methods in fine-tuning to pre-train a 200M student LM, using a 1.8B teacher LM. See Section 3.1 for method and evaluation details. When the training FLOPs are controlled, all KD methods perform similar or worse than Pre-Train w/o KD.
Refer to caption: x3.png

Knowledge Distillation (KD; ), where a small student model learns from a large teacher model, is a promising approach for training high-performing small LMs.
While KD is effective for fine-tuning LMs on specific tasks [], its role in improving pre-training, the critical stage for small LMs to acquire versatile foundation knowledge, remains under-explored.
It is non-trivial to apply KD during pre-training with the methods in fine-tuning stages, which can be categorized as online KD and offline KD.
In online KD, the teacher LM has to perform inference during the student LM pre-training to provide token-level probability supervision(^2^22Pre-computing this supervision is impractical, as it requires 30PB of storage for 50B tokens with a 150K vocabulary. Therefore, teacher LM’s probabilities are typically computed online during the student LM training.) [, ], introducing additional training-time computation overhead.
As shown in Figure 2, while online KD enhances performance within the same training steps, its benefits diminish if the extra computation was instead used to extend pre-training without KD.
In addition, most online KD methods require the teacher and student LMs to share tokenization, limiting its flexibility across different model families.
Offline KD, on the other hand, avoids extra training-time computation and allows for KD across model families, as student LMs are trained on the data offline generated by the teacher LM [, ].
However, ensuring sufficient difficulty and diversity in the teacher-generated data for pre-training is challenging without extensive human expertise.
This often causes the student LM to overfit easy and common language patterns, hindering its generalization to versatile downstream tasks [], as demonstrated in Figure 2.

To address these challenges, we propose MiniPLM, an efficient, flexible, and effective KD framework for pre-training LMs, as illustrated in Figure 3.
Intuitively, MiniPLM distills the teacher LM’s knowledge into the pre-training distribution through Difference Sampling, which samples training instances based on the “difference” between large and small LMs.
Student LMs are then pre-trained from scratch on the refined distribution.
To ensure efficiency, as shown in Figure 3(a), Difference Sampling performs offline teacher LM inference, allowing MiniPLM to distill knowledge into multiple student LMs without incurring additional training-time costs.
For flexibility, MiniPLM operates solely on the training corpus, enabling KD across model families and ensuring seamless integration with highly optimized training pipelines [].
In terms of effectiveness, Difference Sampling samples training instances that the teacher LM prefers but that a small reference LM assigns low probabilities to, promoting data difficulty and diversity.
As depicted in Figure 3(b), this design down-samples easy and common patterns, up-samples hard and diverse instances, and filters out noisy or harmful data points from the pre-training corpus, which encourages student LMs to acquire versatile and sophisticated knowledge, ultimately improving downstream generalization.

We apply MiniPLM to pre-train 200M, 500M, and 1.2B student LMs from scratch, using a 1.8B teacher LM.
We show that MiniPLM surpasses various baselines in improving student LMs’ zero-shot performance on 9 widely used downstream tasks, enhancing language modeling capabilities, and reducing pre-training computation.
By extrapolating the test loss with the Scaling Law [], we observe that MiniPLM’s benefit remains consistent for
LMs trained on ~10T tokens.
MiniPLM also facilitates KD across model families, improving Llama3.1 [] and Mamba [] with a teacher LM from the Qwen family [].
Further analysis shows that MiniPLM enhances pre-training data utilization, reducing the data demand by 2.4 times, which mitigates the quick exhaustion of web-crawled corpora [].

## 2 MiniPLM : KD for Pre-training LMs

We consider pre-training an LM with an output distribution $q_{\bm{\theta}}$ on a large-scale corpus $\mathcal{D}$ consisting of $N$ text sequences, where ${\bm{\theta}}$ represents the model parameters.
KD aids pre-training by incorporating the knowledge of a teacher LM with output distribution $p$, and training ${\bm{\theta}}$ to minimize the discrepancy between $p$ and $q_{\bm{\theta}}$ [].
MiniPLM formulates KD as a reward maximization problem [], which trains the student LM to generate diverse texts receiving high preference from the teacher LM (Section 2.1).
To efficiently and effectively optimize the reward, as illustrated by Figure 3, MiniPLM employs Difference Sampling to refine the pre-training corpus (Section 2.2). This process improves the pre-training distribution in an offline manner with the teacher LM’s knowledge and preserves the data diversity and difficulty with the help of a small reference model.
The student LM is then pre-trained from scratch on the refined corpus (Section 2.3).

Figure: (a) MiniPLM Training Framework
Refer to caption: x4.png

### 2.1 KD as Reward Maximization

Recent works [, , ] have shown the effectiveness of minimizing the reverse Kullback-Leibler divergence (KLD) between $p$ and $q_{\bm{\theta}}$ for KD of LMs in the fine-tuning stage, which avoids $q_{\bm{\theta}}$ from over-estimating the low-probability regions of $p$.
We reformulate this objective as a reward maximizing problem [], which is suitable for designing efficient offline optimization methods []:

$$ $\displaystyle{\bm{\theta}}=\arg\min_{{\bm{\theta}}}\operatorname{KL}\left[q_{ \bm{\theta}}||p\right]$ $\displaystyle=\arg\min_{{\bm{\theta}}}\mathbb{E}_{{\bm{x}}\sim q_{{\bm{\theta} }}}\log\frac{q_{\bm{\theta}}({\bm{x}})}{p({\bm{x}})}$ (1) $\displaystyle=\arg\max_{{\bm{\theta}}}\mathbb{E}_{{\bm{x}}\sim q_{\bm{\theta}} }r(p,q_{\bm{\theta}},{\bm{x}}),$ $$

where the reward $r$ is defined as $r(p,q_{\bm{\theta}},{\bm{x}})=\log\frac{p({\bm{x}})}{q_{\bm{\theta}}({\bm{x}})}$.
Intuitively, Eq. (1) trains $q_{\bm{\theta}}$ to generate text ${\bm{x}}$ that receives high reward values, indicating a preference from the teacher LM (i.e., high $\log p({\bm{x}})$ values), while also ensuring high diversity (i.e., low $\mathbb{E}_{{\bm{x}}\sim q_{{\bm{\theta}}}}q_{\bm{\theta}}({\bm{x}})$ values).
To optimize Eq. (1), a simple yet effective approach is Best-of-N [, , ], where a set $\mathcal{D}_{q_{\bm{\theta}}}$ containing $M$ candidates is first sampled from $q_{\bm{\theta}}$ and $K$ instances with the highest rewards are selected from these candidates to form a new dataset $\mathcal{D}^{\prime}_{q_{\bm{\theta}}}$:

$$ $\mathcal{D}^{\prime}_{q_{\bm{\theta}}}=\text{top-}K\{\left.r(p,q_{\bm{\theta}} ,{\bm{x}})\right|{\bm{x}}\in\mathcal{D}_{q_{\bm{\theta}}}\},$ (2) $$

where $\mathcal{D}_{q_{\bm{\theta}}}=\{\left.x_{m}\right|x_{m}\sim q_{\bm{\theta}},1
\leq m\leq M\}$. The student LM is then trained on $\mathcal{D}^{\prime}_{q_{\bm{\theta}}}$ to learn to generate texts with large $r(p,q_{\bm{\theta}},{\bm{x}})$ values.
Although Best-of-N achieves performing the teacher LM’s inference prior to the student LM training similar to offline KD [, ], it still lacks the efficiency advantage of these methods because obtaining $\mathcal{D}^{\prime}_{q_{\bm{\theta}}}$ requires (1) sampling data from $q_{\bm{\theta}}$ and (2) computing the reward values with $q_{\bm{\theta}}$, making $\mathcal{D}^{\prime}_{q_{\bm{\theta}}}$ non-transferable for pre-training other student LMs.
It is also hard to ensure the diversity and difficulty of the $M$ candidates sampled from $q_{\bm{\theta}}$ without careful prompt engineering with human expertise.
In the following, we show that these issues can be effectively and efficiently addressed by Difference Sampling, which is the basis of the MiniPLM training algorithm.

### 2.2 Difference Sampling

As shown in Figure 3(a), Difference Sampling refines the pre-training corpus $\mathcal{D}$ based on the discrepancy between $p$ and the output distribution $p_{\text{ref}}$ from a tiny reference LM, which eliminates the dependency of Eq. (2) on $q_{\bm{\theta}}$, making the sampled corpus reusable in training multiple student LMs.

#### Top- K 𝐾 K italic_K Sampling From 𝒟 𝒟 \mathcal{D} caligraphic_D , not 𝒟 q 𝜽 subscript 𝒟 subscript 𝑞 𝜽 \mathcal{D}_{q_{\bm{\theta}}} caligraphic_D start_POSTSUBSCRIPT italic_q start_POSTSUBSCRIPT bold_italic_θ end_POSTSUBSCRIPT end_POSTSUBSCRIPT .

To avoid sampling data from $\bm{q_{\bm{\theta}}}$, we sample instances with high $r(p,q_{\bm{\theta}},{\bm{x}})$ values from the pre-training corpus $\mathcal{D}$, rather than from the $\mathcal{D}_{q_{\bm{\theta}}}$ generated by the student LM as in Eq. (2).
This way, changing the student LM does not affect the candidate set, and $\mathcal{D}$ contains enough diverse and hard examples to be sampled for pre-training.
The following proposition offers theoretical support for this approach, showing that the sampled training instances from $\mathcal{D}_{q_{\bm{\theta}}}$ and $\mathcal{D}$ are highly likely to be the same when the sizes of $\mathcal{D}_{q_{\bm{\theta}}}$ and $\mathcal{D}$ are sufficiently large:

###### Proposition 2.1 .

Let $S$ be the sample space of two distributions $p_{1}$ and $p_{2}$, ${\mathbf{X}}_{1},{\mathbf{X}}_{2},\cdots,{\mathbf{X}}_{N}\sim p_{1}$ be $N$ i.i.d random variables, and ${\mathbf{Y}}_{1},{\mathbf{Y}}_{2},\cdots,{\mathbf{Y}}_{M}\sim p_{2}$ be $M$ i.i.d random variables. Let $r(\cdot):S\mapsto\mathbb{R}$ be any injective function. Assume that $\forall{\bm{x}}\in S$, $p_{1}({\bm{x}})>0$, $p_{2}({\bm{x}})>0$. For a fixed $K$ satisfying $1\leq K\leq\min\left\{N,M\right\}$, when $N\rightarrow+\infty$, $M\rightarrow+\infty$, we have

$$ $P\left(\text{top-}K\left\{\left.r({\mathbf{X}}_{n})\right|1\leq n\leq N\right \}=\text{top-}K\left\{\left.r({\mathbf{Y}}_{m})\right|1\leq m\leq M\right\} \right)\rightarrow 1.$ (3) $$

The proof of Proposition 3 is provided in Appendix A. In our context, $p_{1}$ represents the data distribution of $\mathcal{D}$, $p_{2}=q_{\bm{\theta}}$, which is the data sampling distribution of $\mathcal{D}_{q_{\bm{\theta}}}$, and $r({\bm{x}})=r(p,q_{\bm{\theta}},{\bm{x}})$.
Intuitively, Proposition 3 reveals the fact that when $|\mathcal{D}|=N$ and $|\mathcal{D}_{q_{\bm{\theta}}}|=M$ are sufficiently large, the top-$K$ instances selected from both sets tend to overlap.
This is well-suited to our scenario, as the pre-training corpus $\mathcal{D}$ is typically large-scale, and it is advantageous to sample as many candidates from $\mathcal{D}_{q_{\bm{\theta}}}$ as possible to find high-reward instances.

###### Proposition 2.1 .

#### Decoupling the Student and the Reward-Computing LM.

To avoid computing reward values with $\bm{q_{\bm{\theta}}}$, we replace it with $p_{\text{ref}}$, the output distribution of a tiny reference LM, typically smaller than the student LM, for reward computation.
The reference LM is pre-trained on a small subset $\mathcal{D}_{\text{ref}}$, uniformly sampled from $\mathcal{D}$, allowing $p_{\text{ref}}$ to approximate $q_{\bm{\theta}}$ with minimal computation.
This is a reasonable approximation because $q_{\bm{\theta}}$ evaluates the difficulty of instances, and the relative data difficulties generally remain consistent across different models []. As a result, the reward function in Eq. (2) is replaced with $r(p,p_{\text{ref}},{\bm{x}})=\log\frac{p({\bm{x}})}{p_{\text{ref}}({\bm{x}})}$.

In summary, Difference Sampling constructs a pre-training corpus $\mathcal{D}^{\prime}$ from $\mathcal{D}-\mathcal{D}_{\text{ref}}$ as follows:

$$ $\mathcal{D}^{\prime}=\text{top-}K\left\{\left.\log\frac{p({\bm{x}})}{p_{\text{ ref}}({\bm{x}})}\right|{\bm{x}}\in\mathcal{D}-\mathcal{D}_{\text{ref}}\right\},$ (4) $$

which is independent of $q_{\bm{\theta}}$. As shown in Figure 3(b), Difference Sampling essentially refines the data distribution of $\mathcal{D}$ by comparing the “difference” between $p$ and $p_{\text{ref}}$, increasing difficulty and diversity while filtering out noise. A detailed discussion of these effects is provided in Section 2.5.

### 2.3 Pre-Training on Difference-Sampled Corpus

As illustrated in Figure 3(a), we pre-train the student LM from scratch on the difference-sampled corpus $\mathcal{D}^{\prime}$ with the cross-entropy loss for next-token prediction, which is similar to standard pre-training. The loss function $L(q_{\bm{\theta}},\mathcal{D}^{\prime})$ is given by

$$ $L(q_{\bm{\theta}},\mathcal{D}^{\prime})=-\frac{1}{|\mathcal{D}^{\prime}|}\sum_ {{\bm{x}}\in\mathcal{D}^{\prime}}\frac{1}{|{\bm{x}}|}\sum_{t=1}^{|{\bm{x}}|} \log q_{\bm{\theta}}(x_{t}|{\bm{x}}_{<t}),$ (5) $$

where $|{\bm{x}}|$ is the length of ${\bm{x}}$, $x_{t}$ is the $t^{\text{th}}$ token, and ${\bm{x}}_{<t}$ denotes the prefix of ${\bm{x}}$ with $t-1$ tokens.

### 2.4 MiniPLM Training Pipeline

The general training pipeline of MiniPLM is as follows:
(1) Uniformly sample a subset $\mathcal{D}_{\text{ref}}$ from the pre-training corpus $\mathcal{D}$, ensuring $|\mathcal{D}_{\text{ref}}|\ll|\mathcal{D}|$.
(2) Train a reference LM from scratch on $\mathcal{D}_{\text{ref}}$ using the cross-entropy loss to obtain its output distribution $p_{\text{ref}}$.
(3) Perform Difference Sampling with $p_{\text{ref}}$ and the teacher LM’s output distribution $p$, generating a pre-training corpus $\mathcal{D}^{\prime}$ from $\mathcal{D}-\mathcal{D}_{\text{ref}}$ using Eq. (4). The size of $\mathcal{D}^{\prime}$ is controlled by a sampling ratio $\alpha$, where $K=\alpha|\mathcal{D}-\mathcal{D}_{\text{ref}}|$.
(4) Pre-train student LMs from scratch on $\mathcal{D}^{\prime}$ with the cross-entropy loss defined in Eq. (5).

### 2.5 Discussion

#### Efficiency and Flexibility of MiniPLM .

As shown in Figure 3(a), MiniPLM relies only on $p({\bm{x}})$ and $p_{\text{ref}}({\bm{x}})$, the teacher and reference LMs’ probability over the entire sequence, which can be computed and stored offline because each instance in $\mathcal{D}$ is associated with a single floating-point number, amounting to only 200MB storage for 50B tokens with a 1,024 sequence length.
Once $\mathcal{D}^{\prime}$ is difference-sampled based on $p({\bm{x}})$ and $p_{\text{ref}}({\bm{x}})$, multiple student LMs can be efficiently pre-trained under the teacher LM’s guidance without extra computational cost.
In addition, this process modifies only the pre-training data, imposing no restrictions on the architecture or tokenization of the student LM, making MiniPLM highly flexible for integration into optimized pre-training frameworks [] and suitable for KD across model families. In contrast, online KD [, ] relies on per-token distributions, which are infeasible to store offline, as it takes $50B\times 150K\times 4\text{byte}=30\text{PB}$ storage for an LM with 150K vocabulary trained on 50B tokens, making online teacher LM inference necessary. Aligning per-token distributions also demands matching tokenizers between the teacher and student [], which complicates the re-implementation and optimization of pre-training workflows.

#### Effectiveness of MiniPLM .

In essence, as illustrated in Figure 3(b), MiniPLM distills the teacher LM’s knowledge into the student LM’s pre-training distribution via Difference Sampling, producing three effects:
(1) Down-sampling easy and common patterns that both the teacher and reference LM fit well, where $p({\bm{x}})\gtrsim p_{\text{ref}}({\bm{x}})$ and $\log\frac{p({\bm{x}})}{p_{\text{ref}}({\bm{x}})}\gtrsim 0$.
(2) Up-sampling hard and diverse knowledge, which the larger teacher LM has mastered but the smaller reference LM struggles with, and thus $p({\bm{x}})\gg p_{\text{ref}}({\bm{x}})$ and $\log\frac{p({\bm{x}})}{p_{\text{ref}}({\bm{x}})}\gg 0$.
(3) Discarding noisy and harmful instances that the teacher LM assigns lower probabilities to than the reference LM, where $\log\frac{p({\bm{x}})}{p_{\text{ref}}({\bm{x}})}<0$. We provide examples of these effects in Appendix C.5.
Training on the distribution with these effects encourages the student model to focus more on the sophisticated world knowledge learned by the teacher without being distracted by the noise.
Note that without the comparison between $p$ and $p_{\text{ref}}$, the effect (1) disappears, leading to a pre-training corpus dominated by common patterns.
This explains the effectiveness of MiniPLM against prior offline KD methods [, ], where maintaining the data difficulty and diversity, critical to pre-training [], is challenging without extensive human efforts on prompt engineering [].

## 3 Experiments

### 3.1 Experimental Setup

#### Model.

We adopt the Qwen-1.5 [] architecture in our experiments. We use the officially released 1.8B Qwen-1.5 model as the teacher LM and distill its knowledge into students with 200M, 500M, and 1.2B parameters. Detailed model configurations are provided in Appendix B.

#### Pre-Training.

We construct pre-training corpora from the Pile []. To control the computation in experiments, we pre-train all LMs on a maximum of 50B tokens, where documents are merged to construct instances with sequence lengths of 1,024. For online KD methods that incur additional train-time computation, we reduce their training steps to align the total training computation with pre-training without KD or offline KD methods. See Appendix B for more pre-training details.

#### Baselines.

We compare MiniPLM with 4 baselines:

- •
Pre-Train w/o KD pre-trains the student LM on a 50B corpus uniformly sampled from the Pile dataset, without the guidance of the teacher LM’s knowledge.
- •
Vanilla KD [] minimizes the token-level forward KLD between $p$ and $q_{\bm{\theta}}$, which requires online inference of the teacher LM to obtain the token-level output distributions.
- •
SeqKD [] trains the student LM on the teacher-generated data. Since it is infeasible to generate all 50B tokens, we approximate Kim and Rush [] by using the first 768 tokens of each instance from the training corpus in Pre-Train w/o KD as the prompts and let the teacher LM generate the remaining tokens offline.
- •
MiniLLM [] minimizes the reverse KLD between $p$ and $q_{\bm{\theta}}$ with PPO [], which requires online inference of the teacher LM and online sampling from the student LM. We treat the first 768 tokens of the instances from the training corpus of Pre-Train w/o KD as the prompts and sample 256 tokens from $q_{\bm{\theta}}$ during the exploration of PPO.

#### MiniPLM .

We employ a 104M reference LM trained on 5B tokens.
In Section 3.2 and 3.3, we consider a setting where $\mathcal{D}$ is sufficiently large, containing 105B tokens uniformly sampled from the Pile corpus.
We reserve 5B tokens as $\mathcal{D}_{\text{ref}}$, and conduct Difference Sampling as per Eq. (4) on the other 100B tokens by setting $\alpha=0.5$ to construct a 50B-token corpus $\mathcal{D}^{\prime}$.
In this way, the student LM is pre-trained on $\mathcal{D}^{\prime}$ for one epoch.
We use the loss difference of the teacher and reference LM to sampled instances from $\mathcal{D}$, which is equivalent to Eq. (4) as all instances have 1,024 tokens.
In Section 3.4, we evaluate MiniPLM in a data-limited setting, where $\mathcal{D}$ is controlled to contain 50B tokens and the student LM is trained on the difference-sampled data for multiple epochs.
See Appendix C.4 for the ablation studies on the reference LM and sampling ratio $\alpha$.

#### Evaluation.

We assess the zero-shot accuracy of LMs trained with different methods on 9 downstream tasks widely used in examining the foundation abilities of base models [, ].
We also test the language modeling capability of the LMs on a subset of DCLM [], a high-quality corpus carefully curated with complex pipelines, to examine how well LMs capture broad and diverse knowledge. See Appendix B for more evaluation details.

**Table 1: Zero-shot accuracy scores on 9 widely-used downstream tasks and the average scores (Avg.). We use the Qwen-1.5 1.8B LM [] as the teacher and Qwen LMs with 200M, 500M, and 1.2B parameters as the student. Student LMs with the same sizes consume the same training-time computation. The best scores of each model size are boldfaced.**
|  | HS | LAM | Wino | OBQA | ARC-e | ARC-c | PIQA | SIQA | Story | Avg. |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| 1.8B Teacher $\rightarrow$ 200M Student |  |  |  |  |  |  |  |  |  |  |
| Pre-Train w/o KD | 31.1 | 32.4 | 49.9 | 27.6 | 38.9 | 23.1 | 61.8 | 36.4 | 58.1 | 39.9 |
| Vanilla KD | 30.4 | 31.0 | 51.4 | 26.6 | 40.1 | 23.1 | 62.2 | 36.9 | 57.3 | 39.9 |
| MiniLLM | 30.2 | 29.4 | 50.0 | 26.6 | 39.0 | 21.3 | 60.5 | 36.6 | 57.6 | 39.0 |
| SeqKD | 30.5 | 31.0 | 51.3 | 27.4 | 39.3 | 22.4 | 61.3 | 36.9 | 57.4 | 39.7 |
| MiniPLM | 32.7 | 35.4 | 51.4 | 27.2 | 40.6 | 23.7 | 63.3 | 37.0 | 60.0 | 41.3 |
| 1.8B Teacher $\rightarrow$ 500M Student |  |  |  |  |  |  |  |  |  |  |
| Pre-Train w/o KD | 35.8 | 40.1 | 51.0 | 30.2 | 41.7 | 24.4 | 65.4 | 38.2 | 61.4 | 43.2 |
| Vanilla KD | 37.0 | 39.9 | 51.7 | 29.4 | 45.1 | 24.2 | 65.8 | 38.0 | 61.6 | 43.6 |
| MiniLLM | 33.0 | 35.4 | 51.2 | 27.5 | 42.1 | 24.2 | 62.3 | 37.3 | 60.2 | 41.5 |
| SeqKD | 34.9 | 37.9 | 50.7 | 28.6 | 42.7 | 23.6 | 65.0 | 38.4 | 58.9 | 42.3 |
| MiniPLM | 39.0 | 42.6 | 52.2 | 30.2 | 45.8 | 24.9 | 67.0 | 39.0 | 62.2 | 44.8 |
| 1.8B Teacher $\rightarrow$ 1.2B Student |  |  |  |  |  |  |  |  |  |  |
| Pre-Train w/o KD | 39.4 | 44.5 | 51.8 | 28.4 | 46.0 | 25.7 | 67.0 | 39.5 | 62.2 | 44.9 |
| Vanilla KD | 40.7 | 43.3 | 53.2 | 29.8 | 46.1 | 25.5 | 67.3 | 39.2 | 63.5 | 45.4 |
| MiniLLM | 36.1 | 42.5 | 51.2 | 28.5 | 44.1 | 25.3 | 65.8 | 37.9 | 61.4 | 43.6 |
| SeqKD | 38.5 | 41.4 | 51.9 | 29.2 | 46.5 | 25.1 | 66.3 | 39.0 | 61.0 | 44.3 |
| MiniPLM | 42.8 | 46.2 | 53.3 | 31.0 | 46.8 | 26.9 | 68.3 | 39.8 | 64.0 | 46.6 |

### 3.2 Main Results

#### MiniPLM Improves Downstream Performance.

Table 1 shows zero-shot accuracy on downstream tasks for LMs trained by different methods, leading to three key observations.
First, among all the baselines, only Vanilla KD outperforms Pre-Train w/o KD for relatively large student LM, given a constant computation budget. This highlights the room for improvement in KD for pre-training, especially when the gap between the teacher and the student is substantial.
Second, MiniPLM-trained model achieves the best performance across most of the tasks.
Compared to Pre-Train w/o KD, MiniPLM effectively leverages the teacher LM’s knowledge to improve the student LM pre-training.
Compared to online methods like Vanilla KD and MiniLLM, MiniPLM, as shown in Figure 3(a), incurs no additional training-time overhead, allowing the student LM to be optimized for more steps and leading to higher performance.
Compared to offline methods like SeqKD, MiniPLM, as shown in Figure 3(b), uses a reference LM to ensure sufficient difficulty and diversity of the pre-training corpus, which is essential for the student LM to learn versatile sophisticated knowledge during pre-training and generalize across various downstream tasks.
Finally, the improvements of MiniPLM against Vanilla KD scale well with the student LM size, which is also illustrated in Figure 1(b).

Figure: (a) 1.8B Teacher $\rightarrow$ 200M Student
Refer to caption: x6.png

**Table 2: Test loss predictions using the Scaling Law []. $N_{\text{stu}}$: the student LM size. $L_{\text{1T}}$, $L_{\text{10T}}$: the loss when Pre-Train w/o KD and MiniPLM process 1T and 10T tokens, with Vanilla KD consuming the same training FLOPs.**
| $N_{\text{stu}}$ | Method | $L_{\text{1T}}$ | $L_{\text{10T}}$ |
| --- | --- | --- | --- |
| 200M | Pre-Train w/o KD | 3.35 | 3.32 |
| Vanilla KD | 3.39 | 3.35 |  |
| MiniPLM | 3.28 | 3.26 |  |
| 500M | Pre-Train w/o KD | 3.12 | 3.08 |
| Vanilla KD | 3.12 | 3.07 |  |
| MiniPLM | 3.06 | 3.04 |  |
| 1.2B | Pre-Train w/o KD | 2.98 | 2.94 |
| Vanilla KD | 2.95 | 2.91 |  |
| MiniPLM | 2.92 | 2.88 |  |

#### MiniPLM Helps Language Modeling.

Figure 4, compares the language modeling performance of the MiniPLM-trained LMs and baselines on the DCLM [] subset, a diverse and high-quality dataset curated from web corpora.
The results show that, given the same training-time FLOPs, LMs trained with MiniPLM achieve the lowest test losses.
In Table 2, we extrapolate the test losses with the Scaling Law [] to simulate pre-training on 1T and 10T tokens (details in Appendix C.2), showing thatMiniPLM maintains its advantages at the scales of pre-training recent large LMs [, ].
A critical stage in DCLM’s data-cleaning pipeline involves removing easy and common patterns, thereby enhancing challenging and diverse signals.
Therefore, a lower test loss on DCLM suggests that, with the teacher LM’s guidance and the reference LM, the MiniPLM-trained LMs learn the diverse and hard knowledge better due to the up-sampling of the corresponding parts in the pre-training distribution, as illustrated in Figure 3(b).

#### MiniPLM Reduces Training Computation.

We plot the 500M student LM’s average zero-shot accuracy scores on the downstream tasks in Figure 1(a) with respect to its pre-training FLOPs.
MiniPLM achieves the same performance as Vanilla KD while reducing computational costs by 2.2 times.
Similar trends are observed for other student LMs (Figure 7) and in the test loss curves on DCLM corpus (Figure 4). The efficiency gains are more pronounced when comparing MiniPLM with Pre-Train w/o KD on the 500M and 1.2B models.
We attribute this acceleration to Difference Sampling, which down-samples common patterns and filters out noisy signals from the pre-training corpus, as shown in Figure 3(b).
As a result, the model trained with MiniPLM avoids wasting computation on learning the easy knowledge quickly memorized during the early training stage and is less distracted by the noisy outliers that slow the convergence down.

Figure: Table 3: Results of KD across model families. We use the teacher and reference LM from the Qwen family to distill the Llama3.1 and Mamba models. The average zero-shot accuracies on the downstream tasks and the losses on the DCLM corpus are reported. Note that Vanilla KD and MiniLLM cannot be applied when the teacher and student LMs use different tokenizations.
Refer to caption: x9.png

### 3.3 KD Across Model Families

A noticeable advantage of MiniPLM over Vanilla KD and MiniLLM is its flexibility to distill the knowledge of a teacher LM into student LMs with completely different tokenizers and architectures without additional strategies like Boizard et al. [].
In Table 5, we illustrate the performance when using the LMs from the Qwen [] family as the teacher and the reference LM to distill knowledge into a 212M Llama3.1 [] model and a 140M Mamba [] model.
The results demonstrate the promising performance of MiniPLM in KD across model families, outperforming Pre-Train w/o KD and the existing offline KD baseline (SeqKD).
This allows emerging LMs with novel architectures [, ] or advanced tokenization [, ] to inherit knowledge from existing LMs, thereby facilitating the development of more efficient and higher-performed models.

### 3.4 Data-Limited Setting

We evaluate MiniPLM in a data-limited setting where $\mathcal{D}$ is constrained to contain 50B tokens.
To this end, the student LM should be trained on the difference-sampled $\mathcal{D}^{\prime}$ over multiple epochs to ensure the total computation and trained tokens remain consistent with Pre-Train without KD.
We split a $|\mathcal{D}_{\text{ref}}|$ containing 1B tokens to train the reference LM.
Therefore, for a sampling ratio $\alpha$, the student LM should be trained for $\frac{|\mathcal{D}|}{\alpha|\mathcal{D}-\mathcal{D}_{\text{ref}}|}\approx\frac
{1}{\alpha}$ epochs, given that $|\mathcal{D}_{\text{ref}}|\ll|\mathcal{D}|$.
In Figure 5, we plot the loss curves of the 200M student LMs on the DCLM corpus when using $\alpha\in\left[0.5,0.25,0.125\right]$ and training the LM for around 2, 4, and 8 epochs, respectively.
We can see that difference-sampling 25% data (with $\alpha=0.25$) and training the student LM for 4 epochs yields the best performance, which aligns with the observations in Muennighoff et al. [].
The corpus sampled with a higher $\alpha$ does not achieve the best quality and diversity offered by Difference Sampling, while a lower $\alpha$ leads to rapid over-fitting of the student LM.
These findings suggest that MiniPLM is a promising approach to enhance data utilization when high-quality web corpora become scarce [].
By extrapolating the loss curve of Pre-Train w/o KD using the Scaling Law in a data-constrained setting [], we estimate that it would require an additional 68B training tokens to match the performance of MiniPLM ($\alpha=0.25$, 4 epochs), which means MiniPLM reduces the pre-training data requirement by 2.4 times.
See Appendix C.2 for more details on this extrapolation.

### 3.5 Analysis

#### Impact of Teacher Model.

Intuitively, larger teacher LMs will lead to better KD results, which is observed in recent works [].
However, early works have also shown that an excessively large gap between teacher and student models can hinder effective KD [].
In Figure 6, we plot the performance of Vanilla KD and MiniPLM when distilling teacher LMs with different sizes into a 200M student model.
We observe a similar phenomenon to Mirzadeh et al. [] on Vanilla KD and MiniPLM that larger teacher LMs are not necessarily more helpful for KD.
However, we attribute this to different factors for Vanilla KD and MiniPLM.
In Vanilla KD, the overhead introduced by larger LMs during training diminishes the benefits of distillation.
As for MiniPLM, a 500M teacher LM proves most effective for distilling into a 200M student LM.
Smaller teacher LMs (e.g., 300M) lack the capacity to identify the hard but valuable parts of the pre-training distribution, weakening the effectiveness of Difference Sampling as shown in Figure 3(b).
On the other hand, the value scale of $\log p({\bm{x}})$ from oversized LMs (e.g., 4B) becomes too small compared to that of the reference LM, which tends to degenerate Difference Sampling into sampling with $p_{\text{ref}}({\bm{x}})$ only, losing the effect of the teacher LM. Future research could focus on optimizing teacher model size for MiniPLM or mitigating the impact of differing $\log p({\bm{x}})$ and $\log p_{\text{ref}}({\bm{x}})$ value scales.

Figure: Figure 6: Impact of the teacher LM’s sizes on Vanilla KD and MiniPLM, with the pre-training FLOPs aligned. The y-axis represents the average zero-shot accuracy on the downstream tasks.
Refer to caption: x10.png

#### Diversity of Difference-Sampled Data.

To further verify the effectiveness of Difference Sampling on improving the diversity of the pre-training corpus, in Table 6, we follow Friedman and Dieng [] to compute the semantic diversity of the difference-sampled corpus and the data used in other baseline pre-training approaches.
The results show that the difference-sampled corpus has the highest diversity, despite that Difference Sampling is derived from minimizing the reverse KLD, which exhibits the mode-seeking behavior [].
We suspect the reason is that Difference Sampling down-samples the easy parts of the corpus containing repeated contents while up-sampling the hard parts consisting of diverse texts.
These two components, with large $p({\bm{x}})$ values as seen in Figure 3(b), constitute the major modes of the teacher LM that $q_{\bm{\theta}}$ seeks during optimizing Eq. (1).
Therefore, the mode-seeking behavior helps remove the noisy parts of the pre-training distribution, and the loss of diversity due to noise reduction is compensated by the up-sampling of the hard and diverse data points.
We provide a case study in Appendix C.5 to further explain our argument.

**Table 5: Average accuracy on downstream tasks when combining MiniPLM and Vanilla KD. “MiniPLM + Vanilla KD”: applying Vanilla KD to pre-train student LMs on the difference-sampled corpus in MiniPLM. $N_{\text{stu}}$: the size of student LMs.**
| $N_{\text{stu}}$ | Method | Acc. |
| --- | --- | --- |
| 200M | Vanilla KD | 39.9 |
| MiniPLM | 41.3 |  |
| MiniPLM + Vanilla KD | 40.7 |  |
| 500M | Vanilla KD | 43.6 |
| MiniPLM | 44.8 |  |
| MiniPLM + Vanilla KD | 44.9 |  |
| 1.2B | Vanilla KD | 45.4 |
| MiniPLM | 46.6 |  |
| MiniPLM + Vanilla KD | 48.1 |  |

#### Combining Vanilla KD and MiniPLM .

As shown in Table 1 and Figure 4, Vanilla KD improves the 500M and 1.2B student LM performance compared to Pre-Train w/o KD, suggesting the potential of further improving MiniPLM by combining it with Vanilla KD.
This combined approach, termed “MiniPLM + Vanilla KD”, applies Vanilla KD to the difference-sampled corpus used in MiniPLM.
In Table 5, we compare “MiniPLM + Vanilla KD” with the individual use of Vanilla KD and MiniPLM, using a 1.8B teacher LM.
The results show that when pre-training student LMs with 500B and 1.2B parameters, “MiniPLM + Vanilla KD” further improves the performance, given the same training-time FLOPs.
This demonstrates that MiniPLM and Vanilla KD complement each other:
MiniPLM distills the coarse-grained sequence-level knowledge of the teacher LM into the student LM via the pre-training data, while Vanilla KD directly aligns the token-level probability distribution between $p$ and $q_{\bm{\theta}}$, providing fine-grained token-level signals.

## 4 Related Work

#### Language Model Pre-Training.

Pre-training is the critical phase for language models (LMs; , , , , , , ) to obtain their foundation abilities for various downstream tasks.
To improve pre-training, some works focus on data curation, such as adjusting domain mixing [, ], selecting valuable data points relevant to desired tasks [, , ], or transforming the instances based on downstream requirements [, , ].
Another line of work improves the optimization during pre-training by solving better data reweighting strategies [], designing more effective optimizers [, ], or discovering better training recipes [, ].
Variations in model architectures [] and training objectives [] are also explored to boost pre-training stability and final LM performance.
In this work, we utilize the knowledge from existing LMs to enhance pre-training.

#### Small Language Models.

Given the high computational demands of large LMs during inference, there has been growing interest in pre-training small LMs [].
However, achieving high performance with limited parameter sizes remains challenging because the training computation that small LMs need to match the capabilities of large LMs often exceeds Chinchilla-optimal [] and scales as a power law with respect to the model size gap [].
Despite these challenges, recent efforts have made promising progress by data quality improvement [, , ] or model pruning [, ]. We explore knowledge distillation as a complementary approach.

#### Knowledge Distillation.

Knowledge distillation (KD;) uses a large teacher model to improve the performance of a small student model, which is widely used to build efficient neural network systems [, , ]. In NLP, early works primarily apply KD to encoder-only models [, ] for text classification by aligning token-level distribution [], hidden states [], and attention matrices [, ]. For generative LMs, a straightforward KD approach is training small LMs on the texts generated by large LMs [, , ]. Other works [, , , , ] explore better optimization objectives. However, these works focus on KD for fine-tuning LMs, while pre-training is critical for establishing core LM capabilities []. Therefore, we investigate KD in the pre-training stage to develop strong small base LMs.

## 5 Conclusion

#### Summary.

In this work, we find it non-trivial to adapt existing KD approaches for fine-tuning LMs to the pre-training stage because of the high overhead brought by the teacher LM inference in online KD and the tendency of losing data difficulty and diversity in offline KD.
Therefore, we propose MiniPLM to address these issues through Difference Sampling, which refines the training distribution by down-sampling easy patterns, up-sampling hard instances, and filtering out noisy data points, with the knowledge of the difference between the large teacher LM and a small reference LM.
The offline nature of MiniPLM makes it both efficient and flexible to distill student LMs with diverse configurations.
The use of the large-small-model differences ensures the difficulty and diversity of the refined pre-training distribution.
Using a 1.8B LM as the teacher to guide the pre-training of 200M, 500M, and 1.2B LMs, we demonstrate that MiniPLM improves the student LMs’ performance on 9 downstream tasks, enhances their language modeling capability, and reduces pre-training computation.
Additionally, MiniPLM improves the utilization of limited pre-training data and can distill teacher LM’s knowledge into student LMs from completely different families.

#### Limitations.

One limitation of MiniPLM is that it requires the large LMs’ probabilities on texts from the pre-training corpus, making black-box KD for close-source LMs [, ] with MiniPLM challenging. For some APIs [], this issue can be solved by specifying user-provided bias in the softmax operation, which allows obtaining the probabilities of given tokens one by one [], at the expense of a large number of API calls.

#### Future Work.

A promising future direction to explore is applying the difference-sampled corpus to pre-train LMs larger than the teacher LM, enabling weak-to-strong generalization []. Since data properties are critical for pre-training LMs across various sizes, the improvement of data diversity and difficulty is likely to be beneficial for LMs larger than the difference-sampling models.

## Appendix A Proof of Proposition 3

To prove Proposition 3, We start from the $K=1$ case, corresponding to selecting an instance with the maximal reward and Eq. (3) becomes

$$ $P\left(\arg\max_{1\leq n\leq N}r({\mathbf{X}}_{n})=\arg\max_{1\leq m\leq M}r({ \mathbf{Y}}_{m})\right)\rightarrow 1,$ (6) $$

which is equivalent to

$$ $\sum_{{\bm{x}}\in S}P\left(\arg\max_{1\leq n\leq N}r({\mathbf{X}}_{n})={\bm{x} }\text{ and }\arg\max_{1\leq m\leq M}r({\mathbf{Y}}_{m})={\bm{x}}\right) \rightarrow 1.$ (7) $$

Since ${\mathbf{X}}_{n}$ and ${\mathbf{Y}}_{m}$ are independent random variables for $1\leq n\leq N$ and $1\leq m\leq M$, Eq. (7) can be further written as

$$ $\sum_{{\bm{x}}\in S}P\left(\arg\max_{1\leq n\leq N}r({\mathbf{X}}_{n})={\bm{x} }\right)P\left(\arg\max_{1\leq m\leq M}r({\mathbf{Y}}_{m})={\bm{x}}\right) \rightarrow 1.$ (8) $$

We focus on the term $P\left(\arg\max\limits_{1\leq n\leq N}r({\mathbf{X}}_{n})={\bm{x}}\right)$, which can be expanded as follows based on the fact that ${\mathbf{X}}_{1},{\mathbf{X}}_{2},\cdots,{\mathbf{X}}_{N}$ are i.i.d random variables:

$$ $\displaystyle P\left(\arg\max_{1\leq n\leq N}r({\mathbf{X}}_{n})={\bm{x}}\right)$ (9) $\displaystyle=$ $\displaystyle P\left(r({\mathbf{X}}_{n})\leq r({\bm{x}}),\text{ for $1\leq n \leq N$ and at least one variable in ${\mathbf{X}}_{1},{\mathbf{X}}_{2},\cdots ,{\mathbf{X}}_{N}$ equals ${\bm{x}}$}\right)$ $\displaystyle=$ $\displaystyle P\left(r({\mathbf{X}}_{n})\leq r({\bm{x}}),\text{ for $1\leq n \leq N$}\right)-P\left(r({\mathbf{X}}_{n})\leq r({\bm{x}})\text{ and }{\mathbf {X}}_{n}\neq{\bm{x}},\text{ for $1\leq n\leq N$}\right)$ $\displaystyle=$ $\displaystyle\prod_{n=1}^{N}P(r({\mathbf{X}}_{n})\leq r({\bm{x}}))-\prod_{n=1} ^{N}\left[P(r({\mathbf{X}}_{n})\leq r({\bm{x}}))-P({\mathbf{X}}_{n}={\bm{x}})\right]$ $\displaystyle=$ $\displaystyle P(r({\mathbf{X}}_{1})\leq r({\bm{x}}))^{N}-\left[P(r({\mathbf{X} }_{1})\leq r({\bm{x}}))-p_{1}({\bm{x}})\right]^{N}$ $\displaystyle=$ $\displaystyle P_{{\mathbf{X}}}({\bm{x}})^{N}\left[1-\left[1-\frac{p_{1}({\bm{x }})}{P_{{\mathbf{X}}}({\bm{x}})}\right]^{N}\right],$ $$

where $P_{{\mathbf{X}}}({\bm{x}})=P(r({\mathbf{X}}_{1})\leq r({\bm{x}}))$. Note that $0<p_{1}({\bm{x}})=P({\mathbf{X}}_{1}={\bm{x}})\leq P(r({\mathbf{X}}_{1})\leq r
({\bm{x}}))=P_{{\mathbf{X}}}({\bm{x}})$, we have $0<\frac{p_{1}({\bm{x}})}{P_{{\mathbf{X}}}({\bm{x}})}\leq 1$, and thus $\lim\limits_{N\rightarrow+\infty}\left[1-\frac{p_{1}({\bm{x}})}{P_{{\mathbf{X}
}}({\bm{x}})}\right]^{N}=0$. Similarly, by setting $P_{{\mathbf{Y}}}({\bm{x}})=P(r({\mathbf{Y}}_{1})\leq r({\bm{x}}))$, we have

$$ $P\left(\arg\max_{1\leq m\leq M}r({\mathbf{Y}}_{m})={\bm{x}}\right)=P_{{\mathbf {Y}}}({\bm{x}})^{M}\left[1-\left[1-\frac{p_{2}({\bm{x}})}{P_{{\mathbf{Y}}}({ \bm{x}})}\right]^{M}\right],$ (10) $$

and $\lim\limits_{M\rightarrow+\infty}\left[1-\frac{p_{2}({\bm{x}})}{P_{{\mathbf{Y}
}}({\bm{x}})}\right]^{M}=0$.
Let ${\bm{x}}^{*}=\arg\max\limits_{{\bm{x}}\in S}r({\bm{x}})$. Therefore, $P_{{\mathbf{X}}}({\bm{x}}^{*})=P_{{\mathbf{Y}}}({\bm{x}}^{*})=1$ and $P_{{\mathbf{X}}}({\bm{x}})<1,P_{{\mathbf{Y}}}({\bm{x}})<1$ for ${\bm{x}}\neq{\bm{x}}^{*}$.
When $N,M\rightarrow+\infty$, we have $P_{{\mathbf{X}}}({\bm{x}}^{*})^{N}=P_{{\mathbf{Y}}}({\bm{x}}^{*})^{M}=1$, and $P_{{\mathbf{X}}}({\bm{x}}^{*})^{N},P_{{\mathbf{Y}}}({\bm{x}}^{*})^{M}\rightarrow
0$ for ${\bm{x}}\neq{\bm{x}}^{*}$. Therefore, we have

$$ $\displaystyle\lim_{N,M\rightarrow+\infty}\sum_{{\bm{x}}\in S}P\left(\arg\max_{ 1\leq n\leq N}r({\mathbf{X}}_{n})={\bm{x}}\right)P\left(\arg\max_{1\leq m\leq M }r({\mathbf{Y}}_{m})={\bm{x}}\right)$ (11) $\displaystyle=$ $\displaystyle\sum_{{\bm{x}}\in S}\lim_{N,M\rightarrow+\infty}P_{{\mathbf{X}}}( {\bm{x}})^{N}P_{{\mathbf{Y}}}({\bm{x}})^{M}\lim_{N,M\rightarrow+\infty}\left[1 -\left[1-\frac{p_{1}({\bm{x}})}{P_{{\mathbf{X}}}({\bm{x}})}\right]^{N}\right] \left[1-\left[1-\frac{p_{2}({\bm{x}})}{P_{{\mathbf{Y}}}({\bm{x}})}\right]^{M}\right]$ $\displaystyle=$ $\displaystyle\lim_{N,M\rightarrow+\infty}P_{{\mathbf{X}}}({\bm{x}}^{*})^{N}P_{ {\mathbf{Y}}}({\bm{x}}^{*})^{M}+\sum_{{\bm{x}}\in S,{\bm{x}}\neq{\bm{x}}^{*}} \lim_{N,M\rightarrow+\infty}P_{{\mathbf{X}}}({\bm{x}})^{N}P_{{\mathbf{Y}}}({ \bm{x}})^{M}$ $\displaystyle=$ $\displaystyle 1,$ $$

which proves Eq. (6). For $K>1$, the equality of two top-$K$ subsets requires the elements ranked from 1 to $K$ to be equal, respectively. Therefore, the left hand of Eq. (3) can be decomposed using Bayes’s Law. Let ${\mathcal{X}}=\left\{{\mathbf{X}}_{1},{\mathbf{X}}_{2},\cdots,{\mathbf{X}}_{N}\right\}$, ${\mathcal{Y}}=\left\{{\mathbf{Y}}_{1},{\mathbf{Y}}_{2},\cdots,{\mathbf{Y}}_{M}\right\}$, ${\mathbf{X}}^{*}=\arg\max\limits_{1\leq n\leq N}r({\mathbf{X}}_{n})$, and ${\mathbf{Y}}^{*}=\arg\max\limits_{1\leq m\leq M}r({\mathbf{Y}}_{m})$, we have $P({\mathbf{X}}^{*}={\mathbf{Y}}^{*})\rightarrow 1$ (Eq. (6)) and

$$ $\displaystyle P\left(\text{top-}2\left\{\left.r({\mathbf{X}}_{n})\right|1\leq n \leq N\right\}=\text{top-}2\left\{\left.r({\mathbf{Y}}_{m})\right|1\leq m\leq M \right\}\right)$ (12) $\displaystyle=$ $\displaystyle P\left(\left.\arg\max_{{\mathbf{X}}\in{\mathcal{X}}-\{{\mathbf{X }}^{*}\}}r({\mathbf{X}})=\arg\max_{{\mathbf{Y}}\in{\mathcal{Y}}-\{{\mathbf{Y}} ^{*}\}}r({\mathbf{Y}})\right|{\mathbf{X}}^{*}={\mathbf{Y}}^{*}\right)P({ \mathbf{X}}^{*}={\mathbf{Y}}^{*}).$ $$

Since ${\mathbf{X}}_{1},{\mathbf{X}}_{2},\cdots,{\mathbf{X}}_{N}$ and ${\mathbf{Y}}_{1},{\mathbf{Y}}_{2},\cdots,{\mathbf{Y}}_{M}$ are i.i.d. random variables, we can still decompose the term $P\left(\left.\arg\max\limits_{{\mathbf{X}}\in{\mathcal{X}}-\{{\mathbf{X}}^{*}
\}}r({\mathbf{X}})={\bm{x}}\right|{\mathbf{X}}^{*}={\mathbf{Y}}^{*}\right)$ as Eq. (9), which means when $N,M\rightarrow+\infty$:

$$ $P\left(\left.\arg\max_{{\mathbf{X}}\in{\mathcal{X}}-\{{\mathbf{X}}^{*}\}}r({ \mathbf{X}})=\arg\max_{{\mathbf{Y}}\in{\mathcal{Y}}-\{{\mathbf{Y}}^{*}\}}r({ \mathbf{Y}})\right|{\mathbf{X}}^{*}={\mathbf{Y}}^{*}\right)\rightarrow 1.$ (13) $$

Eq. (13) means the elements with the secondary large $r({\bm{x}})$ values are highly likely to be the same. The decomposition in Eq. (12) can be conducted $K$ times for the elements ranked from 1 to $K$, and each decomposed term approaches 1 similar to Eq. (13). So far, we have proved that the probability of ${\mathcal{X}}$ and ${\mathcal{Y}}$ having the same top-$K$ subsets measured by $r({\bm{x}})$ approaches to 1 when $N,M\rightarrow+\infty$, which is formally written as

$$ $P\left(\text{top-}K\left\{\left.r({\mathbf{X}}_{n})\right|1\leq n\leq N\right \}=\text{top-}K\left\{\left.r({\mathbf{Y}}_{m})\right|1\leq m\leq M\right\} \right)\rightarrow 1.$ (14) $$

This completes the proof of Proposition 3.

## Appendix B More Experimental Details

#### Model and Training Configurations.

We mostly follow Brown et al. [] to set the model and learning rate configurations, as summarized in Table 6. The 500M, 1.8B, and 4B teacher models are the officially released Qwen-1.5 checkpoints (^3^33[https://huggingface.co/Qwen](https://huggingface.co/Qwen)) and the 300M teacher model used in Section 3.5 to analysis the effect of teacher LM sizes Is pre-trained on 200B tokens from our pre-training corpus. We train all the LMs with the AdamW [] optimizer, with $\beta_{1}=0.9$, $\beta_{2}=0.98$, and a 0.1 weight decay. We set the batch size to 512 and the max sequence length to 1,024, corresponding to 100K total training steps for roughly 50B tokens in Pre-Train w/o KD, SeqKD, and MiniPLM. For MiniLLM and Vanilla KD, we limit the training steps to align their training computation with Pre-Train w/o KD. Specifically, we assume the computation of a forward and backward pass are $2ND$ and $4ND$, respectively, where $N$ is the model size and $D$ is the number of trained tokens. The training steps of MiniLLM and Vanilla KD are listed in Table 7. We linearly warm up the learning rate for 2K steps and apply cosine learning rate decay until 1/10 of the max values. All experiments are conducted on NVIDIA 40G A100 and NVIDIA 32G V100 GPUs.

**Table 6: Model configurations and corresponding learning rates.**
| Model Size | $d_{\text{model}}$ | $d_{\text{FFN}}$ | $n_{\text{layers}}$ | $n_{\text{head}}$ | $d_{\text{head}}$ | learning rate |
| --- | --- | --- | --- | --- | --- | --- |
| 104M | 512 | 1,408 | 8 | 8 | 64 | $6\times 10^{-4}$ |
| 200M | 768 | 2,112 | 12 | 12 | 64 | $6\times 10^{-4}$ |
| 300M | 768 | 2,112 | 18 | 12 | 64 | $6\times 10^{-4}$ |
| 500M | 1,024 | 2,816 | 24 | 16 | 64 | $3\times 10^{-4}$ |
| 1.2B | 1,536 | 4,224 | 24 | 16 | 96 | $2.5\times 10^{-4}$ |

**Table 7: Training steps in Vanilla KD and MiniLLM, which is set to ensure training-time computation to be the same as Pre-Train w/o KD. $N_{\text{stu}}$ and $N_{\text{tch}}$ are model sizes of student and teacher LMs respectively. $N_{\text{tch}}=1.8B$ in our experiments. $T=100K$ is the training steps in Pre-Train w/o KD.**
|  | Vanilla KD | MiniLLM |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| Formula | $\frac{3N_{\text{stu}}}{3N_{\text{stu}}+N_{\text{tch}}}T$ | $\frac{3N_{\text{stu}}}{4N_{\text{stu}}+2N_{\text{tch}}}T$ |  |  |  |  |
| Student Model Size $N_{\text{stu}}$ | 200M | 500M | 1.2B | 200M | 500M | 1.2B |
| Training Steps | 25K | 45K | 65K | 15K | 30K | 40K |

#### Evaluation Details.

Our downstream datasets for evaluation include Hellaswag (HS; ), LAMBADA (LAM; ), Winograde (Wino; ), OpenbookQA (OBQA; ), ARC-Easy/Challange (ARC-e/c; ), PIQA [], SIQA [], and StoryCloze (Story; ). We apply the LM-Eval-Harness [] (^4^44) framework to conduct zero-shot evaluation. We sample 10K documents from the DCLM [] corpus to construct our test set for language modeling evaluation.

## Appendix C More Results

### C.1 Computation Scaling Curves of More Sizes

In Figure 7, we plot the scaling curves of average accuracy on the downstream tasks with respect to the pre-training FLOPs for student LMs with 200M and 1.2B parameters. We can see that MiniPLM saves pre-training computation for both student LM sizes and constantly outperforms Vanilla KD.

Figure: (a) 1.8B $\rightarrow$ 200M
Refer to caption: x11.png

### C.2 Test Losses Extrapolation with Scaling Laws

**Table 8: Scaling Law constants in Eq. (16) fitted using the loss curves in Figure 4. $N_{\text{stu}}$ means the student LM size. $C_{\text{1T}}$ and $C_{\text{10T}}$ are the compute spent on processing 1T and 10T tokens in Pre-Train w/o KD and MiniPLM, which Vanilla KD aligns with.**
| $N_{\text{stu}}$ | Method | $A_{c}$ | $\alpha_{c}$ | $L_{\infty}$ | $C_{\text{1T}}$ (FLOPs) | $C_{\text{10T}}$ (FLOPs) |
| --- | --- | --- | --- | --- | --- | --- |
| 200M | Pre-Train w/o KD | 2.19$\times 10^{7}$ | 0.41 | 3.30 | 1.26$\times 10^{21}$ | 1.26$\times 10^{22}$ |
| Vanilla KD | 9.77$\times 10^{7}$ | 0.44 | 3.34 |  |  |  |
| MiniPLM | 8.56$\times 10^{10}$ | 0.59 | 3.25 |  |  |  |
| 500M | Pre-Train w/o KD | 2.73$\times 10^{8}$ | 0.45 | 3.06 | 3.14$\times 10^{21}$ | 3.14$\times 10^{21}$ |
| Vanilla KD | 3.14$\times 10^{8}$ | 0.45 | 3.05 |  |  |  |
| MiniPLM | 6.64$\times 10^{9}$ | 0.52 | 3.03 |  |  |  |
| 1.2B | Pre-Train w/o KD | 1.88$\times 10^{8}$ | 0.43 | 2.91 | 7.30$\times 10^{21}$ | 7.30$\times 10^{21}$ |
| Vanilla KD | 1.10$\times 10^{10}$ | 0.52 | 2.90 |  |  |  |
| MiniPLM | 4.29$\times 10^{8}$ | 0.45 | 2.86 |  |  |  |

#### Data-Unlimited Setting.

In Table 2, we follow Hoffmann et al. [] to fit the scaling law curves with the test losses on the DCLM corpus. Then, we used the fitted constants to predict the test losses for 1T and 10T training data in Pre-Train w/o KD and MiniPLM, and that for Vanilla KD when it consumes the same training FLOPs. Specifically, we consider the following power law suggested by Hoffmann et al. []:

$$ $L(N_{m},N_{d})=L_{\text{irr}}+\frac{A_{m}}{N_{m}^{\alpha_{m}}}+\frac{A_{d}}{N_ {d}^{\alpha_{d}}}$ (15) $$

where $N_{m}$ is the model sizes, $N_{d}$ is the number of trained tokens, and $L_{\text{irr}}$ is the irreducible test loss. Since the required computation $C$ of training on a token of Vanilla KD and other methods are different, we re-write Eq. (15) using the fact that $C\propto N_{m}N_{d}$ for a fixed model size:

$$ $L(C)=L_{\infty}+\frac{A_{c}}{C^{\alpha_{c}}},$ (16) $$

where $L_{\infty}$ and $A_{c}$ depends on the model size $N_{m}$, and $\alpha_{c}=\alpha_{d}$. In this way, we can fit the loss curves in Figure 4 with Eq. (16). Table 8 includes the fitted values of $L_{\infty}$, $A_{c}$, and $\alpha_{c}$. We also include the total FLOPs of Pre-Train w/o KD and MiniPLM on 1T ($C_{\text{1T}}$) and 10T ($C_{\text{10T}}$) tokens, which Vanilla KD aligns with. The numbers in Table 8 can be computed with Eq. (16) and the constants in Table 8.

#### Data-Limited Setting.

In Section 3.4, we extrapolate the loss curve of Pre-Train w/o KD with the Data-Constrained Scaling Law [], which is almost the same as Eq. (15), except that $N_{d}$ represents the total number of tokens in the pre-training corpus repeated 4 times, as suggested by Muennighoff et al. []. Therefore, after $N_{d}$ is solved by letting the loss of Pre-Train w/o KD equal that of MiniPLM ($\alpha=0.25$, 4 Eps.), we divide the value by $4$, resulting in a training tokens number of 118B. This means 68B extra training tokens are needed for Pre-Train w/o KD to achieve the performance of MiniPLM using 50B training tokens.

### C.3 Difference Sampling with a Proxy Model

#### Motivation.

In this section, we show that the offline computational overhead of Difference Sampling can be further reduced by conducting teacher and reference LM inference on a small proxy dataset and then transferring the value $r(p,p_{\text{ref}},{\bm{x}})$ to the entire corpus $\mathcal{D}$ with a proxy model.
Specifically, we first uniformly sample a proxy subset $\mathcal{D}_{\text{prx}}$ from $\mathcal{D}$, satisfying $|\mathcal{D}_{\text{prx}}|\ll|\mathcal{D}|$. Then, we compute the $r(p,p_{\text{ref}},{\bm{x}})$ value for each instance ${\bm{x}}$ in $\mathcal{D}_{\text{prx}}$.
Note that since $|\mathcal{D}_{\text{prx}}|\ll|\mathcal{D}|$, the computational overhead of teacher LM inference is significantly reduced. After that, we fine-tune a small proxy model on $\mathcal{D}_{\text{prx}}$ to fit the $r(p,p_{\text{ref}},{\bm{x}})$ values, which is used to infer the values for instances from $\mathcal{D}$. Finally, the Top-$K$ operation in Eq. (4) is based on the inferred values.

#### Method.

To test this approach, we uniformly sample a $\mathcal{D}_{\text{prx}}$ containing 0.1B tokens from the 50B-token $\mathcal{D}$. Computing $\log\frac{p({\bm{x}})}{p_{\text{ref}}({\bm{x}})}$ values on $\mathcal{D}_{\text{prx}}$ takes only 0.2% computation of that on $\mathcal{D}$. Then, we employ the reference LM as the proxy model and fine-tune it with the following regression loss:

$$ ${\bm{w}}^{*},b^{*},{\bm{\theta}}^{*}_{\text{ref}}=\arg\min_{{\bm{w}},b,{\bm{ \theta}}_{\text{ref}}}\frac{1}{|\mathcal{D}_{\text{prx}}|}\sum_{{\bm{x}}\in \mathcal{D}_{\text{prx}}}\left[{\bm{w}}^{\top}\overline{{\bm{h}}}({\bm{x}},{ \bm{\theta}}_{\text{ref}})+b-r(p,p_{\text{ref}},{\bm{x}})\right]^{2},$ (17) $$

**Table 9: Offline FLOPs and average accuracy (Acc.) on downstream tasks of MiniPLM using a proxy model, compared with that of standard MiniPLM and Vanilla KD.**
| Method | FLOPs | Acc. |
| --- | --- | --- |
| Vanilla KD | Online | 39.9 |
| MiniPLM | $2\times 10^{20}$ | 41.3 |
| $\textsc{MiniPLM}_{\text{prx}}$ | $9\times 10^{18}$ | 40.9 |

where $\overline{{\bm{h}}}({\bm{x}},{\bm{\theta}}_{\text{ref}})\in R^{d}$ is the average output hidden states of the reference LM, with $d$ representing the model’s hidden size and ${\bm{\theta}}_{\text{ref}}$ representing the parameters of the reference LM. ${\bm{w}}\in\mathbb{R}^{d},b\in\mathbb{R}$ are the parameters of a linear head outputting a scalar as the predicted $r(p,p_{\text{ref}},{\bm{x}})$ values, given the average hidden states. The inferred values on $\mathcal{D}$ are given by $\hat{r}({\bm{x}})={{\bm{w}}^{*}}^{\top}\overline{{\bm{h}}}({\bm{x}},{\bm{
\theta}}^{*}_{\text{ref}})+b^{*}$. Since this inference process is based on the reference LM, it still saves computation compared to inference with the teacher LM. The difference-sampled corpus is then obtained by selecting $K=\alpha|\mathcal{D}|$ instances from $\mathcal{D}$ with the highest $\hat{r}({\bm{x}})$ values.

#### Results.

We term this method as $\textsc{MiniPLM}_{\text{prx}}$ and compare its performance with Vanilla KD and MiniPLM in Table 9. We can see that $\textsc{MiniPLM}_{\text{prx}}$ requires much less offline inference computation compared to standard MiniPLM while still maintaining substantial improvement over Vanilla KD.

### C.4 Ablation Study

#### Impact of the Reference Model.

Figure: Figure 8: Impact of the reference model size. We use the 1.8B LM as the teacher and the 200M LM as the student. We report the average zero-shot accuracy on the downstream tasks of the LMs trained with MiniPLM and compare it with Vanilla KD.
Refer to caption: x14.png

In Figure 9, we show the impact of the reference model size on the performance of LMs trained with MiniPLM. We can see that larger reference models lead to better performance of MiniPLM, but the improvement saturates as the reference LM grows larger.

#### Difference Sampling Ratio.

In Figure 9, we plot the impact of the difference sampling ratio on the student LM’s performance in the data-unlimited setting. We can see that small sampling ratios result in better zero-shot accuracy on downstream tasks. However, to ensure that the difference-sampled data contains 50B tokens, we need a larger original corpus $\mathcal{D}$. To control the computation overhead, we mainly use $\alpha=0.5$ in our experiments.

### C.5 Case Study

We present a case study in Table 10, 11, and 12 to show the instances corresponding to the three parts of data distribution shown in Figure 3(b).

**Table 10: Easy and common instances down-sampled by Difference Sampling. Instance #1 and #2: HTML and LaTeX code data that contain repeated patterns, which is easy to fit by both the reference and teacher LM. Instance #3: Dialogues from a story, which is relative easy for LMs but are still selected by Difference Sampling to help student LMs learn basic language skills.**
| $p({\bm{x}})\gtrsim p_{\text{ref}}({\bm{x}})$: Easy and common instances |  |  |  |  |
| --- | --- | --- | --- | --- |
| Instance #1 | $-\log p({\bm{x}})=1.24$ | $-\log p_{\text{ref}}({\bm{x}})=1.28$ | $\log\frac{p({\bm{x}})}{p_{\text{ref}}({\bm{x}})}=0.04$ | Discarded |
|  | <node id=’-659’ action=’modify’ visible=’true’ lat=<br>’0.05467069248579575’ lon=’0.014892969670168166’ /><br><node id=’-657’ action=’modify’ visible=’true’ lat=<br>’0.05243834625645345’ lon=’0.023237696125627347’ /><br><node id=’-655’ action=’modify’ visible=’true’ lat=<br>’0.04940873338995851’ lon=’0.03152927145717611’ /><br><node id=’-653’ action=’modify’ visible=’true’ lat=<br>’0.04786735135170292’ lon=’0.040405509151806455’ /><br><node id=’-651’ action=’modify’ visible=’true’ lat=<br>’0.04733584029593642’ lon=’0.04513595918070425’ /> |  |  |  |
| Instance #2 | $-\log p({\bm{x}})=0.44$ | $-\log p_{\text{ref}}({\bm{x}})=0.51$ | $\log\frac{p({\bm{x}})}{p_{\text{ref}}({\bm{x}})}=0.07$ | Discarded |
|  | {\sum\limits_{j = 1}{̂n - 1}\operatorname{size}\left(<br>{\mathcal{X},j} \right) \cdot \operatorname{size}<br>\left( {\mathcal{Z},n - j} \right)\quad} &<br>{\mathcal{I} = \mathcal{M}_{1},} \\<br>{\sum\limits_{j = 1}{̂n - 1}\operatorname{size}\left(<br>{\mathcal{X},j} \right) \cdot \operatorname{size}<br>\left( {\mathcal{X},n - j} \right)\quad} &<br>{\mathcal{I} = \mathcal{M}_{2},} \\<br>{\sum\limits_{j = 1}{̂n - 1}\operatorname{size}\left(<br>{\mathcal{X},j} \right) \cdot \operatorname{size}<br>\left( {\mathcal{A},n - j} \right)\quad} &<br>{\mathcal{I} = \mathcal{M}_{3},} \\<br>{\sum\limits_{j = 1}{̂n - 1}\operatorname{size}\left(<br>{\mathcal{A},j} \right) \cdot \operatorname{size}<br>\left( {\mathcal{Z},n - j} \right)\quad} &<br>{\mathcal{I} = \mathcal{M}_{4},} \\ |  |  |  |
| Instance #3 | $-\log p({\bm{x}})=2.83$ | $-\log p_{\text{ref}}({\bm{x}})=3.86$ | $\log\frac{p({\bm{x}})}{p_{\text{ref}}({\bm{x}})}=1.02$ | Selected |
|  | "I felt bad for Coach Rod to have to deal with it,"<br>said senior tight end Mike Massey, a St. Ignatius<br>grad. "But for the players it was a non-issue. We<br>didn’t even talk about it."<br>But some other people in college football did.<br>"[Boren] was a legacy guy and when he makes comments<br>like [that], that’s like getting kicked square in the<br>shorts when you’re Rich Rodriguez," ESPN analyst and<br>former OSU quarterback Kirk Herbstreit said Thursday.<br>"With all the other things that are happening, when<br>you get that comment from a legacy guy, a guy whose<br>dad started three or four years for Bo, all of a<br>sudden everybody’s tentacles go up a little bit. |  |  |  |

**Table 11: Hard and valuable instances up-sampled by Difference Sampling. Instance #1: High-quality long documents containing versatile world knowledge. Instance #2: instance that contains translation tasks, presented in an in-context learning form. Instance #3: High-quality code data with detailed comments.**
| $p({\bm{x}})\gg p_{\text{ref}}({\bm{x}})$: Hard and valuable instances |  |  |  |  |
| --- | --- | --- | --- | --- |
| Instance #1 | $-\log p({\bm{x}})=1.26$ | $-\log p_{\text{ref}}({\bm{x}})=4.20$ | $\log\frac{p({\bm{x}})}{p_{\text{ref}}({\bm{x}})}=2.94$ | Selected |
|  | Legal along with Environmental Responsibility! Dumpster rentals in the user side may seem as fundamental as placing a phone, having a dumpster sent and hurling all your disposals inside to be carted away. Nonetheless, there are legal issues attached to appropriate disposal connected with certain products which tie up into environmental issues. The 10 Yard Dumpster For Rent in Pocahontas customer or perhaps demolition purchaser should be informed about these issues by means of careful screening so as to reduce a firm’s liability which inturn keeps a firm’s overhead all the way down and makes for prompt fall off, pick up along with disposal of the dumpster and it’s articles. |  |  |  |
| Instance #2 | $-\log p({\bm{x}})=2.36$ | $-\log p_{\text{ref}}({\bm{x}})=5.59$ | $\log\frac{p({\bm{x}})}{p_{\text{ref}}({\bm{x}})}=3.23$ | Selected |
|  | {CJK}UTF8gbsn<br>有利 you3li4 yǒulì advantageous; beneficial<br>谨慎 jin3shen4 jǐnshèn cautious; prudent<br>甲 jia3 jiǎ one; armor (1st Heavenly Stem)<br>犹豫 you2yu4 yóuyù hesitate; hesitant; undecided<br>从此 cong2ci3 cóngcǐ from now on; since then<br>企业 qi3ye4 qǐyè company; business; firm<br>下载 xia4zai3 xiàzǎi to download<br>狮子 shi1zi5 shīzi lion<br>青少年 qing1shao4nian2 qīngshàonián teenager |  |  |  |
| Instance #3 | $-\log p({\bm{x}})=0.16$ | $-\log p_{\text{ref}}({\bm{x}})=2.73$ | $\log\frac{p({\bm{x}})}{p_{\text{ref}}({\bm{x}})}=2.56$ | Selected |
|  | function WritableState(options, stream) {<br>var Duplex = require(’./_stream_duplex’);<br>options = options || {};<br>// the point at which write() starts returning false<br>// Note: 0 is a valid value, means that we always<br>return false if<br>// the entire buffer is not flushed immediately<br>on write()<br>var hwm = options.highWaterMark;<br>var defaultHwm = options.objectMode?16:16*1024;<br>this.highWaterMark = (hwm || hwm === 0) ? hwm :<br>defaultHwm;<br>// object stream flag to indicate whether or not<br>this stream<br>// contains buffers or objects.<br>this.objectMode = !!options.objectMode;<br>...<br>} |  |  |  |

**Table 12: Noisy and harmful instances discarded by Difference Sampling. Instance #1: irregular symbols. Both the reference LM and the teacher LM are hard to fit. Instance #2: a string primarily made of meaningless numbers. Both the reference LM and the teacher LM are easy to fit, by predicting numbers. Instance #3: data collection in scientific paper, but lack of contexts. The reference LM fits the patterns, which the teacher LM finds useless.**
| $p({\bm{x}})\lesssim p_{\text{ref}}({\bm{x}})$: Noisy and harmful instances |  |  |  |  |
| --- | --- | --- | --- | --- |
| Instance #1 | $-\log p({\bm{x}})=9.50$ | $-\log p_{\text{ref}}({\bm{x}})=6.60$ | $\log\frac{p({\bm{x}})}{p_{\text{ref}}({\bm{x}})}=-2.90$ | Discarded |
|  | ]{}<br>*********************************************<br>. ...<br>... . |  |  |  |
| Instance #2 | $-\log p({\bm{x}})=1.01$ | $-\log p_{\text{ref}}({\bm{x}})=0.90$ | $\log\frac{p({\bm{x}})}{p_{\text{ref}}({\bm{x}})}=-0.11$ | Discarded |
|  | (91.00,60.00) (1.00,35.00) (2.00,35.00)[(1,0)[13.00]{}]<br>{} (16.00,35.00) (16.70,34.30)[(1,-1)[8.60]{}]{}<br>(26.00,25.00) (26.00,11.00) (26.00,12.00)[(0,1)[12.00]<br>{}]{} (36.00,35.00) (26.00,45.00) (16.70,35.70)<br>(35.30,34.30)[(-1,-1)[8.60]{}]{} (56.00,35.00)<br>(66.00,25.00) (76.00,35.00) (66.00,45.00) (66.00,46.00) |  |  |  |
| Instance #3 | $-\log p({\bm{x}})=2.53$ | $-\log p_{\text{ref}}({\bm{x}})=0.26$ | $\log\frac{p({\bm{x}})}{p_{\text{ref}}({\bm{x}})}=-2.26$ | Discarded |
|  | *Z* = 8<br>------------------------- ---------------------------------------<br>Data collection #tablewrapdatacollectionlong<br>===============<br>---------------------------------------------------------------<br>Bruker APEXII area-detector diffractometer<br>Radiation source: fine-focus sealed tube<br>graphite<br>p and w scans<br>Absorption correction: multi-scan<br>(*SADABS*; Sheldrick, 2004)<br>T*m̃in=̃ 0.970, *T*m̃ax=̃ 0.981<br>20408 measured reflections<br>--------------------------------------------------------------- |  |  |  |

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: The explanation of Difference Sampling and its theoretical justification could be simplified or supplemented with more intuitive examples for clarity.

The paper should discuss in detail the limitations of the proposed method.

The transferability and generalizability of MINIPLM across different language models and types of downstream tasks should be discussed.
soundness: 3
strengths: The introduction of MINIPLM, which uses offline teacher inference and a unique Difference Sampling method, provides an innovative solution to enhance KD during pre-training. This method effectively balances computational costs and improves student language model performance.

The paper presents extensive experimental results across a variety of downstream tasks and model sizes (200M, 500M, 1.2B parameters), demonstrating the effectiveness of MINIPLM. Comparisons with various baselines show substantial improvements.

The results on computation reduction and performance scalability (e.g., 2.2x computation reduction) validate the practicality of MINIPLM for large-scale pre-training.
summary: The paper studies the field of language model pre-training by proposing an efficient KD framework. MINIPLM’s use of offline inference and Difference Sampling stands out as a method that balances computational efficiency and model performance improvement. The findings are well-supported by robust experiments, although the explanations and practical considerations could be more user-friendly and extensive.
weaknesses: The explanation of Difference Sampling and its theoretical justification could be simplified or supplemented with more intuitive examples for clarity.

The requirement for teacher LM’s probabilities on pre-training corpus data may pose challenges for closed-source models or environments with limited computational resources. This could be highlighted more explicitly as a practical limitation.

The paper would benefit from a more detailed discussion on the transferability and generalizability of MINIPLM across different language models and types of downstream tasks, beyond the ones tested.

While a brief limitations section is included, expanding on potential drawbacks, such as cases where MINIPLM may underperform compared to online KD or specific model configurations where it might not be ideal, would strengthen the discussion.

---

confidence: 3
contribution: 3
initial_rating: 8
presentation: 4
questions: Please refer to weakness.
soundness: 4
strengths: - This paper provides decent theoretical foundations 
- The contribution and experiments designed for this paper is very clear. 
- Evaluations demonstrate improved performance and data efficiency gained from the method. 
- The contribution of this paper is important to scenarios when data is insufficient.
summary: This paper introduces an approach to knowledge distillation (KD) for pretraining LLMs. The proposed method, MINILLM, reformulates the objective of KD with reverse KL divergence as a reward maximization problem, enhancing the student model's focus on the teacher model's primary knowledge modes through difference sampling. Evaluations further show the efficiency and effectiveness of this approach.
weaknesses: Overall, I don't have many criticism toward this paper. 
- There is not enough emphasis on how much of this approach would require computational resources to perform distillation. A discussion on resource consumption, such as training time or memory usage, would help assess its practicality for real-world deployment.
- The complexity of the training process may impose challenges in reproducibility and implementation.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: - While the cross-model family experiments with Llama 3.1 and Mamba are interesting, they lack details about model configurations, training data, and experimental setup. Could the authors provide them?
soundness: 3
strengths: - This paper Introduces "Difference Sampling" as a novel offline KD method that fundamentally changes how knowledge is transferred.
- This paper re-computes teacher model probabilities just once, requiring only 200MB storage for 50B tokens vs 30PB for traditional methods.
- This paper shows clear computational efficiency gains (2.2x reduction for equivalent performance).
- This paper provides solid theoretical justification through Proposition 2.1 and detailed analysis of the sampling approach
summary: MiniPLM introduces a novel approach to knowledge distillation (KD) for pre-training smaller language models that addresses key challenges in efficiency, flexibility, and effectiveness. The core innovation is "Difference Sampling" - a method that intelligently refines the pre-training data distribution by comparing probabilities between a large teacher model and a small reference model. This enables efficient offline KD by selecting training instances where the teacher model shows strong confidence but the reference model struggles, indicating valuable knowledge that smaller models need to learn. Notably, MiniPLM achieves 2.2x computational savings while maintaining performance comparable to traditional KD approaches.
weaknesses: - First, the absence of teacher model (1.8B) performance metrics in Table 1 makes it difficult to gauge the effectiveness of knowledge transfer - readers cannot assess how much of the teacher's capabilities are preserved in the student models. 
- The paper's experimental scope is rather constrained. By focusing primarily on a 1.8B teacher model, it misses the opportunity to demonstrate effectiveness with larger, more modern models. The experiments would be more convincing if they demonstrated success with larger model gaps, such as distilling from Llama 3.1 8B to 1B on Table 1 setting.
-  While the paper introduces "Difference Sampling" as a key innovation, it doesn't thoroughly demonstrate its effectiveness through controlled experiments. The comparison with SeqKD might be not enough because SeqKD's lower performance might be attributed to the authors' implementation choice of using only 768 tokens for approximation.
- Furthermore, the omission of comparisons with recent work in small language models, particularly MobileLLM and its associated baselines in the 200M/500M parameter range, leaves an important gap in understanding how MiniPLM compares to state-of-the-art approaches in efficient model training. 

[MobileLLM] MobileLLM: Optimizing Sub-billion Parameter Language Models for On-Device Use Cases

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 4
questions: - Can you provide theoretical justification for using log(p(x)/p_ref(x)) as the sampling metric?
- How does the method's effectiveness scale with model size differences?
- How would the method perform with larger teacher models (>10B)?
- What modifications are needed for very large scale deployment?
soundness: 3
strengths: - Well-motivated solution approach
- Elegant Difference Sampling design
- Efficient offline computation strategy
- Comprehensive experiments across model scales
summary: This paper introduces MINIPLM, a knowledge distillation framework for pre-training language models. The key innovation is Difference Sampling, which refines the training data distribution by comparing a large teacher model with a small reference model. The framework: (1) performs offline teacher inference to avoid online KD computation overhead, (2) maintains data difficulty through reference model comparison, and (3) enables cross-family model distillation. Experiments on models from 200M to 1.2B parameters demonstrate improvements in downstream performance, language modeling capability, and computational efficiency.
weaknesses: - Limited analysis of finite-sample scenarios
- Teacher model size limited to 1.8B
- Limited cross-family experiments
- Limited exploration of alternative sampling strategies
- Reference model selection criteria not fully justified

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: See weaknesses.
soundness: 3
strengths: 1. The paper uses knowledge distillation to obtain a better subset of samples for training a small PLM which is different from existing methods.

2. The paper is well written and easy to follow.

3. The paper has good experimental results as compared to the existing methods.

4. The problem taken up is an important problem in pre-training the language models.
summary: The paper proposes a method to train a small PLM using knowledge distillation from large PLM and in an offline manner instead of the existing online methods. The methods first trains a small LM from scratch on small subset of the training corpus and then a large PLM is used to perform difference sampling that samples from the large corpus the difficult samples reducing the chances of learning the patterns only from the easy portion of the dataset.
weaknesses: 1. The paper proposes to choose a subset of training dataset that can provide better performance of the small model using knowledge distillation but the title of the paper is ``Knowledge distillation for pre-training language models", here knowledge distillation is not being used for pretraining instead it is being used for a better subset selection from a huge corpus of data that can in turn have better pre-training of the PLMs, this makes it quite misleading.

2. In line 238 paper says that when the log term is less than zero, those instances are harmful, however, this is not true as explained in papers such as [1] where an easy sample can be predicted correctly by a smaller model while it can get wrong on a larger model due to extraction of overly distracted features. This phenomenon is termed as `overthinking' of PLMs.

3. I am not very clear how $p_{ref}$ can mimic the behavior of $q_{\theta}$, even these models can have certain differences and how MINIplm performs better than MINIllm even when MINIllm directly uses the $q_{\theta}$ distribution?

4. The method is computationally very heavy as we need three LLM model: one large model from which distillation is being performed, a smaller version trained from scratch and then training the distilled model using the existing two, which makes it complex.

Meta-review: The paper introduces MINIPLM, a method for simplifying large language models. MINIPLM uses ``Difference Sampling'' to find difficult training examples where a large teacher model and a smaller reference model give different answers. When trained under the same conditions, MINIPLM performs better than other methods and works well with different model sizes.

Reviewer ZWq4 and Reviewer Yqhf find the method and experimental results clear, highlighting the effectiveness of MINIPLM. Reviewer 25k6 notes the computational efficiency gains as a key advantage, while Reviewer rC3o commends the paper for being easy to follow. 
Reviewers 25K6 and Rr1c raised concerns about the practicality. They questioned the efficiency of the method and the criteria for choosing the small reference model. Reviewer 25K6 also questioned the offline inference process and how it affects resource usage. Additionally, Rr1c asked for more proof that the reference model helps improve performance.

During the rebuttal, the authors made great efforts to address these issues. The conducted additional experiments, such as an ablation study on different reference models, teacher model size, and compared MINIPLM with other open-source small language models.

After the rebuttal, all reviewers reached a consensus to accept the paper. The AC agree with the reviewers and recommended accepting the paper. Nonetheless, AC suggests that the authors include additional experiments and discussions from the rebuttal in the final version of the paper.

Decision: Accept (Poster)