MiniPLM: Knowledge Distillation for Pre-training Language Models
Authors: Yuxian Gu
Hao Zhou
Fandong Meng
Jie Zhou
Minlie Huang
Venue: ICLR 2025 Poster (2025)
Keywords: Language Models
Knowledge Distillation
Pre-Training
Primary area: foundation or frontier models, including LLMs
TL;DR: This work introduces an efficient, flexible, and effective knowledge distillation framework for pre-training language models.

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 efficiency, flexibility, and effectiveness issues. 
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.
In this work, we propose **MiniPLM**, a KD framework for pre-training LMs by refining the training data distribution with the teacher LM's knowledge.
For efficiency, MiniPLM performs offline teacher inference, allowing KD for multiple student LMs without adding training 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 training data difficulty and diversity, helping student LMs acquire versatile and sophisticated knowledge.
Extensive experiments demonstrate that MiniPLM boosts the student LMs' performance on 9 common downstream tasks, improves language modeling capabilities, and reduces pre-training computation. 
The benefit of MiniPLM extends to larger training scales, evidenced by the scaling curve extrapolation.
Further analysis reveals that MiniPLM supports KD across model families and enhances the pre-training data utilization. Our code, data, and models can be found at https://github.com/thu-coai/MiniPLM.

Published as a conference paper at ICLR 2025
MINIPLM: KNOWLEDGE DISTILLATION FOR PRE- TRAINING LANGUAGE MODELS
Yuxian Gu!:2; Hao Zhou?, Fandong Meng", Jie Zhou?, Minlie Huang'
The CoAl Group, Tsinghua University WeChat AI, Tencent Inc., China
ABSTRACT
Knowledge distillation (KD) is widely used to train small, gerh-performing student language models (LMs) using large teacher LMs. While effective in fine-tuning, KD during pre-training faces efficiency, flexibility, and effectiveness issues. Exist- ing methods either incur high computational costs due to online teacher inference, require tokenization matching between teacher and student LMs, risk losing the difficulty and diversity of the seacher-generaae training data. In this work, we propose MINIPLM, a KD framework for pre-training LMs by refining the training data distribution with the teacher LM's knowledge. For efficiency. MINIPLM performs offline teacher inference, allowing KD for multiple student LMs without adding training costs. For flexibility, MINIPLM operates solely on the training corpus, enabling KD across model families. For effectiveness, MINIPLM lever- ages the differences between large and small LMs to enhance the training data difficulty and diversity, helping student LMs acquire versatile and sophisticated knowledge. Extensive experiments demonstrate that MINIPLM boosts the student LMs' performance on 9 common downstream tasks, improves language model- ing capabilities, and reduces pre-training computation. The benefit of MINIPLM extends to larger training scales, evidenced by the scaling curve extrapolation. Further analysis reveals that MINIPLM supports KD across model families and enhances the pre-training data utilization. Our code, data, and models can be found https::/// / /gi thub. com ttt/tttucooee coai/ M/MiniMMM.....
Average Accuracy (96)
Average Accuracy (%)
44
2.2x
44
42
Vanil a KD
Vanilla KD MiniPLM (Ours)
MiniPLM (Ours)
40
0.25
0.50
0.75
1.00 25 1.50
200M
500M
1.28
Training Computation (x102 FLOPs)
Student Model Size
(a) Computation Scaling (1.8B + 500M)
(b) Model Size Scaling
Figure I: Computation (a) and model size (b) scaling curves of student LMs pre-trained from scratch with Vanilla KD' and MINIPLM. The teacher LM has 1.8B parameters. "1.8B->000M" means we use 500M student LM. Training-time computation is kept constant for LMs of the same size in model scaling. The y-axis represents the LMs' zero-shot performance on 9 downstream NLP tasks.
INTRODUCTION
Recent advances in language models (LMs; Han et al., 2021; Bommasani et al., 2021: OpenAI, 2023; Touvron et al., 2023a) 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, small models typically far from compute-optimal
"Contribution during internship at Teneent Inc. (guyx21 @mai i 1s wwwnghin..eeeee edu edu.cn) cn) "Correspondig author.
Vanilla KD (Sanh et al., 2019; Muralidharan et al., 2024) minimizes the token-level forward Kullback-Leehllrr divergence between the output distributions of the teacher LM and student L.M.

---

Published as a conference paper at ICLR 2025
configurations according to Scaling Laws (Hoffmann et al., 2022). This has spurred growing interest in exploring the limits of small LMs under the constraint of pre-training computation (Lu et al., 2024).
Knowledge Distillation (KD; Hinton et al., 2015), where a small student LM learns from a large teacher LM, is promising approach for training high-performing small
Average Accuracy (%g
on Downerram fooks
Varilla KD (Oninw)
41
MiniLUM (Grline]
SeerD (Offline)
LMs. While KD is effective for fine-tuning (Xu et al.,
Train
2024), its role in improving pre-training, the critical stage for LMs to acquire foundation knowledge, remains under- explored. Applying KD during pre-trainn me-trrining with the meth- ods in fine-tuning stages is non- non-trivial, which can be cat- egorized as online KD and offiine KD. online KD, the teacher LM has to perform inference during pre-training to provide teven- probability supervision' (Sanh al., 2019; Gu et al., 2024b). introducing additional training- time 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 extend pre-training without KD. In addition, most online KD methods require the teacher and student LMs to share tokenization, limiting their flexibility for knowledge transfer across 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 (Kim & Rush, 2016; Gunasekar et al., 2023). However, ensuring sufficient difficulty and diversity in the generated data is challenging without extensive human expertise. Student LMs tend to overfit easy and common language patterns, hindering their downstream generalization (Shumailov et al., 2024), as shown in Figure 2.
Controlled Training Steps
Contralled
Training-Time FLOPs
Figure 2: Results of applying KD meth- ods in fine-tuning to pre-train a 200M student LM, using a 1.8B teacher LM. See Section 3.1 for method and evalua- tion details. When the training FL.OPs are controlled, all KD methods perform similar or worse than Pre-Train w/o KD.
To address these challenges, we propose MINIPLM, an efficient, flexible, and effective KD frame- work for pre-training LMs,, as illustrated in Figure 3. MINIPI.M typically works in scenarios where the teacher LM's probabilities are available to the users. Intuitively, it distills the teacher LM's knowl- edge 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, 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 (Chowdhery et al., 2022). In terms of effectiveness, Differrnne Sampling samples training instances that the teacher LM prefers but that small reference LM assigns low probabilities to, promoting data difficulty and diversity. As depicted in Figure (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 1.8B teacher L.M. 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 (Hoffmann et al., 2022), we observe that MINIPLM's 's benefit remains consistent for LMS trained on ~10T tokens. MINIPLM also facilitates KD across model families, improving Llama3.1 (Dubey et al., 2024) and Mamba (Gu& & Dao, 2023) with a teacher LM from the Qwen family (Bai et al., 2023). 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 (Villalobos et al., 2022).
2 MINIPLM: KD FOR PRE-TRAINING LMS
We consider pre-training an LM with an output distribution ge on a large-scale corpus D consisting of N text sequences, where 0 represents the model parameters. KD aids pre-training by incorporating
? Prre-computigg this supervision is impractical, as it requires 30PB of storage for 50B tokens with 150K vocabulary. Therefore, teacher LMss probabilities typically computed online during the student LM training.
2

---

Published as a conference paper at ICLR 2025
traninngg Conpus:
Teacher
Cross-Entropy Loss
Sifferece-Ssmmplee Corpus:
Topk fing
Student
40
D
D
Federence
Pref
p(a) Pru(e)
p(a) Pir((a)
p(s) (ra(a)
Easy and Common
Hard and Diverse Up-
Noisy and Harmeul
Offline Difference Sompling
Pre-Training
Down Samplee
Ramawed
(a) MINIPLM Training Framework
(b) Effect of Difference Sampling
Figure 3: MINIPLM. (a): Training framework. MINIPLM distills the knowledge of the teacher LM into the student LM by adjusting the pre-training corpus of the student LM (4g) through offline Difference Sampling, based on the output probability discrepancy between the teacher LM (p) and small reference LM (Pret). (b): Illustration of the effect of Difference Sampling, which down-samples common easy instances, up-samples hard valuable instances, and removes noisy harmful instances.
the knowledge of teacher LM with output distribution p. and training I to minimize the discrepancy between p and qe (Sanh et al., 2019). MINIPLM formulates KD as reward maximization prob- lem (Haarnoja et al., 2017), 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 the refined corpus (Section 2.3).
2.1 KD AS REWARD MAXIMIZATION
Recent works (Gu et al., 2024b: Agarwal al., 2024; Ko et al., 2024) have shown the effectiveness of minimizing the reverse -ulbback-Leibler divergence (KLD) between and qe for KD of LMs in the fine-tuning stage. which avoids qo from over-estimating the low-probability regions of p. We reformulate this objective as a reward maximizing problem (Haarnoja et al., 2017), which is suitable for designing efficient offline optimization methods (Levine al., 2020):
arg min KL [ge [Ip] arg min Ez log (x(x)
p(z)
(I)
arg max E Ex r (p, q0, x),
where the reward is defined r(p, go. x) log (o(æ) Intuitively, Eq. (1) trains qo to generate text
p(w)
a that receives high reward values, indicating a preference from the teacher LM (i.e., high log p(x) values), while also ensuring high diversity (i.e., low Ez~ge (x values). To optimize Eq. (1), simple yet effective approach is Best-of-N (Stiennon et al., 2020; Bai et al., 2022; Touvron et al., 2023b), where set Dgo containing M candidates is first sampled from qe and K instances with the highest rewards are selected from these candidates to form a new dataset D
D' Kop-K(((pp go. x1) x E Dgo
where Dao = (1mmm2. Zm 20, 5 M). The student LM is then trained on D' to learn to generate texts with large r(p, go, .)) values. Although Best-of-N achieves performing the teacher LM's inference prior to the student LM training similar to offline KD (Kim & Rush, 2016; Gunasekar et al., 2023), it still lacks the efficiency advantage of these methods because obtaining D' requires (1) sampling data from qo and (2) computing the reward values with qe. making D' non- transferable for pre-training other student LMs. It is also hard to ensure the diversity and difficulty of the M candidates sampled from de 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.
3

---

Published as a conference paper at ICLR 2025
2.2 DIFFERENCE SAMPLING
As shown in Figure 3(a), Difference Sampling refines the pre-training corpus D based on the dis- crepancy between p and the output distribution Pter from a tiny reference LM, which eliminates the dependency of Eq. (2) on go. making the sampled corpus reusable in training multiple student LMs.
Top-K Sampling From D, not Deo. To avoid sampling data from qo, we sample instances with high r(p, qo, æ) values from the pre-training corpus D, rather than from the Dgo generated by the student LM as Eq. (2). This way, changing the student LM does not affect the candidate set, and 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 Dgo and 'D are highly likely to be the same when the sizes of D4o and D are sufficiently large:
Proposition 2.1. Let S be the sample space of two distributions P1 and pz, X2, .NN Pi be N ii.d random variables, and 1, Yy,'.., Ys ~ P2 be M i.i.d random variables. Let r(.): S R be any injective function. Assume that V' S. p1 (z) >> O. p2(2)) > O. For fixed K satisfying < K S min (N, M), when N + +o0 M - DG, we have
P (top-K {r(X,)]1 S n S N = top-K (r(Ym))l
m M)) +
The proof of Proposition 2.1 is provided in Appendix A. Here, P1 represents the data distribution
of D, P2 = ge. which is the sampling distribution of Doo. and r(æ)) r(p, qa, æ). Intuitively, Proposition 2. reveals that when [D] N and [P = M are sufficiently large, the top-K instances selected from both sets tend to overlap. This suits our scenario well, as D is typically large-scale, and it is advantageous to sample many candidates from Dgo possible to find high-reward instances.
Decoupling the Student and the Reward-Computing LM. To avoid computing reward values with qe. replace it with Pref, the output distribution of tiny reference LM, typically smaller than the student LM, for reward computation. The reference LM is pre-trained on a small subset Diet uniformly sampled from D, allowing Pret to approximate ge with minimal computation. This is a reasonable approximation because qe evaluates the difficulty of instances, and the relative data difficulties generally remain consistent across different models (Ethayarajh al., 2022). As result, the reward function in Eq. (2) is replaced with r(p, Prefs Φ) = log p(a) In Appendix D.3, we p=)
empirically show that rpp, Prets æ) and r(p, frer: æ) have high correlations.
In summary, Difference Sampling constructs a pre-training corpus D' from D D Drer as follows:
p(x)
t'=top-K K log (x)
x E D Diet D
which is independent of ge. As shown in Figure 3(b). Difference Sampling essentially refines the data distribution of 'D by comparing the "difference" between and Prer. 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-
CORPUS
As illustrated in Figure 3(a), we pre-train the student LM from scratch on the difference- slfferenne-sampled corpus D' with the cross-entropy loss for next-token prediction, which is similar standard pre- training. The loss function T.(go, D') is given by
- log go ( (xe)
1
Mgo, D¹)
(5)
D
RED
where [2| is the length of x, Iz is the thh token, and x<e denotes the prefix of a with t tokens.
2.4 MINIPLM TRAINING PIPELINE
The general training pipeline of MINIPLM is as follows: (1) Uniformly sample a subset Dier from the
pre-training corpus D. ensuring 'eerl < (D). (2) Train reference LM from scratch on Dier using the cross-entropy loss to obtain its output distribution Pret- (3) Perform Difference Sampling with Pet and the teacher LM's 's output distribution p. generating a pre-training corpus D' from D Dret using Eq. (4). The size of D' is controlled by a sampling ratio a, where K a (D Dretl. (4) Pre-train student LMs from scratch on D' with the cross-entropy loss defined in Eq. (5).
4

---

Published as a conference paper at ICLR 2025
2.5 DISCUSSION
Efficiency and Flexibility of MINIPLM. As shown in Figure 3(a), MINIPLM relies only on p(ar) and (r) the teacher and reference LMs' probability over the entire sequence, which can be computed and stored offline because each instance in D is associated with a single floating-point number, amounting only 200MB storage for 50B tokens with a 1,024 sequence length. Once based on p(æ) and Pre((a), multiple student LMs can be efficiently pre-
D' is tifference-samplee 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 (Chowdhery et al., 2022) and suitable for KD across model families. In contrast, online KD (Muralidharan et al., 2024; Gu et al., 2024b) relies on per-token distributions, which infeasible store offline, as it takes 50B x 150K x 4byte 30PB 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 (Boizard et al., 2024), 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- pre-training distribution via Difference Sampling, producing three effects: (1) Dawn-sampling easy and common patterns that both the teacher and reference LM p(æ)
fit well, where p(z) Z Pa) and log Pirlar) Z 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(æ)= >> (rer(a)) and log pla)) O. (3) Discarding noisy and harmful instances that the teacher LM
Pre((z)
< 0. We provide examples of these effects in Appendix F. 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 pref. 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 (Kim & Rush, 2016; Peng et al., 2023), where maintaining the data difficulty and diversity, critical to pre-training (Shumailov et al., 2024), challenging without extensive human efforts on prompt engineering (Gunasekar et al., 2023).
assigns lower probabilities to than the reference LM, where log p(x)
Parla)
3 EXPERIMENTS
3.1 EXPERIMENTAL SETUP
Model. We adopt the Qwen-1.5 (Bai et al., 2023) 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- Tre-Training. We construct pre-training corpora from the Pile (Gao et al., 2020). To control the computation in experiments, we pre-traie all LMs on 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 align the total training computation with pre-training without KD or offline KD methods. See Appendix B for more pre- veeirrain details.
Baselines.
We compare MINIPLM with 4 baselines:
Pre-Train w/o KD pre-trains the student LM on 50B corpus uniformly sampled from the Pile dataset, without the guidance of the teacher LM's knowledge.
Vanilla KD (Muralidharan et al., 2024) minimizes the token-level forward KLD between p and ge- which requires online inference of the teacher LM to obtain the token-level output distributions.
SegKD (Kim & Rush, 2016) trains the student LM on the teacher-genen generated data. Since it is infeasible to generate all 50B tokens, we approximate Kim & Rush (2016) 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.
Mi (Gu et al., 2024b) minimizes the reverse KLD between p and go with PPO (Schul- man al., 2017), which requires online inference of the teacher LM and online sampling
5

---

Published as a conference paper at ICLR 2025
HS
LAM
Wino
OBQA
ARC-e
ARC-c
PIQA
SIQA
Story
Avg.
1.8B Teacher
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
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 500M Student
Pre-Train 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
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
18B Teacher > 1.2B Student
Pre- Pra-Train 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
46.6
Table 1: Zero-shot accuracy
9 widely-used downstream tasks and the average (Avg.). We use the Qwen-1.5 1.8B LM (Bai et al., 2023) as the teacher and Qwen LMs with 200M, 500M, and 1.2B parameters the student. Student LMs with the same sizes consume the same training-time computation. The best scores of each model size are boldfaced.
from the student LMM 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 qo during the exploration of PPO.
MINIPLM. We employ a 104M reference LM trained on 5B tokens. Section 3.2 and 3.3, we consider setting where D is sufficiently large, containing 105B tokens uniformly sampled from the Pile corpus. We reserve tokens as Dres and conduct Difference Sumpling as per Eq. (4) on other 100B tokens by setting = 0.5 to construct 50B-token corpus D'. In this way, the student LM pre-trained on D' for one epoch. We use the loss difference the teacher and reference LM to sampled instances from D, which is equivalent to Eq. (4) all instances have 1,024 tokens. In Section 3.4, we evaluate MINIPLM
data-limited setting, where D is controlled to contain 50B tokens and the student LM trained on the difference-sampled data for multiple epochs. See Appendix D.3 for the ablation studies on the reference LM and sampling ratio a.
Evaluation. We assess the zero-shot accuracy of LMs trained with different methods on 9 down- stream tasks widely used in examining the foundation abilities of base models (Touvron al., 2023a; Groeneveld et al., 2024). We also test the language modeling capability of the LMs on a subset of DCLM (Li et al., 2024a), 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.
3.2 MAIN RESULTS
MINIPLM Improves Downstream Performance. Table shows zero-shot accuracy on down- stream 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 constant computation budget. This highlights the room for improvement in KD for pre-training, es- pecially when the gap between the teacher and the student is substantial. Second, MINIPLM-traine 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, incurs no additional training- timin time overhead, allowing the student LM to be optimized for more steps and leading higher performance. Compared offline methods like SegKD, MINIPLM as shown
shown in Figure 3(a),

---

Published as a conference paper at ICLR 2025
3.8
Pre-Triin w/o KD MiniPLM
3.5
Pre-Train w/o KD MIM
3.4
Pre-Tain w/o KD MiniPLM
3.7
Testt Loss
3.4
3.3
Test
3.6
3.3
3.2
3.5
3.4
3.1
0.2
0.6
1.5 Computation (x1020 FLOPs)
0.5
1.0
0.5
2.5
3.5
Computation (x1020 FLOPs)
Computation (x1020 FLOPs)
(a) 188 Teacher + 200M Student
(b) 1.8B Teacher - 500M Student
(c) 1.8B Teacher + 1.2B Student
Figure 4: Language modeling loss on the DCLM (Li et al., 2024a) subset. We distill the knowledge of the 1.8B Qwen model (Bai al., 2023) into student LMs from the Qwen family with 200M, 500M, and 1 1.2B parameters. We control the total training-time FL.OPs of different methods to be the same.
in Figure 3(b), uses 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 MINIPIM against Vanilla KD) scale well with the student LM size, which is also illustrated in Figure 1(b).
MINIPLM Helps Language Modeling. Figure 4, compares the language modeling performance of the MiNiPLM-traine
Natu
Method
LiT
T L
LMs and baselines on the DCLM (L.i et al., 2024a) subset, a diverse and high- quality dataset curated from web corpora. The re- sults 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 (Hoffmann et al., 2022) to sim- ulate pre-training on IT and 10T tokens (details in Appendix D.2), showing thatMINIPL.M maintains its advantages at the scales of pre-training recent large LMs (Touvron et al., 2023a; Dubey et al., 2024). crit- ical stage in DCLM's data-cleaning pipeline involves removing easy and common patterns, thereby enhanc- ing challenging and diverse signals. Therefore, 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).
Pre- Train w/o KD
3.35
3.32
200M
Vanilla KD
3.39
3.35
MINIPLM
3.28
3.26
Pre- vre-Train w/o KD
3.12
3.08
500M
Vanilla KD
3.12
3.07
MINIPLM
3.06
3.04
Pre-Train w/o KD
2.98
2.94
1.2B
Vanilla KD
2.95
2.91
MINIPLM
2.92
2.88
Table 2: Test loss predictions using the Scal- ing Law (Hoffmann et al., 2022). Naw: the student LM size. LIT, LIOT: the loss when Pre- Train w/o KD and MINIPLM process IT and 10T tokens, with Vanilla KD con- suming the same training FLOPs.
MINIPLM Reduces Training Computation. We plot the 500M student LM's average zero-shot accuracy scores on the downstream tasks in Figure I(a) with respect to its pre- 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, shown in Figure 3(b). As result, the model trained with MINIPLM avoids wasting computation on learning the easy knowledge quickly memorized during the early training stage and less distracted by the noisy outliers that slow the convergence down.
3.3
KD ACROSS MODEL FAMILIES
A noticeable advantage of MINIPLM over Vanilla KD and MiniLLM is its flexibility to distill the knowledge of teacher LM into student LMs with completely different tokenizers and architectures without additional strategies like Boizard et al. (2024). In Table 3, illustrate the performance
7

---

Published as a conference paper at ICLR 2025
Llama3.1
Mamba
Pre-Trinn wlo KD
3.6
MiniPLM (ar- 0.5, Eps)
Acc.
Loss
Acc.
Loss
Test Loss
MinIPLM 0.25, 4 Epp)
MInIPLM (ar 0.125, Eps)
Pre-Traie TraiTrain w/o KD
41.0
3.52
41.6
3.24
3.5
1.p.
Ep.
SeqKD
40.8
3.54
41.0
3.27
Ep.
MINIPLM
41.8
3.43
42.6
3.15
3.4
4Ep
Ep
12.58
25B
37.5B
508
Trained Takens
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 mod- els. The average zero-shot accuracies on the down- stream tasks and the losses on the DCLM corpus are reported. Note that Vanilla KD and MiniLLM cannot be applied when the teacher and student
Figure 5: MINIPLM in the data-constrained setting. We fix DD to contain 50B tokens and alter the sampling ratio a to obtain D' with Difference Sumpling, which will be trained on for multiple epochs to achieve the constant total trained tokens. The y-axis represents the test loss on the DCLM corpus.
LMs use different tokenizations.
when using the LMs from the Qwen (Bai et al., 2023) family as the teacher and the reference LM to distill knowledge into a 212M Llama3.1 (Dubey et al., 2024) model and 140M Manba (Gu & Dao, 2023) 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 (D) This allows emerging LMs with novel architectures (Lieber et al., 2024; Sun et al., 2023) or advanced tokenization (Tay et al., 2022; Godey et al., 2022) to inherit knowledge from existing LMs, thereby facilitating the development of more efficient and
models.
3.4 DATA-LIMITED SETTING
We evaluate MINIPLM in a data-limited setting where D is constrained to contain 50B tokens. To this end, the student LM should be trained the difference-sampledd sampled D' over multiple epochs to ensure the total computation and trained tokens remain consistent with Pre- Pre-Train without KD. We split |Derl containing 1B tokens train the reference LM. Therefore, for sampling ratio o, the
student LM should trained for Darl 1 epochs, given that [Dirr| < [D]. In Figure 5, we plot the loss curves of the 200M student LMs on the DCLM corpus when using E [0.5, 0.25, 0.125] and 25%
training the LM for around 2, 4, and 8 epochs, respectively. We can see that difference-samplinn data (with 0.25) and training the student LM for 4 epochs yields the best performance, which aligns with the observations in Muennighoff et al. (2023). The corpus sampled with a higher a does not achieve the best quality and diversity offered by Difference Sampling, while lower leads to rapid over-fitting of the student LM. These findings suggest that MINIPLM is promising approach to enhance data utilization when high-quality web corpora become scarce ( (Villalobos et al., 2022). By extrapolating the loss curve of Pre-Train w/o KD using the Scaling Law in a data-constrained setting (Muennighoff et al., 2023), we estimate that would require an additional 68B training tokens to match the performance of MINIPLM (a 0.25, 4 epochs), which means MINIPLM reduces the pre- pre-training data requirement by 2.4 times. See Appendix D.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 observed in recent works (Gu et al., 2024b). However, early works have also shown that an excessively large gap between teacher and student models can hinder effective KD (Mirzadeh et al., 2020). 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. (2020) 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 200M student LM. Smaller teacher LMs (c.g.. 300M) lack the capacity to identify the hard but valuable parts of the pre-training distribution, weakening the effectiveness of Difference Sampling shown in Figure 3(b). On the
8

---

Published as a conference paper at ICLR 2025
Average Accuracy (%)
Pre- Tre-Training Corpus
Usage
Diversity
41
Pre-Train Vanilla KD
Tre-Trrai w/o KD
Original
32.25
& Vaniilla KD
40
MiniPLM
Teacher-Generatee
SeqKD MINIPLM
30.16
Difference -aanleed
36.70
39
300M
500M
1.8B
4B
Size of Teacher Model
Table 4: Semantic diversity of the original pre-training corpus DD used in Pre-Train w/o KD and Vanilla KD, the teacher- -eeeherrgeeertedd corpus used SeqKD, and the cufference-samplee corpus DD' in MINIPLM. Differ- ence Sampling increases the diversity of the refined pre-training distribution, which helps LM pre-training.
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.
other hand, the value scale of (og((a)) from oversized LMs (e.g., 4B) becomes too small compared to that the reference LM, which tends to degenerate Difference Sampling into sampling with Pee only, losing the effect of the teacher LM. Future research could focus optimizing teacher model size for MINIPLM or mitigating the impact of differing log p and log Pie((z) value scales.
Diversity of Difference-Sampled Data. To further verify the effectiveness of Difference Sampling on improving the diversity of the pre-training corpus, in Table 4, we follow Friedman & Dieng (2023) corpus and the data used in other baseline
to compute the semantic diversity of the difference-samppl pre-training approaches. The results show that the difference--amplee corpus has the highest diversity, despite that Difference Sampling is derived from minimizing the reverse KLD, which exhibits the mode-seeking mode-seed seeking behavior (Minka et al., 2005). 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(x) values as seen in Figure 3(b), constitute the major modes of the teacher LM that yo seeks during optimizing Eq. (1). Therefore, the modc-secking 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 case study in Appendix F to further explain our argument.
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 Pre-Train w/o KD, suggesting the potential of fur- ther 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. 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- gaasseerained sequence-level knowledge of the teacher LM into the student LM via the pre-training data, while Vanilla KI) directly aligns the token-level probability distribution between p and qo, providing fine-grained token-level signals.
Nam
Method
Acc.
Vanilla KD
39.9
200M
MINIPLM
41.3
MINIPLM + Vanilla KD
40.7
Vanilla KD
43.6
Table 5, we
500M
MINIPLM
44.8
MINIPLM + Vanilla KD
44.9
Vanilla KD
45.4
1.2B
MINIPLM
46.6
MINIPLM + Vanilla KD
48.1
Table 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-sampless corpus in MINIPLM. N the size of student LMs.
RELATED WORK
Language Model Pre- Pre-Training Pre-training is the critical phase for language models (LMs; Brown al., 2020; OpenAl, 2023; Team al., 2024; Chowdhery et al., 2022; Touvron et al., 2023a;b; Dubey

---

Published as a conference paper at ICLR 2025
et al., 2024) to obtain their foundation abilities for various downstream tasks. To improve pre training, some works focus on data curation, such as adjusting domain mixing (Xie et al., 2024; Ye et al., 2024), selecting valuable data points relevant desired tasks (Brandfonbrener et al., 2024; Engstrom et al., 2024), or transforming the instances based on downstream requirements (Cheng et al., 2024; Gu et al., 2022; 2023). Another line of work improves the optimization during pre-training by solving better data reweighting strategies (Gu et al., 2024a), designing more effective optimizers (Liu et al., 2024; Shazeer Stern, 2018), or discovering better training recipes (Hu al., 2024; Hoffmann et al., 2022). Variations in model architectures (Xiong et al., 2020) and training objectives (Tay et al., 2023) are also explored to boost pre- -reiiian 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 (Sardana al., 2024). 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 (Hoff- mann et al., 2022) and scales as a power law with respect to the model size gap (Kaplan et al., 2020). Despite these challenges, recent efforts have made promising progress by data quality improve- ment (Mehta et al., 2024; Bellagente et al., 2024; Zhang et al., 2024) or model pruning (Muralidharan al., 2024; Xia et al., 2024). We explore knowledge distillation a complementary approach.
Knowledge Distillation. Knowledge distillation (KD;Hinton et al., 2015) uses a large teacher model improve the performance of a small student model, which is widely used build efficient neural network systems (Park et al., 2019; Czarnecki et al., 2019; Salimans & Ho, 2022). In NLP, early works primarily apply KD to encoder-only models (Devlin et al., 2019; Liu al., 2019) for text classification by aligning token-level distribution (Sanh et al., 2019), hidden states (Sun et al., 2019), and attention matrices (Wang et al., 2020; 2021). For generative LMs, a straightforward KD approach is training small LMs on the texts generated by large LMs (Chiang et al., 2023; Peng et al., 2023; Hsich et al., 2023). Other works (Gu et al., 2024b; Agarwal et al., 2024; Wen et al., 2023; Li et al., 2024b; Wu et al., 2024) explore better optimization objectives. However, these works focus on KD for fine-tuning LMs, while pre-training is critical for establishing core LM capabilities (Allen-Zhu & Li, 2024). Therefore, 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 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-mont and diversity of the refined pre-training distribution. Using 1. 1.8B LM 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.
differences ensures the difficulty
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 (OpenAI, 2023; Team et al., 2023) with MINIPIM challenging. For some APIs (OpenAI, 2022)) 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 (Carlini et al., 2024), at the expense of 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 (Burns et al., 2023). Since data properties are critical for pre-training LMs across various sizes, the improvement of data diversity and difficulty is likely be beneficial for LMs larger than the difference-sampling models.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENTS
This work is supported by the National Science Foundation for Distinguished Young Scholars (with No. 62125604) and Tsinghua University Initiative Scientific Research Program.
REFERENCES
Rishabh Agarwal, Nino Vieillard, Yongchao Zhou, Piotr Stanczyk, Sabela Ramos Garea, Matthieu Geist, and Olivier Bachem. On-policy distillation language models: Learning from self- generated mistakes. In Proceedings of ICLR, 2024. URL http:///ppnneeeeeeetteeeeeeeeeee
extraction. In Proceedings of ICML, 2024. URL /
Gozum?idd3<Ktassill
Zeyuan Allen-Zhu and Yuanzhi Li. Physics of language models: Part 3.1, knowledge storage and net / wtt.mmooummmm
5x788rgboj.
Jinze Bai, Shuai Bai, Yunfei Chu, Zeyu Cui, Kai Dang, Xiaodong Deng, Yang Fan, Wenbin Ge, Han, Fei Huang, et al. Qwen technical report. arXiv preprint urXiv:2309.16600 2023. URL https:/ /aaxivvvvvv.. wwt/oog/////ooooo.ootttptt..wttt 2tt.///399........... 16609.
Yuntao Bai, Saurav Kadavath, Sandipan Kundu, Amanda Askell, Jackson Kernion, Andy Jones, Anna Chen, Anna Goldie, Azalia Mirhoseini, Cameron McKinnon, al. Constitutional ai: Harmlessness from ai feedback. arXiv preprini arXiv:2212.08073, 2022. URL https: www.iiiogg/..o.ooo..ooooo.o /bs// 2212. 08073.
Marco Bellagente, Jonathan Tow, Dakota Mahan, Duy Phung, Maksym Zhuravinskyi, Reshinth Adithyan, James Baicoianu, Ben Brooks, Nathan Cooper, Ashish Datta, et al. Stable LM 1.6B technical report. arXiv preprint arXiv:2402. 17834, 2024. URLhttte / 2402 17834.
Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, al. Piqa: Reasoning about physical commonsense in natural language. In Proceedings of AAAI, 2020. URL htttt::////jj aaai. org/index.php/ AAAIAAAAAASt....AA htt///aaticetevieeeeeeeeeeeeeeee 6239/ 6095.
Nicolas Boizard, Kevin El-Haddad, Céline Hudelot, and Pierre Colombo. Towards cross-tokenizer distillation: the universal logit distillation loss for llms. arXiv preprint arXiv:2402.12030, 2024.
URL https:tt////..... /tt//aiivvv./////////aa/////atva orgg abs/2402 12030.
Rishi Bommasani, Drew A Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney von Arx, Michael S Bernstein, Jeannette Bohg, Antoine Bosselut, Emma Brunskill, et al. On the op- portunities and risks of foundation models. arXiv preprint arXiv:2 108.:2008..7288, 2021. URL https: org/ abss 2108 07258.
David Brandfonbrener, Hanlin Zhang, Andreas Kirsch. Jonathan Richard Schwarz, and Sham Kakade. CoLoR-Filter Conditional loss reduction filtering for targeted language model pre- pre-training. arXiv preprint arXiv:2406. 10670, 2024. URL httpp:://aettte....eeeeee /taaxi//axxiiv wrggabas.......oooo...ee..eeeebe 2406 10670.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, et al. Language models are few-shot learners. In Proceedings of NeurlPS, 2020. URL https: //aper//apppeaeee/ee nips cass.ce/paa cc/paper/
01144418fbBaccccc 42f6:ac4444444 wbstteaaaaaaaaatto..oo html.
Collin Burns, Pavel Izmailov, Jan Hendrik Kirchner, Bowen Baker, Leo Gao, Leopold Aschenbrenner, Yining Chen, Adrien Ecoffet, Manas Joglekar, Jan Leike, et al. Weak-to-strong, generalization: Eliciting strong capabilities with weak supervision. arXiv preprint arXiv:2312.093390 2023. URL https: / com/
weak-to- wtt.---tttttto--eeeeaaaaa.aooo.t
Nicholas Carlini, Daniel Paleka, Krishnamurthy Dj Dvijotham, Thomas Steinke, Jonathan Hayase, A Feder Cooper, Katherine Lee, Matthew Jagielski, Milad Nasr, Arthur Conmy, et al. Stealing part of a production language model. Proceedings of ICML, 2024. URL https: petttps://axxv...o/ //axiiv.ogg/ abs/24033 06634.
11

---

Published as a conference paper at ICLR 2025
Daixuan Cheng, Yuxian Gu, Shaohan Huang, Junyu Bi, Minlie Huang. and Furu Wei. Instruction pre-training: Language models are supervised multitask learners. In Proceedings of EMNLP, 2024.
URL https: /atpss///rxaae wwtaxviivv.oo.////////aa org/ abs//400 220.61449950 14491.
Wei-Lin Chiang, Zhuohan Li, Zi Lin, Ying Sheng, Zhanghao Wu, Hao Zhang, Lianmin Zheng, Siyuan Zhuang, Yonghao Zhuang, Joseph E. Gonzalez. lon Stoica, and Eric P. Xing. Vicuma: An open-source chatbot impressing gpt-4 with 90%* chatgpt quality, March 2023. URL https:
// //t.m/sssorr/////.......... og/ -103--033--iiiii
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung. Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. arXiv preprint arXiv:2204.02311, 2022. URL https:
//arxiv. wrwpvv.orr//aaa abs///ab///// 4204. 02311.
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge.
/03.vv/800.0555700 2018. URL https: //arivv///rvvvaaa/aa/ org abs 1800. 0
Wojciech M Czamecki, Razvan Pascanu, Simon Osindero, Siddhant Jayakumar, Grzegorz Swirsze7, and Max Jaderberg. Distilling policy distillation. Proceedings of AISTATS, 2019. URL htttp://proceeding/ mlr. press / /v////89/////
http///cannnnnaa/ee/eea wwtp:///waa.aaaa.aiaaaaa.a pdf.
Jacob Devlin, Ming-Wer Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of NAACL-HLT, 2019. URL
Owg///g/ /N1 9 14. pdf.
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, al. The llama herd of models. arXiv preprint arXiv:2407.21783, 2024. URL https: http//ixvvvv///// org/ borg/abl aaaaabaa4000 21783.
Logan Engstrom, Axel Feldmann, and Aleksander Madry. Dsdm: Model-aware dataset selection with datamodels. arXiv preprint arXiv:2401.12926, 2024. httpstts: http::::///xxx..o//////// abs/ 2401. 12926.
Kawin Ethayarajh, Yejin Choi, and Swabha Swayamdipta. Understanding dataset difficulty with v-usable information. In Proceedings of ICML, 2022. URL press/vieee (tt.://etaaaaaaaaa222
mlr.
pdf.
Dan Friedman and Adji Bousso Dieng. The vendi score: A diversity evaluation metric for machine learning. Transactions on Machine Learning Research, 2023. ISSN 2835-8856. URL https:
// //openrevieww.om net/ forummid=gge hotp://ooorm22/gg000000///
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, et al. The pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2101.00027 2020. URL /ettpss//arxiv...o abs/2101. 00027.
Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, al. A framework for few-shot language model evaluation, 07 2024. URL
Nathan Godey, Roman Castagné, Éric de la Clergerie, and Benoît Sagot. MANTa: Efficient gradient- based tokenization for end-to-end robust language modeling. In Findings of EMNLP, 2022. URL
https://aclle wnthocllaltooooooooooooooooooooo ogy. org/ 2022. fi andiiingg nniigg--nnenm 1p.207.
Dirk Groeneveld, Iz Beltagy, Pete Walsh, Akshita Bhagia, Rodney Kinney, Oyvind Tafjord, Ananya Harsh Jha, Hamish Ivison, Ian Magnusson, Yizhong Wang, et al. OLMo: Ac- celerating the science of language models. arXiv preprint arXiv:2402.00838 2024. URL attps://azze
org/ abs/ 220224222.00 00838.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.0073 2023. URL httttss://rxxvv.errtt
http://wwwww..o////////////////a /2312. 00752.
12

---

Published as a conference paper at ICLR 2025
Yuxian Gu, Xu Han, Zhiyuan Liu, and Minlie Huang. PPT: Pre-trained prompt tuning for few-shor
learning. Proceedings ACL, 2022. URL https://aaxxiyy org/ abs/2109.04m
Yuxian Gu, Li Dong, Furu Wei, and Minlie Huang. Pre-training to learn in context. In Proceedings ACL, 2023. URL https:tpp/... http://:s..ccaaaaooooooo./////..
Yuxian Gu, Li Dong, Yaru Hao, Qingxiu Dong, Minlie Huang, and Furu Wei. Towards optimal learning of language models. arXiv preprint arXiv:2402.17759, 2024a. URL
org/abs / 2400217755
Yuxian Gu, Dong, Furu Wei, and Minlie Huang. MiniLLM: Knowledge distillation large lan- guage models. In Proceedings of ICLR, 2024b. URL https:://tppeeeeeeeeeeee 1d=5h0gf7IB22.
www. forum?
Suriya Gunasekar, Yi Zhang, Jyoti Aneja, Caio César Teodoro Mendes, Allie Del Giorno, Sivakanth Gopi, Mojan Javaheripi, Piero Kauffmann, Gustavo de Rosa, Olli Saarikivi, et al. Textbooks are all you need. arXiv preprint
2023.
Tuomas Haarnoja, Haoran Tang, Pieter Abbeel, and Sergey Levine. Reinforcement learning with deep eeery-bbased policies. In Proceedings of ICML, 2017. URL http: httt::::////cceiiggs.... mlr. prw.sppesseeee v0 wwwaaaaaaaaaaa al
/ttp:/://ggtttttt/gg/g com.
Xu Han, Zhengyan Zhang, Ning Ding, Yuxian Gu, et al. Pre-trained models: Past, present and future. A/ Open, 2021. URL https:: htwp///wwww///////////////t/ www.ssceeeeeerertt...c com//science//rttcciee....
026666510210020 0220002319.
Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in neural network. arXiv preprint arXiv:1503.0253: 2015. URL
wwwoorggg/o..... /10/3/503..0 02531.pdE
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Flena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, al. Training compute-optimal large language models. In Proceedings of NeurIPS, 2022. URL https::/prooceedings neurips cc/paper foeer_fille fileflllee fiwes/ pttpa///pppepe/2///ea..../
Cheng-Yo Yu Hsich, Chun-Liang Li, Chih-kuan Yeh, Hootan Nakhost, Yasuhisa Fujii, Alex Ratner, Ranjay Krishna, Chen- Yu Lee, and Tomas Pfister. Distilling step-by-step! outperforming larger language models with less training data and smaller model sizes. In Findings of the ACL, 2023.
URL https:////// / /acllallleee anthology .rgg 2023. findings-ack.
Shengding Hu, Yuge Tu, Xu Han, Chaoqun He, Ganqu Cui, Xiang Long, Zhi Zheng, Yewei Fang, Yuxiang Huang, Weilin Zhao, et al. MiniCPM: Unveiling the potential of small language models with scalable training strategies. In Proceedings of COLM, 2024. URL https: pttetts://pperreee /opp//epee....... net/ forum? ddmmmmi1==d222F--m wwtp://rrrrtoooeeeeeoeooeea httpdd--wwwwppcch....
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprini arXiv:2001.08361. 2020. URL https: http://axvvvvoo/// httpp//////xxvv..//bbb..... wwwp:r/g/aab///ooo/o////////// /atpbb/2011.3330 08361.
Yoon Kim and Alexander M Rush. Sequence-level knowledge distillation. In Proceedings of EMNLP, 2016. URL htttps://cct http://aaaaattooooo..ooooo www.hhhlll....ggg/eeeeoo
pdf.
Jongwoo Ko, Sungnyun Kim, Tianyi Chen, and Se- Young Yun. Distillm: Towards stream- lined distillation for large language models. In Proceedings of ICML, 2024. URL https:
/ ///ppn/eview.w.ee.mee net forum? ( sHZNNoC7r.
Hector Levesque, Ernest Davis, and Leora Morgenstem. The winograd schema challenge. In Procedd- 5555/303111111
ings of KR, 2012. URL https: //11 acm. org/ /5//5//53/35.3333311133 080003333330030000 3031 909.
Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.0164) 2020. URL
https://arxte wwtp://////p/////o/oo/ 01643.ppd
13

---

Published as a conference paper at ICLR 2025
Jeffrey Li. Alex Fang, Georgios Smymiss, Maor Ivgi, Matt Jordan, Samir Gadre, Hritik Bansal, Etash Guha, Sedrick Keh, Kushal Arora, et al. DataComp-LM: search of the next generation of training sets for language models. arXiv preprint 11Xi::2406.1199,, 2024a. URL https:
//arxiv. wwwpxvv.oorr.eoooo abs 22404 11794.
Yixing Li, Yuxian Gu, Li Dong, Dequan Wang, Yu Cheng, and Furu Wei. Direct preference knowledge distillation for large language models, arXiv preprint arXiv:2406. 19774, 2024b. URL https: arxiv..iiivve.... org/ abs 2406.197744
Opher Lieber, Barak Lenz, Hofit Bata, Gal Cohen, Jhonathan Osin, Itay Dalmedigos, Erez Safahi, et al. Jamba: A hybrid transformer-
Shaked Meirom, Yonatan Belinkov, Shai mamba language model. arXiv preprint arXiv:2403. 2024. URL httptpp://axi........ abs/2405 19887.
Hong Liu, Zhiyuan Li, David Leo Wright Hall, Percy Liang, and Tengyu Ma. Sophia: A scalable stochastic second-order optimizer for language model pre-training. In Proceedings of ICLR, 2024. URL https:://////mc http///ppeeeeeeeeeeeeeeeeeeeeeet net forrmm ? i =d= 3xHDeABNoi.
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. RoBERTa: A robustly optimized BERT pretraining approach. arXiv preprint arXiv:1907.11692 2019. URL https: 11692.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In Proceedings of ICLR, 2019. URL https:t////.mm wttp//ppnweeeeeeeeeeeee forum? Big.rrmmmiiddggkkiccs...
Zhenyan Lu, Xiang Li, Dongqi Cai, Rongjie Yi, Fangming Liu, Xiwen Zhang, Nicholas Lane, and Mengwei Xu. Small language models: Survey, measurements, and insights. arXiv preprint arXiv:2409. 15790, 2024. URL https: http://axxvvvr//o /RLaattes///aravvvorrr/bbb/////
15790.
Sachin Mchta, Mohammad Hossein Sekhavat, Qingqing Cao, Maxwell Horton, Yanzi Jin, Chenfan Sun, Iman Mirzadeh, Mahyar Najibi, Dmitry Belenko, Peter Zatloukal, et al. OpenELM: An efficient language model family with open-source training and inference framework. arXiv preprint arXiv:2404. 140.99 14619, 2024. URL https: /ttp/aaaavvv/////////////e//aaaa org/absg///..oooooooooeo.ee 2404 14619.
Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. Can a suit of armor conduct electricity? a new dataset for open book question answering. In CProceedings of EMNLP, 2018.
URL https: ///paapp..p app. wttpppp.....a.aa.mmmaammmmamaaaa semanti org/ CorpusID: 52183757.
Tom Minka et al. Divergence measures and message passing. Technical report, Cite- seer, 2005. URL htttpp::/www.mmcrr/// uploads/201s 20lddd////o.// 001000/6602/11220-17722
Seyed Iman Mirzadeh, Mehrdad Farajtabar, Ang Li, Nir Levine, Akihiro Matsukawa, and Hassan Ghasemzadch, Improved knowledge distillation via teacher assistant. In Proceedings of AAAI, 2020.
URL httpt::t///ooo //jss aaai. http:///w..aaaaaaaa/aaaaaaaa/aa/
Nasrin Mostafazadeh, Nathanael Chambers, Xiaodong He, Devi Parikh, Dhruv Batra, Lucy Van- derwende, Pushmeet Kohli, and James Allen. corpus and cloze evaluation for deeper un- derstanding of commonsense stories. In Proceedings of NAACL-HLT, 2016. URL https:
//aclanthology
Niklas Muennighoff, Alexander M Rush, Boaz Barak, Teven Le Scao, Nouamane Tazi, Aleksandra Piktus, Sampo Pyysalo, Thomas Wolf, and Colin Raffel. Scaling data-constrained language models. In Proceedings of NeurlPS, 2023. URL https: //tpsr////ppeeevieww net net/forum?id= J5BuTrEj35.
Saurav Muralidharan, Sharath Turuvekere Sreenivas, Raviraj Joshi, Marcin Chochowski, Mostofa Patwary, Mohammad Shoeybi, Bryan Catanzaro, Jan Kautz, and Pavlo Molchanov. Compact language models via pruning and knowledge distillation. arXiv preprint 47X1v:2407.14679, 2024. URL https:://// /tww//wwww....//../////ww att.a/aa.ooorg//baaaaaaaa.o..aa. 2407.14679.
OpenAl. OpenAl: Introducing ChatGPT, 2022. URL openai.ee.......oo. com http://eeaaaa..cooooooooo//////e
14

---

Published as a conference paper at ICLR 2025
OpenAl. GPT-4 technical report, 2023. URL URLhtts:://cdn.opp http://ppp.aaa.ccc.pa//p///pppap pdf.
Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Ngoc-Quan Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. The lambada dataset: Word prediction requiring a broad discourse context. In Proceedings of ACL, 2016. URL https:
//aclancholog. org/P 6-11--1-4 pdff.
Wonpyo Park, Dongju Kim, Yan Lu, and Minsu Cho. Relational knowledge distillation. In
CVPR http//pppes/aaaaaeeeee Baar//arra_aa/a Relational /ttp://ccp//cssse/////s/ssssssse wtheccff...... com/ content_
Proceedings of CVPR, 2019. URL https: WttLthht//ppeeeeeeeeee wwwp/pppees//eeeeeeeteappaeeastt pdf.
Disililaggpitttlliiiino _VPR
Baolin Peng, Chunyuan Li, Pengcheng He, Michel Galley, Jianfeng Gao. Instruction tuning with GPT-4. arXiv preprint arXiv:2304.03277, 2023. URL https: : //tpxa/xivv//////////////////aia wrg/abs//3330 03277.
Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. ZeRO: Memory optimizations toward training trillion parameter models. In Proceedings of SC20, 2020. URL https:// arxiv.
02054. pdf.
Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of dif- fusion models. In ICLR, 2022. URL https:///ppeeevieet 3830857a62809988
Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. DistilBERT, distilled version of bert: smaller, faster, cheaper and lighter. arXiv preprint arXiv:1910.00 108, 2019. URL https: http::s://axxi.orr///ppoo
Maarten Sap, Hannah Rashkin, Derek Chen, Ronan Le Bras, and Yejin Choi. Social IQa: Commonsense reasoning about social interactions. In Proceedings of EMNLP, 2019. URL
https://acllathht
org/D19-143
Nikhil Sardana, Jacob Portes, Sasha Doubov, and Jonathan Frankle. Beyond chinchilla-optimall Accounting for inference in language model scaling laws. In Proceedings of ICML, 2024. URL
pettss:/copenreeee
net/ forum? deeerummdd-bbdiie odO
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov, Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347, 2017. URL https: httttps://aaxi..oo/ pdf/1 06340.6348. pdf.
Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In Proceedings of ICML, 2018. URL https::/proceed:: mlr. press/v80/
shazeer18al www.e//aaaaaaaaaaaaaaaaaaaaaaaea
Ilia Shumailov, Zakhar Shumaylov, Yiren Zhao, Nicolas Papernot, Ross Anderson, and Yarin Gal. Ai models collapse when trained on recursively generated data. Nature, 63 (31(8022)7555555, 2024.
URL https:tt////mm /www///wwww/... nature........ wtweturee.ooeee... http//aricclsseeee 022-51888---2-----------
Nisan Stiennon, Long Ouyang, Jeffrey Wu, Daniel Ziegler, Ryan Lowe, Chelsea Voss, Alec Radford, Dario Amodei, and Paul Christiann. Learning to summarize with human feedback. In Proceedings of NewrIPS, 2020. URL https: /ttttpp://proceeding neurips cc/ cc/paper/2020 file/
Siqi Sun, Yu Cheng, Zhe Gan, and Jingjing Liu. Patient knowledge distillation for BERT model com- pression. In Proceedings EMNLP, 2019. URL
Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma. Yuqing Xia, Jilong Xue, Jianyong Wang, and Furu Wei. Retentive network: A successor to transformer for large language models. arXiv preprint arXiv:2307.0866/1 2023. URL https: a/aaxi/aavv wwt.vvoorrgab/oooeeooooo.o 2307 00011007........ 68621.
15

---

Published as a conference paper at ICLR 2025
Yi Tay, Vinh Q Tran, Sebastian Ruder, Jai Gupta, Hyung Won Chung, Dara Bahri, Zhen Qin, Simon Baumgartner, Cong Yu, and Donald Metzler. Charformer: Fast character transformers via gradient- based subword tokenization. In Proceedings of ICLR, 2022. URL https://openzeeveeee bet//fouum 2 i dd= BRnr 1 IEFN.
Yi Tay, Mostafa Dehghani, Vinh Q Tran, Xavier Garcia, Jason Wei, Xuezhi Wang, Hyung Won Chung, Dara Bahri, Tal Schuster, Steven Zheng, et al. U12: Unifying language learning paradigms. In Proceedings of ICLR, 2023. https httpas:///ppeneee //op/pppeeviiieeee wwtp::.weteedddddeeeeeee 6ruVLB 7 7MC.
Gemini Team, Rohan Anil, Sebastian Borgeaud, Yonghui Wu, Jean- Jean-Baptiste Alayrac, Jiahui Yu, Radu Soricut, Johan Schalkwyk, Andrew M Dai, Anja Hauth, et al. Gemini: family of highly capable multimodal models. arXiv preprint 1r2iv:2312./1805, 2023. URL https: //te////rxiii
org/abs/ 2212.311222 11805.
Gemma Team, Thomas Mesnard, Cassidy Hardin, Robert Dadashi, Surya Bhupatiraju, Shreya Pathak, Laurent Sifre, Morgane Rivière, Mihir Sanjay Kale, Juliette Love, al. Gemma: Open models based on gemini research and technology. arXiv preprint 2xXi::443300295, 2024. URL https:
/ arxiv. org/ alss 2403. 08295.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurelien Rodriguez, Armand Joulin, Edouard Grave, and Guillaume Lample. LLaMA: Open and efficient foundation language models. arXiv preprint (rXXi:2302.139771 2023a. URL https: /ttpiv...vvv///////////////..a org http/dd//2300 13971. pdf.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.nm 09288, 2023b. URL htt.t::/tp//ttttt. org/ /ars/ab//aa 2307.000 09288.
Pablo Villalobos, Jaime Sevilla, Lennart Heim, Tamay Besiroglu. Marius Hobbhahn, and Anson Ho. Will we run out of data? an analysis of the limits of scaling datasets in machine learning. arXiv preprint arXiv:2211. 010222111.4322. 04325, 2022. URL https: : //tp/aaavvv////////////////..av. www.wiiww.g.....aao.ooooooooooot abs///bs//// 2211.2211122 04325.
Wenhui Wang, Furu Wei, Li Dong, Hangbo Bao, Nan Yang. and Ming Zhou MiniLM: Deep self- attention distillation for task- compression of pre-trained transformers. In Proceedings of NeurlPS, 2020. URL https: httttp:://pooceediggg neurips
wwteriis..ccpppppee wwtpppppp///////////// file// Paper. pdf.
Wenhui Wang, Hangbo Bao, Shaohan Huang, Li Dong, and Furu Wei. Minil Multi-head self-attention relation distillation for compressing pretrained transformers. In Findings of ACL, 2021. URL https: http///aaaaanttotoeee/ee org/ 2221
188.
Yuqiao Wen, Zichao Li, Wenyu Du, and Lili Mou. f-divergence minimization for sequence-level knowledge distillation. In Proceedings of ACL, 2023. URL https: htttt::://aacaaathtaa a00 long. 605.. .df.
org/
Taiqiang Wu, Chaofan Tao, Jiahao Wang. Zhe Zhao, and Ngai Wong. Rethinking kullback-leibler divergence in knowledge distillation for large language models, arXiv preprint 2rXiv:2404.02657 2024. URL httph:t////o /ttp://aaavvvv/...///i//e.....aa /rg/abs/////eeaaooea/////////bar 2404 02657.
Mengzhou Xia, Tianyu Gao, Zhiyuan Zeng, and Danqi Chen. Sheared Ilama: Accelerating language model pre-training via structured pruning. In Proceedings of ICLR, 2024. URL https:/ ///
openreview httprreeewwwweeeeeeeeeeeee http:/wwwwweee//aiidie...... 77h 77
Sang Michael Xie, Hieu Pham, Xuanyi Dong, Nan Du, Hanxiao Liu, Yifeng Lu, Percy S Liang, Quoc V Le, Tengyu Ma, and Adams Wei Yu. DoReMi: Optimizing data mixtures speeds up language model pretraining. In Proceedings of NewrlPS, 2024. URL https: /ttttts://opereeeeee /ppp///eeeeeeeeeeeeeeeeee net/forum?in fid- 1XuByUeHhd.
Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tieyan Liu. On layer normalization in the transformer archi- tecture. In Proceedings of ICML, 2020. URL https: http::::////peeeeee..... ffrtt//iiimd
R1x8anVFPr.
16

---

Published as a conference paper at ICLR 2025
Xiaohan Xu, Ming Li, Chongyang Tao, Tao Shen, Reynold Cheng, Jinyang Li, Can Xu, Dacheng Tao, and Tianyi Zhou. A survey on knowledge distillation of large language models. arXiv preprint arXiv:2402.13116, 13116, 2024. URL https: wttpaaawwww..aaaaaaaooooooooooee org/ a 220.../116...16
Jiasheng Ye, Peiju Liu, Tianxiang Sun, Yunhua Zhou, Jun Zhan, and Xipeng Qiu. Data mixing laws: Optimizing data mixtures by predicting language modeling performance. arXiv preprint arXiv:2403.16952, 2024. URL https: / aaw.aiiiiii..oo org/ abs/ 16952.
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can a machine really finish your sentence? In Proceedings of ACL, 2019. URL https: htap//t//ocla
org/P19-144 14722/
Peiyuan Zhang, Guangtao Zeng, Tianduo Wang, and Wei Lu. Tinyllama: An open-source small language model. arXiv preprint 24X0:2401..02885, 2024. https:t:://aaa 2401. 02385.
17

---

Published as a conference paper at ICLR 2025
A PROOF OF PROPOSITION 2.1
To prove Proposition 2.1, We start from the K case, corresponding to selecting an instance with the maximal reward and Eq. (3) becomes
P arg max r(X,)) arg max ISREN
r(Y ((m)) 1,
(6)
IsmSM
which is equivalent to
EP(R
arg max r(X,) and arg max r(Y m) x +1 ISMSM
(7)
ISASN
mES
Since X, and Y, m are independent random variables for S N and I m S M. Eq. (7) can be further written as
Σ
P
arg max ISnSN
r(Xn) x P arg max r r(Ym) x 1.
We focus on the term P arg max r(X,) which can be expanded as follows based on the ISnSN
fact that 11, 22, XN are i.i.d random variables:
P arg max rX,,,) ISREN
-P(((X,) r(w)) for 5 N and at least one variable in X1, X2, X equals a:)
(r(((X,.) S r(x), for 1 S n S N) - (r,(X,,, & r(x) and X, + x, for I n I N)
N
N
= III P(r(X,) rr())) II (x)( P(X, x]]
(9)
==1
=P(r(X,) S <r(z)) (P(r(XX) S (x)=))
1-()a (1()))
-P)(x)) N
1 1- 11(1- (x(
where (x(x)= P(r(Xi) S rax)). Note that < (1x)
= P(X = x) P(r(Xi) S r(x)))
(x(x), have 0 (=)
[x < 1. and thus
lim N
1 (1(m)
0. Similarly, by setting
[x
Px P(r(Y)) S r( r(x)), we have
p2(x) M
arg
max r(Y (Ym) = x (r)
(10)
ISm$M M
Px
M
and lim
p2( Py( (z)
= Let x= = arg max r(x) Therefore, Px(x") = (xx) = 1 and #ES for x + x. When N + +oo, we have Px(x*)N Px
M-+-0
Px (x) 1, (x
and Px ) Py x
MMM + for x + x.. Therefore, we have
lim Σ P
arg max r r(x)) x P arg max r r(Y ISnSN
N,+++00
Isms
#ES
Pi (x))
P2 (2(x)
E lim
Px (r Px(r))PP(()"'
lim
1
N-M
Px (a)
Pr(x))
#ES
lim
Σ
lim
Px (xx) Py (x)
NMM++00
N.M
2ES
=1,
(11)
18

---

Published as a conference paper at ICLR 2025
which proves Eq. (6). For K > 1, the equality of two top- subsets requires the elements ranked from to K to be equal, respectively. Therefore, the left hand of Eq. (3) can be decomposed using
Bayes's 's Law. Let X' X2, .x), y {Y,, Y2. Y Y.). b.r.X. X arg max r (X.n). IgngN
and Y = =agg max r(Ym). we have P(X" = Y*) + (Bq. (6)) and ISMSM
P (top-2 [r(X.)) S Sn N = top-2 (r(Ym)) S S M))
(12)
=P
arg
max (X-)
r(X) arg max YEN-(Y-)
r(Y) x = P(x-=Y)
Since X1, X XN and Y1, Y2, "..... Y M are i.i.d. random variables, we can still decompose
the term P arg max r(X) a x"= Y*) as Eq. (9), which means when N, M - +o0: (xx-
P arg max r(X) = arg max ((Y) x = YEN-
+1 )
(13)
(XX-(X-)
Eq. (13) means the elements with the secondary large r(æ) values are highly likely be the same. The decomposition in Eq. (12) can be conducted K times for the elements ranked from I to K, and each decomposed term approaches 1 similar Eq. (13). So far, we have proved that the probability of X' and Y having the same top- K subsets measured by r(æ:) approaches to I when N, M ~> +00, which is formally written as
P (top-K K [r(X,)] I n N) top-K (r(Ym)) 1 $ m S Mt) - 1.
(14)
This completes the proof of Proposition 2.1.
B MORE EXPERIMENTAL DETAILS
Model and Training Configurations. We mostly follow Brown et al. (2020) 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 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 Adam (Loshchilov & Hutter, 2019) optimizer, with 31 = 0.9, 82 0.98, and 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 2N D and 4N D. 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.
Configurations of KD Across Model Families The model configurations of LLaMA3.1 (Dubey et al., 2024) and Mamba (Gu Dao, 2023) for the cross- family distillation experiments are listed in Table 6. The pre-training data and optimization paprr-parameters are the same as the configurations in our main experiments on Qwen, as listed above.
Evaluation Details. Our downstream datasets for evaluation include Hellaswag (HS; Zellers
et al., 2019), LAMBADA (LAM; Paperno et al., 2016), Winograde (Wino; Levesque et al., 2012),
OpenbookQA (OBQA; Mihaylov et al., 2018 88, ARC-Easy/Challange (ARC-e/c; Clark et al., 2018), PIQA (Bisk et al., 2020), SIQA (Sap et al., 2019), and StoryCloze (Story; Mostafazadeh et al., 2016). We upply the sample 10K documents from the DCLM (Li et al., 2024a) corpus to construct our test set for language modeling evaluation.
(Gao et al., 2024) framework to conduct zero-shot evaluation. We
19

---

Published as a conference paper at ICLR 2025
Family
Size
Vocab,
dmodel
dFFN
faayrrss
TChcad
dhead
learning rate
104M
151,936
512
1,408
8
8
64
6 x 10
200M
131,936
768
2,112
12
12
64
x 10 -4
Qwen
300M
151,936
768
2,112
18
12
64
6 10
500M
151,936
1,024
2,816
24
16
64
3 x 10
1.2B
151,936
1,536
4,224
24
16
96
2.5 x 10
LLaMA3.1
212M
128,000
768
3,072
12
12
64
6 X 10
Vocab.
dmodel
dFFN
YMaayer
conv
rank
learning rate
Mamba
140M
50.280
768
1,536
24
48
10
Table 6: Model configurations and corresponding learning rates.
Vanilla KD
MiniLLM
Formula
3 T
3N 4NNo
T
3.Neu + Nuc
22N
Student Model Size Nau
200M
500M
1.2B
200M
500M
1.2B
Training Steps
25K
45K
65K
15K
30K
40K
Table 7: Training steps in Vanilla KD and MiniLLM, which is set to ensure training-time computation to be the same Pre-Train w/o KD. Nsm and Nich are model sizes of student and teacher LMs respectively. Nich = 1.8B in our experiments. T 100K is the training steps in Pre-Train w/o KD.
C DISCUSSION ON RESOURCE CONSUMPTION
In this section, we discuss the asymptotic complexity, actual runtime, and memory usage different methods. Table C summarizes the offline computation requirements for Difference Sampling, which only need to be performed once for any LMS be distilled from specific teacher L.M. Table C provides the training-time asymptotic complexity, actual runtime, and memory usage for various pre-training methods. For these results, we assume the LMs uses ZeRO-2 (Rajbhandari et al., 2020) and FP16 precision during training to calculate space complexity. The "Actual Time" and "Max Batch Size" metrics are derived from experiments conducted on A100 80GB GPUs. We can see that the inference time of the teacher and reference LMs constitutes the majority of the runtime for Difference Sampling. This cost can be further reduced by leveraging a proxy model for sampling, detailed in Appendix E. Regarding training-time resource consumption, MINIPLM requires the same time and memory as pre-training without KD), whereas Vanilla KD incurs significantly higher runtime and GPU memory usage due the online inference of the teacher LM.
D MORE RESULTS
D.1 COMPUTATION SCALING CURVES MORE SIZES
In Figure 7, we plot the scaling curves of average accuracy 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.
D.2 TEST LOSSES EXTRAPOLATION WITH SCALING LAWS
Pata-Unlimited Setting. In Table 2, we follow Hoffmann et al. (2022) to fit the scaling law curves with the test losses the DCLM corpus. Then, we used the fitted constants to predict the test losses for IT and 10T training data Pre- TreiTrai w/o KD and MINIPLM, and that for Vanilla KD when
co/gee
20

---

Published as a conference paper at ICLR 2025
Time Complexity
Space Complexity
Actual Time
Max Batch Size
Reference LM Training
(6 (2.Nich D)
((3.75Nra)
2.6h
32
Teacher LM Inference
((2N
68h
16
Reference LM Inference
0 D2 (D)
O(2NNa)
7h
128
Top-k Sampling
10min
Table Asymptotic complexity, actual time and memory use of Difference Sampling. Near and Nich are the reference and teacher model sizes. D and Dier are the sizes of the original corpus and the corpus for training the reference model. "Actual Time" the running time of each stage. "Max Batch Size" represents the max single GPU batch size, which measures the memory use of each stage.
Time Complexity
Space Complexity
Actual Time
Max Batch Size
Pre- Train Vanilla KD
(6 D)
033.75N
68h
16
(6 D + 2N
03.75.N + 2Na))
139h
8
MiniPLM
033.75N
68h
16
Table 9: Training- asymptotic complexity, actual time and memory use of different pre-training methods. Nsw and Nich are the student and teacher model sizes. D is the size of the pre-training corpus. "Actual Time" is the training time of each method. "Max Batch Size" represents the max single GPU batch size, which measures the memory of each method. Unlike the setting in Section 3.1, we align the total training steps (token numbers D) of the three methods to compare their training-time complexities.
consumes the same training FLOPs. Specifically, we consider the following power law:
Am As NO
L(N Na) Lun +
(15)
N
where Nn the model sizes, Na is the number of trained tokens, and Lan 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' x Nm Na for fixed model size:
Ac
L(C) = L. +
(16) fit the loss
Co.
where Loe and A depends the model size Nn, and 02 = Odd. In this way, curves in Figure with Eq. (16). Table 10 includes the fitted values of Los Ac, and ag. We also include the total FLOPs of Pre-Train w/o KD and MINIPLM which Vanilla KD aligns with. The numbers in Table 10 can be computed with Eq. (16) and the constants in Table 10.
1T (Cit) and 10T (Cior) tokens,
Data-Limited Setting. In Section 3.4, we extrapolate the loss curve of Pre-Train w/o KD with the Cata-Constrained Scaling Law (Muennighoff al., 2023), which is almost the same as Eq. (15), except that N₂ represents the total number of tokens in the pre-training corpus repeated 4 times, as suggested by Muennighoff al. (2023). Therefore, after Na is solved by letting the loss of Pre-Train w/o KD equal that of MINIPLM (a = 0.25, 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.
D.3 ABLATION STUDY
Impact the Reference Model. Figure 8, we show the impact of the reference model size on the performance of LMs trained with MINIPLM. We 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-sampll data contains 50B tokens, we need a larger original corpus D. To control the computation overhead, mainly use 0.5 in experiments.
21

---

Published as a conference paper at ICLR 2025
45
41-
(89
46
Average Acuraty
43
Accurscy
44
39
42
41
42
37
Vanils KD
Po Trai
40
39
Pee Tan Vanila
Vandaa
PIM (Oun))
38
(
1Ouns)
35
1.0
2.5 3.0
Training Computation (x10" FLOPs)
Training Computation (x10 FLOPs)
Training Computation (x10* FLOPS]
(a) 1.8B + 200M
(b) 1.8B - 500M
(c) 1.8B + 1.2B
Figure 7: Computation scaling curves of student LMs when trained with Pre-Train w/o KD, Vanilla KD, and MINIPLM. The y-axis represents the average zero-shot accuracy on downstream tasks.
Nau
Method
Aa
Lio
Cir (FLOPs)
Cior (FLOPs)
Pre-Traii w/o KD
2.19x107
0.41
3.30
200M
Vanilla KD MINIPLM
9.77x10"
0.44
3.34
1.26x 10
1.26x10
8.56x10'0
0.59
3,25
Pre- Pre-Train w/o KD
20
0.45
3.06
500M
Vanilla KD
3.14x10*
0.45
3.05
30
3.14x10
MINIPLM
6.64x10
0.52
3.03
Pre-Trrai w/o KD
108
0.43
2.91
1.2B
Vanilla KD
1.10x10
0.52
2.90
2.
x.30x
MINIPLM
x.29x10
0.45
2.86
Table Scaling Law constants Eq. (16) fitted using the loss curves in Figure 4. Naw means the student LM size. CIT and Clor are the compute spent on processing ITT and 10T tokens in Pre-Train w/o KD and MINIPLM, which Vanilla KD aligns with.
D.4 CORRELATIONS BETWEEN log p(x) AND log
p(a) gu(z)
To examine how replacing go with prer in Difference Sampling works, we compute the correlations be-
Nist
Pearson
Spearman
Sampling Acc.
tween log p(z) (using 45M, 104M, and 200M LMs
(x))
45M
0.743
0.774
80.6
to generate Pref) and log p(a)
go(z) (using the 500M stu-
104M
0.826
0.840
87.4
200M
0.856
0.879
89.5
dent LM to generate qa) on a holdout set. Addi- tionally, we calculate the sampling accuracy, which measures how many instances are correctly classi- p(a) fied as selected or discarded when using log Par( and setting the sampling threshold 0.5, with the
Table 11: Correlations between log phe[aa) and
p(.)
[z]) .Nrer represents the size of the refer- model. qo is generated from the 500M student LM. We report the Pearson/Spearman Correlation and the sampling accuracy (Sam-
log p(a)
selection using log [u(z] serving as the ground truth. As shown in Table D.3, using Peef achieves strong correlations with using qu and high data sampling accuracy. This supports the rationale for using Pref in Difference Sampling as described in Section 2.2.
p(a)
pling Acc.) using log Par(z))
p(z).
E 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 Ppe,
and then transferring the value r(p, Pret, x)) to the entire corpus D with a proxy model. Specifically,
we first uniformly sample proxy subset Dpr from D, satisfying [Daxx << (D). Then, we compute
the r(p, Pret, æ) value for each instance in Dpxx- Note that since pDpm << DD], the computational
22

---

Published as a conference paper at ICLR 2025
Average Accuracy (%)
Average Accuracy (%)
41.5
MiniPLM
Vanilla KD
41.0
MiniPLM
41.0
40.5
Vanilla KD
40.5
40.0
40.0
45M
104M
200M
0.4
0.6
0.8
Size of the Reference Model
Figure 8: Impact the reference model size. We 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 that of Vanilla KD.
Figure 9: Impact of the difference sam- pling ratio We report the average zero- shot accuracy on the downstream tasks of the LMs trained with MINIPLM, using [0.3, 0.4, 0.5, 0.6, 0.7, 0.9] and compare it with that of Vanilla KD.
overhead of teacher LM inference is significantly reduced. After that, we fine-tune small proxy model on D pex to fit the r(p, Prer, x) values, which used to infer the values for instances from D. Finally, the Top-K operation in Eq. (4) is based on the inferred values.
Method. To test this approach, uniformly sample Ppxx containing O.IB tokens from the p(z) values on Pprx takes only 0.2% computation that on D. Then, we employ the reference LM as the proxy model and fine-tune with the following regression loss:
50B-token D. Computing log Par(#)
w", w",b", O = arg min 0.
Σ mED
[w Tilx, O b r(p, Pref, x)}²
(17)
D pex
where h(x, 0ret) Rd is the average output hidden states of the reference LM, with d representing the model's hidden size and O representing the parameters of the reference L.M.
Method
FLOPs
Acc.
Vanilla KD
Online
39.9
w E Fd, E IR are the parameters of a linear head outputting scalar as the predicted r(p, Prets x) values, given the average hidden states. The inferred values on 'D are given by f(x:) w" law 0 +-b" Since this inference process is based on the reference LM, still saves computation compared inference with the teacher LM. The sffference-samplee corpus is then obtained by selecting K = a|Di instances from D with the highest f(x) values.
MINIPLM
2 10
41.3
MINIPLMprr
9 x 1
40.9
Table 12: Offline FLOPs and aver- age accuracy (Acc.) on downstream tasks of MINIPLM using proxy model, compared with that of stan- dard MINIPLM and Vanilla KD.
Results. We term this method MINIPLM and compare its performance with Vanilla KD and MINIPLM in Table 12. We can see that MINIPI. requires much less inference computation compared to standard MINIPLM while still maintaining substantial improvement over Vanilla KD.
F CASE STUDY
We present a case study in Table 13, 14, and 15 to show the instances corresponding to the three parts of data distribution shown in Figure 3(b). The explanation of the three parts list follows:
p(=) Z (r))(:): As shown in Table 13, the instances whose log par(#) value is close to 0 is discarded (Instance #1 and #2). These instances full of repeated patterns, which is easily learned by both the small and large LMs. Instance #3 a dialog from a story, where the teacher LM learns slightly well than the student LM, which is retained for the student L.M to
p(m)
23

---

Published as a conference paper at ICLR 2025
learn basic language skills. In this way, these easy and common patterns are down-sampled but not fully discarded.
p(æ) > Prer(z)): As shown Table 14, the examples where p(æ) much larger than Prer( (z) is retained. These examples include long and showledge-intensife documents (Instance #1), in-context learning examples (Instance #2), and high-quality codes (Instance #3). These examples reveals the key large LM's advantage the small LM, which should up-sampled for the student LM to learn.
p(æ) S (r)( As shown in Table 15, the instances where the larger LM learns worse than the small LM are all discarded. We can see that Instance #1 and Instance #2 are irregular noises, which will be harmful to the training of LMs. The small reference LM easily over-fits their surface patterns like symbols and numbers, while more powerful large LM better recognizes their unnatural features. Instance #3 is a data collection an academic paper, but lack of contexts, which is useless for LMs to learn the dependency in contexts. Note that these examples only accounts for 6.3% of the pre-training corpus before Difference Sampling, which quite small proportion.
24

---

Published as a conference paper at ICLR 2025
p(œ) 2 (r) Easy and common instances
Instance #1
-lgg((>) 1.24
Pa) = 1.28 log Par(m)
p(x)
0.04
Discarded
<node id=' -659' action='modife visible='irue" 1at=
'0 '0.05-6706224879555"" 1. 0101
/>
<node id=' action='modify" visible=ftrue' 1at=
' 00004408822551111111100002222211
/>
<node id=' action wtt.:l/eyyy.isillee==tttteeee lat=
r0. "1.08080339981111000003111111111
/>
<node id='-653' mctionn ortptoornmnddiiyyisiislleeeooorn tre lat=
10. 1r.088155h21100gg11oooo044455555
/>
<node id=' action= visible='true' lat=
00.04733884295596688 1on= 00t00:.0000111991101000000000000
/>
Instance #2
logp(æ) 0.44 Phr = 0.51
log
p(a)
0.07
Discarded
Par(#]
/ssmm/immtts_ - 1 1: 'operatorname (size) \left
\mathcal(x), \right) \cdot \operatorname {size)
Aleft ( (ammmmmmmmaahaal1))1m1a11
1} \right) \quad) (\mathcal(I) = [M] V
[Nsum)limits_(] = [) [
Vleft (
(Amathcal \right) \cdot \operatorname (size)
A /\aatcaaaaaaaainnnnnnnn j \right) \quad) &
\ AAmathaal[I] mmthhaal(Mi_(22.) \
(ksum\limits_() s Aleft
(\mathcal((),.)) \right) \cdt Veperatorname (size)
Aleft (\mathca(()))mm i} \right) \quad) & (mmathcal(T) \M) \\
[ssum\imits__1 1 2/1operatorrame (size) \left
(m)t \right) \cdot Voperatorname (size)
lleft (\mathcal 5} \right) \quad) &
{ (Imathcal(I) = (M)
\
Instance #3
-lgg((:) = 2.83
(og = 3.86 log p()) Par(#)
1.02
Selected
"I felt bad for Coach Rod to have to deal with it, said senior tight end Mike Massey, St. Ignatius grad. "But for the players it was non-issue. We didn't even talk about it."
But some other people college football did.
"[Boren] was a legacy guy and when he makes comments like [that], that's like getting kicked square in the shorts when you're re Rich Rodriguez, ESPN analyst and former OSU quarterback Kirk Herbstreit said Thursday. "With all the other things that are happening, when you get that comment from legacy guy, guy whose dad started three or four years for Ro, all of a sudden everybody's s tentacles go up little bit.
Table 13: Easy and common instances down-sampled by Difference Sampling. Instance #1 and #2: HTML and LaTeX code data that contain repeated patterns, which easy to fit by both the reference and teacher LM. Instance #3: Dialogues from story, which is relative easy for LMs but still selected by Difference Sampling to help student LMs learn basic language skills.
25

---

Published as a conference paper at ICLR 2025
p(x:) >> (r): Hard and valuable instances
Instance
logp(()) 1.26 -oogPe(())) = 4.20
log p(w) 2.94
Selected
par(
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 sisposals inside to be carted away. Nonetheless, there are legal issues attached to aporopriate disposal connected with certain products which tie up into environmental issues. The 10 Yard Dumpster For Rent in Pocahentas customer perhaps demolition purchaser should be informed about these issues by means of careful screening so as to reduce firm's s liability which inturn keeps firm's s overhead all the way down and makes for prompt fall off, pick up along with disposal of the dumpster and it's s articles.
Instance #2
llgg((z) = 2.36
(x) = 5.59 log p(æ)
Pal 3.23
Selected
* you31i4 yöuli advantageous; beneficial
Wi jin3shen4 jinshèn cautious; prudent
F jia3 jiä one; armor (1st Heavenly Stem)
tth you2yu4 youyù hesitate; hesitant; undecided
Hut cong2ci3 cóngeī from now on; since then
Etlk qi3ye4 qiyè company; business; firm
F xia4zai3 xidză1 to download
$ shilzi5 shizi lion
# ging1shao4nian2 qingshāonián teenager
Instance #3
logg(x)) 0.16 - (z) 2.73 log
p(w) 2.56 Selected
function WritableStee (options stream) var Duplex requi ('. [''/_stremm_ddple"")>>>>
options options I1 O) // the point at which write() starts returning false
// Note: 0 is valid value, means that we always return false
// the entire buffer is not flushed immediately on write()
var hwm options highMaterMark; var defaultHwm options. objeetModee16:: 16x11114;; this WiiishiihWWtteMar = (hwm hwm === 0) hwm HefaultHwm;
// object stream flag to indicate whether or not this stream
// contains buffers or objects.
this objectMode !!options. objtMoode;
Table 14: 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 an in-context learning form. Instance #3: High-quality code data with detailed comments.
26

---

Published as a conference paper at ICLR 2025
p(x) S (a)): Noisy and harmful instances
Instance #1
p(z) = 9.50
log Pref (a) = 6.60
p(æ) log Pir((n) -290
Discarded
Instance #2
-lgg(()) 1.01
log (x) 0.90
log p(x)
0.11
Discarded
Par(
(91.00,60.00) (1.00, 35.
())
(16.70, [,,.:,.88.000..1..00000000000110
(25.0,,55.00 (20,0,000000 (26.00,12.00)
[12.00]
(26.0,)))))))) (26.000,,, (16.70,35.70
115::0088..011111111111110111101 (56.00,35.00)
(66.00,25.00) (76.00,35.00) (66.00,45.00), (66.00,46.0000
Instance #3
logp(x) 2.53
Pod Poe (2: 0.26
log p(x)
2226
Discarded
(ar(())
* 8
Data collection stableewrapattaoolleeet...ooo
Bruker APEXII area aree-detector dif iiffractimmterr
Radiation source: fine-focus sealed tube graphite
p and scans
Absorption correction: multi-scan (*SADABS*; Sheldrick, 2004) * 0.970, *T*max= 0.981 20408 measured reflections
Table 15: Noisy and harmful instances discarded by Difference Sampling. Instance #1: irregular symbols. Both the reference LM and the teacher LM are hard fit. Instance #2: 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 scientific paper, but lack of contexts. The reference LM fits the patterns, which the teacher LM finds useless.
27


Reviews:
[{"review_id": "TIArboBBcv", "reviewer": "Reviewer_ZWq4", "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.", "questions": "The explanation of Difference Sampling and its theoretical justification could be simplified or supplemented with more intuitive examples for clarity.\n\nThe paper should discuss in detail the limitations of the proposed method.\n\nThe transferability and generalizability of MINIPLM across different language models and types of downstream tasks should be discussed.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 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.\n\nThe 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.\n\nThe results on computation reduction and performance scalability (e.g., 2.2x computation reduction) validate the practicality of MINIPLM for large-scale pre-training.", "weaknesses": "The explanation of Difference Sampling and its theoretical justification could be simplified or supplemented with more intuitive examples for clarity.\n\nThe 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.\n\nThe 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.\n\nWhile 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "pwqeJtO0e0", "reviewer": "Reviewer_Yqhf", "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.", "questions": "Please refer to weakness.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- This paper provides decent theoretical foundations \n- The contribution and experiments designed for this paper is very clear. \n- Evaluations demonstrate improved performance and data efficiency gained from the method. \n- The contribution of this paper is important to scenarios when data is insufficient.", "weaknesses": "Overall, I don't have many criticism toward this paper. \n- 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.\n- The complexity of the training process may impose challenges in reproducibility and implementation.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Q7P4e4Af76", "reviewer": "Reviewer_25K6", "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.", "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?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- This paper Introduces \"Difference Sampling\" as a novel offline KD method that fundamentally changes how knowledge is transferred.\n- This paper re-computes teacher model probabilities just once, requiring only 200MB storage for 50B tokens vs 30PB for traditional methods.\n- This paper shows clear computational efficiency gains (2.2x reduction for equivalent performance).\n- This paper provides solid theoretical justification through Proposition 2.1 and detailed analysis of the sampling approach", "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. \n- 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.\n-  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.\n- 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. \n\n[MobileLLM] MobileLLM: Optimizing Sub-billion Parameter Language Models for On-Device Use Cases", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "j9bstCDXh7", "reviewer": "Reviewer_Rr1c", "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.", "questions": "- Can you provide theoretical justification for using log(p(x)/p_ref(x)) as the sampling metric?\n- How does the method's effectiveness scale with model size differences?\n- How would the method perform with larger teacher models (>10B)?\n- What modifications are needed for very large scale deployment?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- Well-motivated solution approach\n- Elegant Difference Sampling design\n- Efficient offline computation strategy\n- Comprehensive experiments across model scales", "weaknesses": "- Limited analysis of finite-sample scenarios\n- Teacher model size limited to 1.8B\n- Limited cross-family experiments\n- Limited exploration of alternative sampling strategies\n- Reference model selection criteria not fully justified", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "u81QJWYxFj", "reviewer": "Reviewer_rC3o", "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.", "questions": "See weaknesses.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "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.\n\n2. The paper is well written and easy to follow.\n\n3. The paper has good experimental results as compared to the existing methods.\n\n4. The problem taken up is an important problem in pre-training the language models.", "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.\n\n2. 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.\n\n3. 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?\n\n4. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)