Distilled Pretraining: A modern lens of Data, In-Context Learning and Test-Time Scaling
Authors: Sachin Goyal
David Lopez-Paz
Kartik Ahuja
Venue: ICLR 2026 Poster (2026)
Keywords: Pretraining
Primary area: foundation or frontier models, including LLMs
TL;DR: We study LLM pretraining with logit distillation and show that it comes with a tradeoff: it boosts test-time scaling but impairs in-context learning. We analyze this tradeoff in detail and develop mitigation strategies.

Abstract: In the past year, distillation has seen a renewed prominence in large language model (LLM) pretraining,
exemplified by the Llama-3.2 and Gemma model families. While distillation has historically been
shown to improve statistical modeling, its effects on new paradigms key to modern LLMs—such as
test-time scaling and in-context learning—remain underexplored. In this work, we make three main
contributions. First, we show that pretraining with distillation yields models that exhibit remarkably
better test-time scaling. Second, we observe that this benefit comes with a trade-off: distillation
impairs in-context learning capabilities, particularly the one modeled via induction heads. Third, to
demystify these findings, we study distilled pretraining in a sandbox of a bigram model, which helps
us isolate the common principal factor behind our observations. Finally, using these insights, we shed
light on various design choices for pretraining that should help practitioners going forward.

Published as a conference paper at ICLR 2026
DISTILLED PRETRAINING: A MODERN LENS OF DATA, IN-CONTEXT LEARNING AND TEST-TIME SCALING
Sachin Goyal"
David Lopez-Paz
Kartik Ahuja FAIR at Meta ka rt irttikkauujeee kaaauaaametaaaa... com
FAIR at Meta
FAIR at Meta
Carnegie Mellon University
dl pêmeta. com
sach. goya sach i n@ com
ABSTRACT
In the past year, distillation has seen renewed prominence in large language model (LLM) pretraining, exemplified by the Llama-3.2 and Gemma model families. While distillation has historically been shown to improve statistical modeling, its effects on new paradigms key to modern LL.Ms-uuch as test-time scaling and in-context leanning-remain underexplored. In this work, we make three main contributions. First, we show that pretraining with distillation yields models that exhibit remarkably better test-time scaling. Second, we observe that this benefit comes with trade-off: distillation impairs in-context learning capabilities, particularly the one modeled via induction heads. Third, to demystify these findings, we study distilled pretraining in sandbox of bigram model, which helps us isolate the common principal factor behind our observations. Finally, using these insights, we shed light various design choices for pretraining that should help practitioners going forward.
INTRODUCTION
Knowledge distillation, first proposed by Buciluă et al. (2006) for compressing ensembles, was later popularized by seminal works of Ba & Caruana (2014) and Hinton et al. (2015). However, distillation didn't trickle into the pipelines of early large language models (LLMs)- such as GPT-2/3 and Llama 1/2. But more recently, distillation has resurged a prominent method in the LLM landscape, not just during post-training, but also pretraining as seen in the Llama-3.2 (Meta AI, 2024b) and Gemma (Gemma et al., 2024; 2025). This shift reflects growing reality: extremely large models (e.g., Maaaa--Behheoote (Meta AI, 2024a)) are too costly to deploy widely and will increasingly serve solely as teachers for distilling smaller, more practical models. Going forward, these deployed models are likely to be pretrained entirely via distillation as seen in -lama-4-Maverit (Meta AI, 2024a) that was distilled from 4eama---Behemoth
Despite its growing role. the science of distillation (using soft labels) in modern LLM pretraining has remained largely unexplored. Gemma-3 and Llama-3.2 models show clear empirical benefits on standard benchmarks from pretraining with distillation. However, these models typically leverage teachers trained on far more data than the students. This raises fundamental question: are the gains from distillation merely result of additional teacher data, do they reflect unique benefits beyond extra data exposure? As we hit the data wall, will distillation continue to be beneficial? Moreover, modern LLMs are no longer limited to evaluation in-context learning and test-time scaling are key to current LLM frontiers, yet the effect of pretraining with distillation on these paradigms remains largely unexamined.
standard benchmarks. New paradigms such as
In this work, we uncover key trade-offs associated with distilled pretraining (DPT). First, we show that DPT remains beneficial standard language modeling tasks, even in the data-constrained regime where the student and the teacher models are trained on the same data. This suggests promise for scaling DPT further. However, in contrast, we observe that naively scaling pretraining with distillation (DPT) hurts the in-context learning performance (Figure lb). In particular, distillation impairs the learning of induction heads (Olsson et al., 2022) the transformer circuits that enable models to search and copy from context(Figure Ic).
Work done Meta

---

Published as a conference paper at ICLR 2026
Distilled pretraining improves
Distilled pretraining burts
teet-time scaling
In-context learning
Learning Dynamics for Induction Teak
Pass@k on GSM8K
25-
Dietalatton goves diversity
36
Dirzilied Primaadag 12:)
Indorttion Accurey
Fretroniing
Protralaing
TruckerMeee Madd 1916
20
0.15
15-
Accuracy
30
0.10
Standard Postraining (SP1)
IsoData
Delilativo burte
Indiaction Swice
Distilled Pretreintng (1P7)
0.00
. Standard 1TC 2t dats (5/71223
24-
0
0.00
1258
500B
17
K (#Altemptss
Data
Training Epoch
(b)
(c)
Figure 1: Distilled pretraining in modern LLM regime (a) Comparing standard pretraining (SPT) with distilled pretraining (DPT). On reasoning tasks like GSM8k, although both the models have similar pass@1, DPT substantially outperforms SPT on pass@k for higher & (27% vs 23% for k= == 16). Infact, DPT matches the pass@ 16 of a standard pretrained model trained on twice the data. (b) Distilled pretraining hurts in-context learning capabilities when the student and teacher model see the same data. In the figure, as we scale the student data to IT (data seen by the teacher), the gains of distillation over standard pretraining on in-context learning tasks diminish (Figure 3 for details). (c) We demystify these findings by analyzing a bigram sandbox, where we show that training with distillation can impair the learning of induction heads (Bietti al., 2023), which form the key mechanism behind in-context learning.
Strikingly, the very process of distillation that undermines in-context learning, at the same time also yields models that demonstrate markedly better test-time scaling capabilities. We study this through pass@k, where the model is allowed multiple attempts per question. Distilled models outperform standard pretraining on pass&k at larger k, even when pass@1 is the same (Fig. On GSM8K, for example, both models have the same p but the distilled model achieves a much higher pass@ 16-27% versus 23%. Remarkably, even matches the pass@ 16 of a standard- preerraiddd model trained on twice the data, despite lower p Similar patterns hold on MATH and MBPP, where distilled pretraining consistently improves test-time scaling by enhancing generation diversity(Dang et al., 2025).
Interestingly, the mechanisms through which distillation undermines in-context learning are the same ones that enhance test-time scaling. We study this tradeoff in simple yet expressive sandbox of a bigram model (Bietti et al., 2023; Edelman al., 2024). A bigram model is characterized by a matrix in which each row represents the next token probability distribution over the vocabulary. Pretraining with distillation is beneficial in learning the high-entropy rows. These rows basically model prompts like "I work at",, which admit multiple valid completions (e.g.. "gym", "hospital", "restaurant"). In contrast, distillation does not help in learning low-entropy rows which model the deterministic state transitions (prompts), e.g., induction heads where the next-token probability distribution is one-hot. For these cases, distillation does not provide any information beyond what is already there in ground truth one-hot labels. Worse, an imperfect teacher can hurt the learning of these low-entropy rows by introducing noise via soft probability Eestribution(Figury Ic).
Finally, borrowing insights from our analysis, we discuss various design choices for improving pretraining in $5. These include distillationsspeiifff data curation, teacher selection, and comparisons with other recent advances such as multi-token prediction (Gloeckle et al., 2024), which we hope will aid practitioners going forward. We summarize our key contributions in this work below:
Test-time scaling: We show that distilled pretraining produces models with markedly stronger test-time scaling, often matching standard pretraining on up to twice the data.
In-context learning trade-off: We find that these gains come at a cost, as distillation impairs in-context learning, particularly by weakening induction heads.
Bigram analysis: We isolate the common mechanism that drives the improvements in test-time scaling but impairs in-context learning at the same time.
2

---

Published as a conference paper at ICLR 2026
well
Natural Questions
MBPP (Coding)
Average Accoss All Tasks
Clasification Acc.
Distilled Standard PT
47.5
Distilied FT
Match
Pass@11
Exact
02
525
1350
5300
1250
2050
1238
2018
Data
Data
Data
Data
Figure 2: IsoData Distillation 2): Will distilled pretraining remain effective when the student and teacher are trained on the same data? To explore this, we use an 8B model trained on IT tokens as teacher. Using this teacher, we train various student models, with and without distillation, scaling up the data to the exact same IT tokens. We observe that even in the IsoData case where both teacher and student have seen the same tokens, the distilled model generally outperforms standard pretraining on standard language modeling tasks. Thus distillation generally remains beneficial even in dots-constr regime. See Figure 12 for more tasks.
Practitioner Takeaways: We translate these insights into concrete design choices for improving pretraining with distillation, including distillation specific data curation, teacher selection, etc.
1.1 PRELIMINARIES
We recall the distillation with soft label objective (Hinton et al., 2015), which interpolates between fitting hard labels y, and teacher soft labels s; (hoeachhe (x:))) (Tx)/T)::
h E arg min KEH
[
o((,))((x.))) Σ (ss (((x))))
(1)
In LLM pretraining, the same objective applies next-token prediction with yı replaced by x++1- Full derivations and notation are deferred to the Appendix B.
2 No EXTRA DATA: DOES DISTILLATION STILL IMPROVE PERFORMANCE?
Recent pretrained LLM families- such as the Gemma-3 and Llama-3.2 series have shown clear benefits from distillation compared to training from scratch. However, these models typically leverage teachers trained significantly more data than the students ultimately question: Are the gains from distillation simply due to this additional teacher data? Does distillation offer unique benefits beyond merely seeing extra data via the teacher in the modern web-data regime?
raising fundamental
We begin this work by answering the basic question raised above via a set "IsoData Distillation" experiments. We first train 8B teacher model on 1T tokens. We then train 1B students- and without distillation
the same IT tokens to test if distillation still helps when both see identical data. Figures and 12 compare the performance of the two 1B models standard language modeling tasks like COPA, HellaSwag, NaturalQA, TQA, GSM8K, etc. We observe that distillation continues benefit even when training is scaled to the same data the teacher (IT tokens).
Distilled pretraining (DPT) continues to be generally beneficial even in the data-constrained regime, when the student shown the same amount data as the teacher.
Theoretically, Mobahi et al. (2020); Nagarajan et al. (2024) study self-distillation and understand why it improves performance, which is kind a IsoData distillation. We refer the reader to Appendix H for detailed discussion on theoretical works in IsoData distillation. In the next section, will analyze distilled pretraining on new paradigms centric to modern LLMs, beyond the standard language modeling tasks: in-context learning and test-time scaling.
3

---

Published as a conference paper at ICLR 2026
Context: based QA
Neodle Haystack
Counterfactual Context QA
Average Across ICL Tesks
43-
$7.5
Accurscy
35
Teacber Midel 2381
Ziuane I tokoss)
Accuracy (%)
52
Accuracy
92
Accuracy (%6
12.7
50.3
IsoData
Data
24
IsoData
32-
75
IsoData
1350
536n Data
1250
53tn Data
1258
5018 Data
12
1238
2008 Dala
(a)
(b)
(c)
(d)
Figure 3: Distilled pretraining impairs in-confext learning, especially in the IsoData setting
($ 3.1): We train IB models with and without distillation using 8B teacher trained on IT tokens. We observe that the advantages of distillation on in-context learning tasks diminish as the amount of training data increases (each scatter is a separate model trained with a full LR scheduler). Eventually, the distilled model underperforms in the IsoData setup. where both the teacher and student are trained on the same data. This is because induction heads which form key mechanism behind in-context learning (Olsson et al., 2022) are built on low-entropy mappings, requiring the model to copy a specific token from earlier in the sequence. For these cases, distillation can't 't help-it it can only match the hard label at best, and at worst, it actively hinders learning for such copying tasks by softening the supervision. This is in contrast to performance on standard language modeling tasks where distillation continues to help even in Ista setting (Figure 2).
DISTILLED PRETRAINING THROUGH THE MODERN LENS: IN-CONTEXT LEARNING AND TEST-TIME SCALING
Knowledge distillation has long been shown to improve in-weights learning (IWL), resulting in stronger performance in standard evaluation tasks and benchmarks (Gemma et al., 2024; 2025). However, in modern L.Mss, the desired capabilities extend much beyond the classical setting of IWI.. The ability to generate diverse solution paths is critical for skills like test-time scaling and search at inference (Chow al., 2024; Dang et al., 2025; Chen al., 2025), but more crucially, to also enable better post-training with reinforcement learning with verifiable rewards (RLVR). Likewise, in-context learning (CL)-wher models learn and adapt from inference time prompts is especially desirable. In this section, we examine how pretraining with distillation shapes these two capabilities key to the current LLM frontiers: test-time scaling and in-context learning (ICL).
3.1 DISTILLATION IMPAIRS IN-CONTEXT LEARNING
The seminal work of Olsson et al. (2022) introduced induction heads as key mechanism behind in-context learning in modern L.MMs. They enable models to "copy" tokens from earlier positions in the input into later parts of the output (Olsson et al., 2022; Edelman et al., 2024; Bietti et al., 2023). For instance, given prompt like "I work at Gym." " an induction head helps the model replicate "Gym" in follow-up question about workplace. This copying ability is critical for tasks requiring models to attend to and reuse information from the context.
We train an 8B teacher on IT tokens and IB students with and without distillation-on on the same IT tokens. We call this the "IsoData" setup, which ensures fair comparison by removing any indirect data advantage for the distilled model. To measure model's ability to copy from context which is hallmark for induction head learning- we use 3 benchmarks: (a) context-based QA (DROP (Dua et al., 2019), RACE (Lai et al., 2017)) (b) needle-in-a-haastern tasks (babilong (Kuratov et al., 2024)); and (c) counterfactual context QA (Goyal et al., 2025), where the correct answer as per the context contradicts factual knowledge (i.c. answer based on model's memory weights), forcing the model to rely solely contextual cues. Conterfactual evaluations thus give clearer signal regarding model's abilities to follow the input context (see samples in Appendix 1.1).
Experimental Setup:
Observations: Figure 3 compares the in-context learning performance of the two IB models trained with and without distillation, as the training data is scaled to IT tokens. We observe consistent
4

---

Published as a conference paper at ICLR 2026
pattern that as the training tokens are increased, the relative advantage of distillation the standard pretrained model keeps on diminishing. Infact the distilled model eventually underperforms in the IsoData setup (IT tokens) on needle-in-hage -halssitnhhhaackkkkka and counterfactual-QA tasks. These observations in stark contrast to the observations on standard language modeling tasks (e.g.. Hellaswag, GSM8k, NaturalQA) in Figure 2, where distillation continues to offer advantage even in the "IsoData" setup.
For counterfactual context-based QA, models often default to using their parametric knowledge which is incorrect. This tendency grows with scale when confidence in the context-based answer low. As a result, distilled models show a accuracy drop compared to the smaller-scale (500B tokens) model. The distilled model is expected to outperform standard pretraining in the non-Iso setting (e.g., 125B and 500B training in Figure 3) due to the teacher's 's indirect data ad advantage.
Why does distillation hurt in-context learning (ICL)? ICL is driven by induction heads that implement low-entropy copy mappings, where the model must reproduce a token from earlier in the sequence. For such deterministic targets, distillation adds no signal perfect teacher's soft labels collapse to the one-hot ground truth, so DPT can only match hard supervision. In practice, imperfect teachers assign non-zero mass to distractors, softening supervision and injecting noise into an otherwise clean mapping, which can hinder learning of the copy circuit. As result, DPT often fails to help help-and can hurthhuduttee -inductionn head formation and in-context learning. We formalize this effect in our bigram sandbox in $ 4.
3.2 DISTILLATION HELPS DIVERSITY
Experimental Sctup: We train IB models on 125B tokens, with and without distillation, using
the Llama-3.1-88 base as teacher. We later also consider an IsoData setting to isolate the effect of extra data the teacher may have seen. For distilled pretraining (DPT), we study two settings: DPT-50
and DPT-90, where the distillation loss is weighted at 50% and 90% respectively (a in Eq. 2). We compare models under two settings: (1) using sampling temperature that maximizes pass@16 and (2) sweeping temperature from O to 1.5 (in increments of 0.1) and plotting pass@1 vs. pass@ This lets distinguish whether model simply stronger overall (higher pass I and pass@16). or whether it has higher generation diversity aatrsttt-achivingg better pass 16 despite similar pass@i 1. We clarify that this work we focus on generation diversity as measured by pass@k, a standard metric in the LLM reasoning literature (Chen et al., 2025; Dang et al., 2025; Chow et al., 2024).
Distilled pretraining unlocks superior test-time scaling. In Figure 4 (top row), we first compare the pass@k curves for standard pretraining (SPT model) and distilled pretraining (DPT-50 model with 50% weight of distillation). We begin by selecting the sampling temperature that maximizes pass@ 16 performance full temperature sweep analysis follows next). Observe in Figure 4(a that while the DPT-50 model has slighlty worse pass@1 compared to the SPT model, the DPT-50 model obtains much higher pass@ 16 (e.g., 28% vs. 23% on GSM). Infact on MATH (Figure 4b), the DPT-50 model even starts off worse than SPT on pass@k at k = 1, but clearly outperforms it as increases exhibiting a striking crossover phenomenon.
Distilled pretraining gives diversity worth seeing 2x data. We now evaluate DPT against a harder baseline of standard pretraining on 2x data (SPT-2xx 250B tokens), increasing the distillation weight to 90% (DPT-90) ). As shown Figure 4 (top row), DPT--0 achieves higher p than SPT-2x across all three benchmarks-an -eeen though it is trained on half the data and has lower pass@1. 1. This highlights the strong diversity gains in generations from distillation.
In Figure 4 (bottom row), we plot pass@1 I vs. 16 across temperatures from to 1.5. Across all benchmarks GSM8K, MATH, and MBPP the DPT-90 curve consistently lies vertically above the SPT-2x curve. That is, for any fixed pass@1, 1, the distilled model achieves a higher pass@16. 16. Note that both the models have the same maximum pass@1 1 (if one optimizes the temperature for pass@1 11), but the distilled model always has higher maximum 16, or infact higher pass@16 @16 for any reasonable pass@1. 1. This reinforces that distilled pretraining enables stronger test-time scaling.
Diversity gains even in Isoa setting For the results in Figure 4, we use L.ama-3.1 as the teacher, trained on more data than our BB students. Importantly, the gains persist in the IsoData setting
5

---

Published as a conference paper at ICLR 2026
gives diverdty date
17.5
40
P on GSM8K
@ K on MATH
15.0
Pass@K on MBPP
35
20
12.5
30
25
10.0
20
7.5
15-
Staalad Prvernaing (SPT) 50% Boiil Luna 181224
10
50% Balid Lown 10P230
2.5
Stadlrd 5272211
10
15
5
10
15
5
10
15
K (#Attempts)
K (# Attempts)
K (#Attempts)
(a)
(b)
(c)
MATHH Poosage 16
MRPP Pase
GSMOK
set poat Les 1017-0 Mhosste yt: 2t date neriad
GSMBK Pasag]
MATH Paseg]]
MBPP Pasagl
(d)
(f)
Figure 4: Distilled pretraining improves generation diversity and enables superior test-time
scaling($ 3.2): Top-row (a-c): We plot pass curves with temperature first optimized for pass 16 performance. Distilled pretraining with 50% weight of distillation (DPT-50) consistently outperforms standard pretraining (SPT) on pass@16, even though it has worse pass@1 on GSM8k and MATH. Bottom-row (d-f): We increase the distillation weight to 90% (DPT-90) and compare against stronger baseline of standard pretraining on 2x data (SPT-2x). We plot pass@1 I vs. pass 16 across temperatures 0-1.5 (step 0.1). DPT-90 consistently achieves higher pass@ 16 for any reasonable pass@1 1, despite using half the data, and attains the top pass @ 16 on all three benchmarks. This shows that distilled pretraining yields models with greater generation diversity and stronger test-time scaling.
where both teacher and student are trained on the same IT tokens (Figure 12): distilled pretraining (orange curve) still outperforms standard pretraining on GSM8K and MBPP Pass@ 16. This shows that DPT's test-time scaling advantage cannot be explained solely by the teacher's access more data.
Why does distillation help with diversity? When prompts admit multiple plausible continua- tions-like "I work -te ground truth data provides only one answer (e.g., hospital), but a teacher model distributes probability mass across many valid completions (e.g.. hospital, gym, cafe). Distilla- tion exposes the student to this richer signal, which intuitively explains why it improves the model's diversity it's generations inference time. We discuss this more formally the next section.
While pass@1 demands only that the top prediction be correct, pass@k evaluates whether any of the outputs are rewirddenn breadth over precision. This subtle shift means that correctly ranking option is enough; the model must distribute probability across multiple plausible answers. Distilled models excel at this, helping them exhibit better test-time scaling. We will discuss this in more detail in Appendix F.
4 BUILDING INTUITION VIA A BIGRAM SANDBOX
In the previous section, saw that distilled pretraining improves test-time scaling but hurts in- context learning performance by hurting the learning of induction heads. In this section, we try to dissect the reasons behind this using simple yet powerful sandox of bigram model.

---

Published as a conference paper at ICLR 2026
Learnnng Dynamics
Learning Dynamios for
High Eutropy Bigram Rows
Low Entropy Bigram Rows
Learning Dynamics for Induction Task
KL. Div. with True Distribution
0.175
0.150
Distllation Peveruining
Pretraning
Diw Distribution
1.50
Pretreiing Distllotion Peccruiing
Induction Thok Acnurncy
0.20
Standard Restrannug
Pestratting
Yeucher Modal
1.25
Thacher Modal
0.125
2.00
0.15
0.100
0.075
0.75
0.10
0.050
0.50
iscucticn S660e
0.25
0.05
0.025
0.00
0.00
Training Spoch
Training Epoch
Training Epoch
(a)
(b)
(c)
Figure 5: Understanding distillation through the lens of a bigram model ($ 4): To dissect why distillation enhances diversity yet impairs in-context learning, we examine these phenomena in simple yet expressive sandbox-a bigram model (Bietti et al., 2023; Edelman et al., 2024). A bigram models first-order Markov chain represented via transition probability matrix. (a) We illustrate that distillation particularly aids the learning of high-entropy rows, corresponding to prompts such as "I work at", which admit multiple plausible completions (e.g.. "gym". "hospital", "restaurant"), (b, e) Conversely, distillation offers no advantage for learning low-entropy rows (b), which not only represent deterministic state transitions (prompts), but are also essential for induction head formation as described by Bietti et al. (2023). Moreover, distillation with an imperfect teacher may even slow or hinder learning these low-entropy, ddduction-headllike rows (c).
4.1 BIGRAM MODEL: LOW-ENTROPY Vs. HIGH-ENTROPY Rows
To build intuition for our results, consider two illustrative prompts: (i) Low Entropy Prompts: "2 + =" with completions: a) 5. b) 4. c) where a) occurs with probability in natural data; and (ii) High Entropy Prompts: "I go to" with completions a) office, b) gym, e) restaurant, d) 33 where a), b), and c) each occur with probability 1/3 in natural data.
Bigram data generation process: A bigram model captures a first-order Markov process, where the next token depends only on the current token. Mathematically. it is represented by a matrix R where each element Tyy denotes the transition probability from token to token j. Our dataset consists of sequences generated from the above bigram model, and the first token is sampled uniformly from the vocabulary. We categorize each row of the transition matrix π as either low- entropy high-entropy, based on the entropy of that row relative to a fixed threshold. High-entropy rows are akin to prompts that have a diverse completion set (recall "I go to" example from above). Low entropy rows then correspond to prompts with less-diverse completions (e.g.. "2+3=").
Distillation accelerates learning of high-entropy bigram rows In Figure 5(a.b), we present the results of the experiments. The teacher is a bigger model trained on 2x more data than the students (details in the Appendix E). We observe that models trained from scratch and models trained via distillation are both at par when it comes to the low entropy rows (Figure 5b). A real distinction appears in how well they approximate the high entropy rows, where the distilled model performs better, i.e., it requires fewer samples to achieve a better approximation of the high-entropy rowgFig 5a). We now formalize the intuitions behind the above arguments.
Sample complexity analysis for bigram model Each row of the bigram matrix и is p-sparse, i.e., contains most p non-zero entries. We consider sequences of length two. Both the scratch- -rratch-trainedd and distilled student models are parameterized by bigram matrices n scrath and distill respectively, while the teacher is parameterized by yf teacher
7

---

Published as a conference paper at ICLR 2026
Proposition 1. (informal) In bigram learning with p-sparse rows
Sample complexity when training with distillation is Saisult O log k).
Sample complexity when learning bigram without distillation is Sstandirdd Sdistill, where is the upper bound on the the approximation error.
We refer the reader Appendix D for the proof. Consider first the high-entropy setting where the row sparsity = (k)), where denotes the vocab size. The standard pretrained model requires O log k) samples. whereas the distilled model needs only O(k log k). In contrast, in the low- entropy setting where p is constant, both models have sample complexity at most O log k). This reflects the empirical observations from Figure 5(a,b) where observe distillation accelerating the learning of high-entropy rows but no difference for low-entropy rows.
4.2
WHY DOES INDUCTION HEAD LEARNING SLOW DOWN FOR DISTILLED MODELS?
Recall from $ 3.1 that distillation impairs the learning of induction heads-key circuits for in-context learning. We revisit this phenomenon by detailing the induction head setup in bigram sandbox. Following Bietti et al. (2023), we modify the bigram model to embed an induction-styy pattern using trigger tokens. A trigger token is special token such that whenever appears, it is always followed by a fixed token within that sequence. This fixed token differs across sequences but remains the same for all trigger occurrences within given sequence.
Formally, before generating each sequence, we randomly choose a "copy target" token c E (1....,k, k). We then alter the bigram transition matrix so that whenever the current token is the trigger (denoted i == l), the next token is deterministically c. Mathematically:
バラコ 3 [ c) if
it
Sampling from म produces setting where the optimal strategy is to learn to copy the token (c) following a trigger token (the token t in the above mamic-mi the behavior of induction heads in real LLMs (Olsson al., 2022; Bietti et al., 2023). The difference between standard pretraining and distillation emerges in the supervision signal.
In standard pretraining, encountering trigger yields one-hot ground-trute -ruth label for the next token- clenn and unambiguous supervision.
In distillation with a perfect teacher, the soft label distribution is also exactly one-hot, so the supervision identical. In practice, however, teachers are imperfect: they may assign non-zero probability mass to distractor tokens. This produces slightly higher-entropy target distribution, effectively injecting noise into what should be deterministic mapping.
4.3 WHY DOES pass@k IMPROVE FOR DISTILLED MODELS?
We earlier observed in Figure that distilled models have higher pass @ k despite having lower or a similar pass @ compared to a standard pretrained model. In Appendix F, we show that optimal pass@k performance requires accurate estimation the underlying probability distributions rather than just correct ordering of class probabilities (which suffices for pass @ 1). Distilled models, by incorporating soft supervision from teachers, better approximate these probability distributions, especially in high-entropy settings where multiple valid answers exist. We study this theoretically in the setting of estimating a Bayes optimal classifier in Appendix F.
TOKEN ROUTING: MITIGATING THE DROP IN IN-CONTEXT LEARNING
In Section 3.1, we saw that distilled models underperform on ICL tasks because they are based on low-entropy mappings where distillation doesn't 't help. To mitigate this, we propose simple yet effective strategy: token routing. Recall from Equation that during distilled pretraining, there two terms in the loss-one for loss with ground truth labels and the other with teacher's label
8

---

back to Published as a conference paper at ICLR 2026
Contert based
Needle Haystack
Counterfactual Context
Average Across ICL Tasks
Suacar Frenaiiggi riB
07.5
60
Dhat wit Lag
325
Accuracy (%Q
33.0
92.5-
Accuracy
Accurac
33.0
223
IsoData
IsoData
350-
45.5-
IsoData
IsoData
42.3
Figure 6: Token Routing: Mitigating the Drop in In-Context Learning($ 5): Distilled models often struggle on ICL. tasks due softening of supervision on low-entropy (near-deterministic) tokens. To mitigate this, we apply token routing: for each input, we skip the distillation loss on the 15% lowest-entropy tokens, using only ground-truth supervision there. This strategy (red curve) improves over vanilla distillation (orange). As shown in Table 1, these gains come without hurting standard language modeling performance.
(distillation loss term). Rather than applying distillation loss with the teacher's label on all tokens, we dynamically adjust the supervision based on the entropy of the teacher's output. Specifically, given input sequence, we first compute the teacher's 's soft labels for the sequence. We then drop the distillation loss term for x% of the positions with lowest entropy in teacher's only the standard hard-label supervision with the ground truth here.
In Figure 6, we show results with æ = 15% token routing. On two of three tasks-Needle a vanilla DPT, partially closing the
Haystack and Counterfactual QA QA-routing yields clear gains gap with standard pretraining. On Context- based QA, no gain is expected since vanilla DPT already outperforms standard pretraining. Importantly, routing low entropy tokens to standard pretraining objective does not hurt standard LM tasks (Table 1), reinforcing the view that distillation gains arise from high-entropy teacher labels. Appendix G.1 reports results with 30% routing, which offers no further ICL benefit and degrades standard benchmarks.
Practitioners guidelines: While preliminary, token-routing demonstrates how distillation centric data curation pipelines can help and we hope our work motivates future research in this direction. Beyond token routing, we explore several additional design choices crucial for effective distilled pretraining in Appendix G. These include comparing distillation against other diversity-enhancing methods like multi multi-token prediction, investigating the choice of teacher model (base vs. instruction- tuned vs. RL-trained), and analyzing the impact of top-k sampling during distillation. These important practical considerations that can significantly impact the effectiveness distilled pretraining.
RELATED WORKS
Modern paradigm of distillation In the past year, we've witnessed resurgence of distillation in the context of modern LLMs. Both the Llama-3.2 (1B and 3B models) (Meta AI, 2024b) and Gemma model families (sizes ranging from 3B 27B) (Gemma et al., 2024; 2025) rely heavily on pretraining distillation mechanisms. These models primarily employ the prominent weighted loss introduced by Hinton al. (2015). Both Llama-3.2 and Gemma series of models use teacher models that have been trained on way more data than the what the student model is ultimately trained on. In this work, we first show that distilled pretraining continues help even in the IsoData setting, More crucially, later we isolate intriguing tradeoffs of distilled pretraining.
Synthetic data, generated by teacher models, now commonly used to enrich pretraining corpora, effectively constituting another form of hard-label distillation. Cha & Cho (2025) analyzed distillation using synthetic data generation (hard-label distillation), where students learn from samples drawn directly from the teacher model. In contrast, our analysis focuses on Hinton et al. (2015)-style pretraining distillation, where students learn from soft labels provided by the teacher, leading us to distinct conclusions from Cha & Cho (2025) regarding prediction diversity. We believe that this happens because of the difficulty in sampling diverse synthetic pretraining data (hard labels) from the teacher. Recently, Busbridge al. (2025) discuss how distillation might not be helpful under certain

---

Published as a conference paper at ICLR 2026
compute- matched settings. However, in $ we argue that incorporating teacher logit computation cost might not be the correct setting, and is more important to consider data-cons settings.
Improving distillation for LLMs Li et al. (2021) discuss using small teacher models for tokens where the student model predictions are less confident. Cho & Hariharan (2019); Zhang et al. (2024a); Mirzadeh al. (2019); Zhang et al. (2023); Beyer et al. (2022) highlight bigger teacher is not always better and propose various ways to mitigate capacity mismatch between student and the teacher. Our practitioner's 's guidelines in this work are complementary to the above findings.
We refer the reader to Appendix H for additional related works and Appendix I for future directions.
7 REPRODUCIBILITY STATEMENT
To ensure the reproducibility of our findings, we provide comprehensive implementation details and experimental specifications throughout the paper and appendices. First, the full details of our training equation we used is explained in detail in Appendix B. Our experimental setup, including model architectures, hyperparameters, training procedures, and evaluation protocols are detailed in the "experimental setup" subsection in $3.1 and $3.2. The dataset composition and other hyper- parameters detailed in Appendix C. All theorems and proofs detailed in the relevant sections in the Appendix.
REFERENCES
Rishabh Agarwal, Nino Vieillard, Yongchao Zhou, Piotr Stanczyk, Sabela Ramos, Matthieu Geist, and Olivier Bachem. On-policy distillation of language models: Learning from self-generated mistakes, 2024. http: org/ alos 22306 1369..
AlphaEvolve. Alphaevolve: A gemini-powered coding agent for designing advanced algorithms, 2025.
Zhangir Azerbayev, Hailey Schoelkopf, Keiran Paster, Marco Dos Santos, Stephen McAleer, Albert Q. Jiang. Jia Deng. Stella Biderman, and Sean Welleck. Llemma: An open language model for mathematics, 2023.
Jimmy Ba and Rich Caruana. Do deep nets really need to be deep? Advances in neural information processing systems, 27, 2014.
Edward Beeching. Lewis Tunstall, and Sasha Rush. Scaling test-time compute with copppacss HuggingFaceet!]
open models, 2024. URL https:/ blogpost- sca wwwlllinngtsttt test-etttee time- compute....e.e
Lucas Beyer, Xiaohua Zhai, Amélie Royer, Larisa Markeeva, Rohan Anil, and Alexander Kolesnikov. Knowledge distillation: A good teacher is patient and consistent, 2022. URL piltts:://rrivv org/abs// 2106. 05237.
Alberto Bietti, Vivien Cabannes, Diane Bouchacourt, Herve Jegou, and Leon Bottou. Birth of transformer: A memory viewpoint. Advances in Neural Information Processing Systems, 36: 1560-1588, 2023.
Cristian Buciluă, Rich Caruana, and Alexandru Niculescu-Mizil, Model compression. In Proceedings of the 12th ACM SIGKDD international conference on Knowledge discovery and data mining, pp. 535-541, 2006.
Dan Busbridge, Amitis Shidani, Floris Weers, Jason Ramapuram, Etai Litwin, and Russ Webb. Distillation scaling laws, 2025. URL http://awwwi.oora/////aaa/////// 08666.
Sungmin Cha and Kyunghyun Cho. Why knowledge distillation works generative models: A minimal working explanation. arXiv preprint 23x..:2505.33/11 /1, 2025.
Feng Chen, Allan Raventos, Nan Cheng, Surya Ganguli, and Shaul Druckmann. Rethinking fine- tuning when scaling test-time compute: Limiting confidence improves mathematical reasoning, 2025. URL https: /ttpas:///rxiiiiii /arxivivvv.. wwweiiioorrraaae abs 22502 07154.
10

---

Published as a conference paper at ICLR 2026
Guoxin Chen, Minpeng Liao, Chengxi Li, and Kai Fan. Alphamath almost zero: Process supervision without process, 2024. URL https: /atsss///axivi www.iiii.....oooeooooooooo abs 2405. 03553.
Jang Hyun Cho and Bharath Hariharan. On the efficacy of knowledge distillation, 2019. URL 1910.
https: / httpaaiiv...//..aaa///////oo wwtp//ggg/////ooo/eo/.////.e///t /1900 01348.
Yinlam Chow, Guy Tennenholtz, Izzeddin Gur, Vincent Zhuang, Bo Dai, Sridhar Thiagarajan, Craig Boutilier, Rishabh Agarwal, Aviral Kumar, and Alcksandra Faust. Inference-aware fine-tuning for best-of-n sampling in large language models, 2024. URL https://arxii:. //tp://aaaivvvvi org/ abs 2412. 15287.
Xingyu Dang, Christina Back, Kaiyue Wen, Zico Kolter, and Aditi Raghunathan. Weight ensembling improves reasoning in language models, 2025. URL 10478.
Dheeru Dua, Yizhong Wang, Pradeep Dasigi, Gabriel Stanovsky, Sameer Singh, and Matt Gardner. DROP: A reading comprehension benchmark requiring discrete reasoning over paragraphs. In Proc. of NAACL, 2019.
Benjamin I.. Edelman, Ezra Edelman, Surbhi Goel, Eran Malach, and Nikolaos Tsilivis. The evolution of statistical induction heads: In-context learning markov chains, 2024. URL https:
//t.xxi..rggaab////// /2002 11004.
Tommaso Furlanello, Zachary Lipton, Michael Tschannen, Laurent Itti, and Anima Anandkumar. Born again neural networks. In International conference on machine learning, pp. 1607-1616.
PMLR, 2018
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy. The Pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2 001000027 2020.
Team Gemma, 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 arXiv:2403.08295, 2024.
Team Gemma, Aishwarya Kamath, Johan Ferret, Shreya Pathak, Nino Vieillard, Ramona Merhej, Sarah Perrin, Tatiana Matejovicova, Alexandre Ramé, Morgane Rivière, al. Gemma 3 technical report. arXiv preprint arXiv:2503. 2025.
Anastasios Gerontopoulos, Spyros Gidaris, and Nikos Komodakis. Multi prediction needs registers, 2025. URL https: /ttp://aaavv.//aa//aa/eeeea.aa.a org/ http/bb//25000 22002 10518.
Fabian Gloeckle, Badr Youbi Idrissi, Baptiste Rozière, David Lopez-Paz, and Gabriel Synnaeve. Better faster large language models via multi-token prediction, 2024. URL https: //axxiv.. org/abs/ 2404 19737.
Jianping Gou, Baosheng Yu, Stephen Maybank, and Dacheng Tao. Knowledge distillation: A survey, International Journal of Computer Vision, 1789))1189-1189, 2021
Sachin Goyal, Christina Baek, J. Zico Kolter, and Aditi Raghunathan. Context-parametrie inversion: Why instruction finetuning can worsen context reliance, 2025. URL abs/2410 10796.
Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in neural network. arXiv preprint arXiv:1503.023: 2015.
Yuri Kuratov, Aydar Bulatov, Petr Anokhin, Ivan Rodkin, Dmitry Sorokin, Artyom Sorokin, and Mikhail Burtsev. Babilong: Testing the limits of llms with long context reasoning-in-a-hay-ee 2024.
Guokun Lai, Qizhe Xie, Hanxiao Liu, Yiming Yang, and Eduard Hovy. RACE: Large-scale ReAding comprehension dataset from examinations. Proceedings of the 2017 Conference on Empirical Methods in Natural Language Processing, pp. 785-794, Copenhagen, Denmark, September 2017. Association for Computational Linguistics. doi: 10.18753/ URL https: //acle hot ogy.org/D17-119..
11

---

Published as a conference paper at ICLR 2026
Jeffrey Li, Alex Fang. et al. Datacomp-Im: In search of the next generation of training sets for language models, 2025. URL https: : /ttp/aix.v...//// org/ http////////////// 11759
11794.
Lei Li, Yankai Lin, Shuhuai Ren, Peng Li, Jie Zhou, and Xu Sun, Dynamic knowledge distillation 2b0. 11295.
for pre-trained language models, 2021. URL https: :
Shalev Lifshitz, Sheila A. Mcllraith, and Yilun Du. Multi-agent verification: Scaling test-time compute with multiple verifiers, 2025. URL
Wang Ling, Dani Yogatama, Chris Dyer, and Phil Blunsom. Program induction by rationale generation : Learning to solve and explain algebraic word problems, 2017. URL URL.ttpa://arxivvooo abs/1705. 04 6.
David Lopez-Paz, Léon Bottou, Bernhard Schölkopf, and Vladimir Vapnik. Unifying distillation and privileged information. arXiv preprint arXiv:1511.0364 2015.
Anton Lozhkov. Loubna Ben Allal, Leandro von Werra, and Thomas Wolf. Fineweb-edu: the finest collection of educational content, 2024. URL Huggi www.ginggaaFFFFiimm f eeweb-edu.
http://aaa......aaaaeeeeeeeeeeae /atasets/
Aditya K Menon, Ankit Singh Rawat, Sashank Reddi, Seungyeon Kim, and Sanjiv Kumar. A statistical perspective on distillation. In International Conference on Machine Learning, pp. 7632- 7642. PMLR, 2021.
Meta AI. The Ilama 4 herd: The beginning of new era era natively multimodal ai innovation. https: a hwwm///llg/lllln in //ai //ai.meta.
2024a https://
Meta
AI.
Llama 3.2:
Revolutionizing
edge ai and vision with
https: / meta. httpaaaaataa...oooooeaeamea com/blog/ 1lama-3 3- http::/:eeeeeee20//////---o-aa-1 ai.meta. combblog/llama 3- wwwnectoonccooeeeoooooooo 2024 20-- -eo--eeddeemm-eeee--eeeee---eeee mobile- detiiieeeeeiii//.....
open, customizable models.
2024b. https://
Seyed- Mirzadeh, Mehrdad Farajtabar, Ang Li, Nir Levine, Akihiro Matsukawa, and Hassan Ghasemzadeh. Improved knowledge distillation via teacher assistant, 2019. URL https:
/ www.xii.ogg/bba.....ooaoe 1902 03393.
Seyed Iman Mirzadeh, Mehrdad Farajtabar, Ang Li, Nir Levine, Akihiro Matsukawa, and Hassan Ghasemzadeh. Improved knowledge distillation via teacher assistant. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pp. 5191-5198. 2020.
Hossein Mobahi, Mehrdad Farajtahar, and Peter I.. Bartlett. Self-distillation amplifies regularization in hilbert space, 2020. URL https: www.
2002.05715.
Niklas Muennighoff, Zitong Yang, Weijia Shi, Xiang Lisa Li, Li Fei-Fei, Hannaneh Hajishirzi, Luke Zettlemoyer, Percy Liang. Emmanuel Candès, and Tatsunori Hashimoto. sl: Simple test-time scaling, 2025. URL htt.t:p:://axxx..orrgooo
www.iiv..og/aa....oooeaa 22011. .99993.
Vaishnavh Nagarajan, Aditya Krishna Menon, Srinadh Bhojanapalli, Hossein Mobahi, and Sanjiy Kumar. On student-teacher deviations in distillation: does it pay to disobey?, 2024. URL
https:/ http://xxxiw...a.a/.////
Vaishnavh Nagarajan, Chen Henry Wu, Charles Ding, and Aditi Raghunathan. Roll the dice look before you leap: Going beyond the creative limits of next-token prediction, 2025, URL
https: //tt://aaivv..// org/ abs 152666......
neogithub. Github code dataset, 2022.
https://haggginffccccoo//
codeparrot//iithubeee..
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac -oddield-Dodds, Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse. Dario Amodei, Tom Brown. Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. In-context learning and induction heads, 2022. URL https: / //axivv.. org/ abs/2209 11..118955 1895.
12

---

Published as a conference paper at ICLR 2026
Mary Phuong and Christoph Lampert. Towards understanding knowledge distillation. In International conference on machine learning, Pp. 5142-5151. PMLR, 2019.
Mher Safaryan, Alexandra Peste, and Dan Alistarh. Knowledge distillation performs partial variance reduction. Advances in Neural Information Processing Systems, 36:75229-75258, 2023.
David Saxton, Edward Grefenstette, Felix Hill, and Pushmeet Kohli. Analysing mathematical 904. 01557.
reasoning abilities of neural models, 2019. URL https:/ htt.://s//aaxv..orr////////
Pier Giuseppe Sessa, Robert Dadashi, Léonard Hussenot, Johan Ferret, Nino Vieillard, Alexandre Ramé, Bobak Shariari, Sarah Perrin, Abe Friesen, Geoffrey Cideron, Sertan Girgin, Piotr Stanczyk, Andrea Michi, Danila Sinopalnikov, Sabela Ramos, Amélie Héliou, Aliaksei Severyn, Matt Hoffman, Nikola Momchev, and Olivier Bachem. Bond: Aligning lims with best-of-n distillation, 2024. URL https: arxiv. org. abs///////// 2407 1422..
Amrith Setlur, Chirag Nagpal, Adam Fisch, Xinyang Geng, Jacob Eisenstein, Rishabh Agarwal, Alekh Agarwal, Jonathan Berant, and Aviral Kumar. Rewarding progress: Scaling automated process verifiers for IIm reasoning, 2024. URL https: arriv. hwt.:/orgg.///aeeeeeeeeeeesse//e
2t//a//ss......///// 08146.
Charlie Snell, Jaehoon Lee, Kelvin Xu, and Aviral Kumar. Scaling IIm test-time compute optimally can be more effective than scaling model parameters, 2024. URL httph:t/t/aaaattt /arxiviiiiiv///ii//aaaa//aaaaava www.i..gg/...............t 2408. 03314.
Abitha Thankaraj, Yiding Jiang, J. Zico Kolter, and Yonatan Bisk. Looking beyond the next token, 2025. URL https: www.ivivvvv... wwtariivvoro. org/ /abs////////// 22004 .5504.13366.............11.33366
An Yang, Anfeng Li, al. Qwen3 technical report, 2025. URL http:::/aarwww.wog/g/////////// 2505. 09386.
Weizhe Yuan, Jane Yu, Song Jiang, Karthik Padthe, Yang Li, Dong Wang, Ilia Kulikov, Kyunghyun Cho, Yuandong Tian, Jason Weston, and Xian Li. Naturalreasoning: Reasoning the wild with 2.8m challenging questions, 2025. URL ht http///xxw...ggbb//////aaa/a
13124.
Yang Yue, Zhiqi Chen, Rui Lu, Andrew Zhao, Zhaokai Wang, Yang Yue, Shiji Song, and Gao Huang. Does reinforcement learning really incentivize reasoning capacity in Ilms beyond the base model?, 2025. URL https: / /awxivivvvvvm////////aa.//aaaava org/ /abbbb/2544..88
Chen Zhang, Yang Yang, Jiahao Liu, Jingang Wang. Yunsen Xian, Benyou Wang, and Dawei Song. Lifting the curse of capacity gap in distilling language models, 2023. URL https: //arxiv. wwg.bbggg/a..................... abb.///5...... 12129.
Chen Zhang, Dawei Song, Zheyu Ye, and Yan Gao. Towards the law of capacity gap in distilling language models, 2024a. URL
Yiming Zhang, Avi Schwarzscilld, Nicholas Carlini, Zico Kolter, and Daphne Ippolito. Foreing diffuse distributions out of language models, 2024b. URL https: /twxxarxvvv...o. org/ /rg/abs/////ee /2404. 10859.
13

---

Published as a conference paper at ICLR 2026
A LLM USAGE.
We used large language model (LLM) as a general-purpose tool to assist with polishing the writing style. The LLM was involved research ideation, experiment design, analysis. All technical content, results, and conclusions are entirely our own, and we take full responsibility for the final manuscript.
B PRELIMINARIES
We start by revisiting the setup of distillation from Hinton et al. (2015). We are given a dataset {( y)]== of inputs I; E Rt's and the labels y/ss A where k is the num- ber of classes and R is probability simplex over those classes. Let begin with the ob- jective of training a model from scratch on the above data using cross-entropy loss f, h" E > Zi=1 (yi, [hh(z)))) where h is a candidate function drawn from the hypothe- sis class H. R A is the softmax function dz(2) Σi=1 exp(2;)
arg minkeH 12
exp(s,) and E(y, 9)
-ド= My log(ú;)-
are now ready to define the standard objective used in distillation:
ht arg min NEH
Σ eyu (((x))(()))) + Ee ((ss))((())))))
(2)
where a
[0, 1] and 5; = o (hteacher (x:)/T) is a soft label generated by the teacher using temperature T.
This form distillation has recently been adopted in pretraining language models as well. We first start by describing the next-token prediction objective over sequence (z 21)
Σ ((HHI))(h(()))))
(3)
The objective function used in pretraining distillation is
1-1 a){ ( + Σ [(Hnno((MM)s)))
+
(4)
where 83+1 σ (fiteachhr (IT)
C GENERAL EXPERIMENTAL DETAILS
Pretraining dataset composition Our pretraining corpus consists tokens drawn from diverse domains to ensure broad coverage of knowledge and reasoning capabilities. The majority of the data comes from the DCLM (Li et al., 2025) like baseline dataset and GitHub repositories (neog- ithub, 2022). In addition, we include range of specialized sources spanning mathematies, coding, scientific literature, and high-quality web content. Specifically, our mixture includes DeepMind Math- ematics problems (Saxton et al., 2019), Proof Pile 2 collections (ArXiv, Open Web Math, Algebraic Stack) from Azerbayev et al. (2023), Stack Exchange from the pile (Gao et al., 2020), Fine Web- Edu (Lozhkov et al., 2024), and smaller curated sets such as Natural Reasoning Dataset (Yuan et al., 2025) and AQuA (Ling et al., 2017).
Pretraining Hyperparameters For temperature(T) in distillated pretraining, we do a gearchearch over T {0.5, 3). We select the temperature which gives the best performance standard benchmarks, In our experiments, = 1 worked the best. We pretrain with cosine scheduler using a learning rate of 36-3 for 1B models and 3 for 8B models.
14

---

Published as a conference paper at ICLR 2026
D PROPOSITION 1 (FORMAL).
Proposition 1.
If the number of sequences observed grow as O(k log k + ls klog())) then the distill π teacher with probability at least 1 6.
If the number of sequences observed grow as O( [klog k+ (p/e2t1ppp)e/eeeeee log
then for each row
i E E of the dataset.
Ti|l1] S with probability 1 6, where is computed over the entire draw
Proof. To prove the first part, let us recollect standard result.
The coupon collector problem studies the following question. Suppose each box contains a coupon, and there different types of coupons, What is the number of boxes we need to see 7'' before we have collected all coupons? Assuming each coupon is drawn uniformly at random,
P(T 3k log k) -
Substitute / = 1 log log k
we obtain
P(T > k + log( ())) < <8
Translated to our setting, this means if we observe & log h + log( then with probability at least 1 - 8 each of the distinct & tokens have been observed at the first position in the sequence. This completes the proof for the first part.
We now turn to the model trained from scratch. The log-likelihood of a model is written as Eig May log(#t)), where n, is the number of times we see a token j appear after token i. The solution to maximum likelihood is simply ăij- Biy where n, = ΣjEi Rij. Rig is an unbiased estimator of Taj. Define
For this model, we need to ensure that each row in the estimated matrix is close to the true row. Next, we want to bound the distance between || T1, where we particularly use & distance to emphasize the role sparsity. Observe that the variance each element of the row is [
Tay )²]
Ta:(1 (1)7,))
E, m.y
Observe that
E B m5ll) S (fiaa n1z)
Tir(1 -ii)) Σ nij
E[
Tay
Tij
mij) nij
(5)
To compute, || TT/l,, we only need to sum over the terms that are non-zero owing to the sparsity assumption. Suppose that without loss of generality first p terms are non-zero. Hence, we obtain
Bli+
Σ
Tay
Tty)
(6)
Maj
SSP
JSp
can arrive at a simple upper bound for Ejsp Schwarz inequality. We express
Tij (1 - F.z) as follows. We again apply Cauchy-
Σ 35p
Tey n.y) = (1,1 Tn
n,1), v TTz T22), V Tip(1 (1 - Tzp)))
(1, [1,[v nal (1-Ta), T1 - T:2).
Tip - S VD. Σ (mT))( - TT)) S vp
15

---

Published as a conference paper at ICLR 2026
We substitute this in equation 6 to obtain
E Bl
Tellil S
Tzy(1 Tzg)
P
V
74.
74
JSP
From the above, we can observe that if 2 then
Blll -. S [A]
Hence, if each token i observed the first position of the sequence least 2, then we should obtain the desired outcome we set out to prove in this part.
We now recollect the generalized version of coupon collector's problem. In the generalized version one is interested in computing the number of boxes to collect defined as Tm before collecting m copies of each coupon. In this case,
[[Tm] = k log k + (m 1)k log log le
we apply Markov inequality on the above, we obtain simple bound
P Tm
Thus from the above, we gather that if the number of boxes collected least 1 (klog k + (m
1) log log k) then with probability at least å we have collected m copies each coupon.
We can now substitute m to obtain our bound of k (p/e 11) This completes the proof.
E EXPERIMENTAL DETAILS FOR BIGRAM SANDBOX AND INDUCTION HEAD LEARNING
Our bigram sandbox experiments designed to provide simple, controlled testbed for under- standing how distillation influences test-time scaling and in-context learning. All results in Section are derived from this setup.
Data generation. The vocabulary consists of & = 64 tokens. The bigram transition matrix & was constructed to include mix of low-, medium-, and high-entropy rows: low-entropy rows concentrated probability mass on 3-5 tokens; high-entropy rows were nearly uniform; medium- entropy rows had intermediate profile. Trigger tokens were randomly selected (5, 10, or 20 triggers per experiment), with trigger-output mappings varying across sequences to induce induction head learning (following Bietti et al. (2023)). Sequences with these bigram transitions, with special logic to ensure copying behavior for trigger tokens.
generated using first-order Markov chain
Both teacher and student models were implemented as small Transformers with 2-4 layers, causal masking, and a fixed sequence length of 64. Teacher models used 128 228-dimensional embeddings; students used 64 -4-dimensional embeddings. Training was performed with Adam optimizer and cosine learning rate schedule.
Models.
Training. Teacher models were trained on datasets of size 16k sequences. Student models were trained with either cross-entropy (CE) loss or knowledge distillation (KD), using soft logits from the teacher. Dataset sizes for students were 8k sequences i.c. half the data. The KD objective used temperature T = 2.0 and mixing coefficient 0.5(Equation 2).
16

---

Published as a conference paper at ICLR 2026
Evaluation. All models were evaluated fixed held-out dataset of 4k sequences. Metrics included: Induction head accuracy (trigger - copy) as shown in Figure lc: and KL-divergence between the ground-truth distribution (bigram rows) and the learnt distribution for low-, medium-, and high-entropy rows as shown in Figure
Our full codebase will be released for reproducibility.
F
WHY DOES pass@k IMPROVE FOR DISTILLED MODELS?
Demistifying pass@k k trends: In Figure 1(b), start with a worse pass@1 and can have much better pass@k. Is this a mere accident, or does there exist a deeper principle behind the observations?
puzzling finding. The distilled model can
Suppose that our data consists of one fixed prompt v. which is followed by three options y = {0. 1, The true probabilities ply 0/ær)= 1 + c, p(y a
and 0 with 0. Define three
classifiers:
Bayes optimal classifier, Cl: Assigns probability 1 to class 0 and achieves the optimal pass@1 accuracy of 1 +.
Diverse classifier with right coverage, C2: Assigns a probability of to both classes 0 and 1. This classifier achieves suboptimal pass accuracy of 2-
Diverse classifier with wrong coverage, C3: Assigns probability of classes 0 and 2. This
classifier achieves suboptimal pass01 accuracy of _ + 5.
Interestingly, observe that the pass0k accuracy of C1 is & +e for all kt The pass@k accuracy of C2 is ( for all k. The pass accuracy of C3 is ( + c) 1 (1)* for all As shown in Figure 7, the classifier C2 exhibits crossover over the Bayes optimal classifier C1. Thus, the Bayes optimal classifier is suboptimal higher pass@A: Further, C3's support does not contain the support of the true distribution, highlighting the importance of right coverage over the correct solution space.
Accuracy
0.6
0.4
0.2
Cl: Bayes Optimal
The above example leaves us with the question that the Bayes optimal classifier is not optimal for pass@k, then what is? We derive this classifier below.
Right Coverage
C3i Wrong Coverge
0.0
K (Number of Attempts)
Figure 7: Bayes optimal for pass@1 is not optimal for passukk. A diverse classifier with correct coverage (C2) outperforms the Bayes optimal classifier (C1) at higher k, while incorrect coverage (C3) remains suboptimal. Cover- age-not just pass@l-is key to improving pass@k.
Generalized Bayes optimality for pass@k this section, we restrict ourselves to binary classification tasks with the true probability distribution over labels conditional on denoted as p(ylz)).
Recall the definition of a Bayes optimal classifier for binary classification. For each in the support of the training distri- bution, the classification rule
{
plu = ())) p④y = (r) K
(8)
Define general classifier which assigns probability axx) to class and (((r) class O.
17

---

Published as a conference paper at ICLR 2026
Theorem The generalized Bayes optimal classifier that achieves the optimal pass@k assigns for each a in the training distribution
p(g=lla) p(y=0(a)
o² (3r)
(9)
p(y=ll)) p(
I
1+
Proof. pass@k accuracy of a classifier checks at least one of the attempts of the classifier predicts the label correctly. For fixed in the support of the training distribution, the pass@k accuracy of this classifier is stated as
pvu = 10 (1- + plu = -
(10)
To understand the above expression, let look at the first term. Conditional on 1, (1 (B(x)))) is the probability that at least one of the attempts by the model says class 1.
To simplify notation, let us write p(y 1 as p. ax(x) as and rewrite the above
p(1 (1 a)*) + p) 111-at a*).
(11)
The function is concave in for E j0. 1] and k > 1, with second derivative given by - (k) (p(1 - aa^- + - 2-) Setting the first derivative to zero gives
I-P
a"
1+
Thus, the generalized Bayes optimal classifier is as given in Eq. 9. Observe that as h: approaches from the right, the expression reduces to the standard Bayes optimal classifier: p(y = 1(x) > 1/2, then a* (x) = 1; otherwise, a a²(x)= 0. This completes the proof.
few key remarks follow. For 1, the Bayes optimal classifier is a (x) 1(p(y 1|x) ))
Optimal pass@1 requires only correct ordering of class prohabilities not precise estimates of ply 1|x). In contrast, optimal pass@kk demands accurate estimation of p(y 1|x). Distilled models better approximate these distributions, especially in high-entropy settings. While this may improve pass 1, it yields superior pass@k performance.
G PRACTITIONERS GUIDELINES
G.1 TOKEN ROUTING: MITIGATING THE DROP IN-CONTEXT LEARNING
We introduced token routing in $ as simple yet effective strategy to mitigate the drop in in-context learning observed with distilled pretraining. Figure 6, we showed results when distillation loss is skipped on 15% of the tokens in each in the teacher's soft labels. This routing improves in-context learning on 2 out of the evaluated benchmarks. In Figure 8, we first share additional results when routing 30% of the tokens. We observe that 30% token routing improves performance only on task compared to the tasks when routing 15% tokens. Moreover, too much token routing hurt performance standard tasks as shown in Table 1.
those with the lowest entropy
We share the performance with token routing standard language modeling tasks and reasoning benchmarks Table 1. We observe that routing 15% of the tokens preserves the performance on standard language modeling benchmarks. However, if further increase the tokens which
18

---

Published as a conference paper at ICLR 2026
Contert based
Needle Haystack
Counterfactual Context
Average Across ICL Tacks
Saacar Treniiig iiB
07.5
60
Dhat D1g
Toatee Maly thi
325
Accurss (%0
92.5
33.0
Accuracy
Accurac
33.0
233
IsoData
IsoData
250-
45.5
IsoData
IsoDala
42.3
Figure 8: Token Routing: Mitigating the Drop in In-Context Learning In Figure 6 we presented results when routing 15% of the tokens. Here present results when routing 30% of tokens.
HtrlbSoog
TQ4
MISTY
Clamstbaale
GSMIK
CSMBK
ARCE
CDDS
SQJAD
Arg
(Acouray)
(F3)
(pai
(paa8
(B30
(poss
(pass
(pass
(Acc.)
(Acc.)
(Aec)
(pess(i))
(passl1d)
(FI)
9.75
13.99
Discilation
64
33.64
17.08
47.64
9.76
9.15
4.25
31.59
54
66. 55
79
15.6
55.31
34.05
Ditiiationn Bosting (1995)
53
52.23
17.16
45.91
9.15
9.16
32.9
59.21
67.57
79
55.48 15.6
Rosting (305)
Table Token Routing ($ 5) does not significantly hurt performance on standard benchmarks. Doing distillation only tokens for which teacher label has high-entropy mitigates the drop in ICL performance (Figure 6) while preserving the performance on standard language modeling tasks and reasoning tasks, shown in the table. This also reinforces the fact that gains in reasoning tasks come primarily from tokens where teacher label has high-entropy, and removing the distillation loss term for tokens where teacher label has low-entropy does hurt standard tasks. As expected, routing lot of tokens (c.g., 30%) hurts the standard benchmark performance.
GSMOK 25
MATEF PaaNeE
Mrnp
NTP
NTP
NTP
KOP
GSMIK Pass@?
MATH Pass@1
MBPP Pasa@
Figure 9: NTP vs MTP vs Distillation: We compare 1B models trained on IT tokens via (1) standard next-token prediction (NTP), (2) multi-token prediction (MTP), and (3) distillation from an 8B teacher trained on the same IT tokens (IsoData setting). We plot pass@1 vs pass@ 16 curve. Distillation curve lies generally above MTP on GSM8k and MBPP, and matches on HATH-despite no data advantage. In real-world setups, where teachers have seen more data, the gains from distillation expected to even larger.
distillation is not performed to 30%, there is drop in performance, although it still remains above standard pretraining would expect.
G.2 NTP vs. MTP vs. DISTILLATION: WHICH YIELDS BETTER DIVERSITY?
In this work, we showed that distillation produces models particularly well-suited for test-time scaling--rimmrrlyy due to their richer generation diversity. In parallel, recent works on multi-token prediction (MTP) (Gloeckle et al., 2024) have also emerged promising way train inherently diverse models (Nagarajan et al., 2025). This raises a natural question for practitioners: given the choice, should one invest in MTP or in distillation?
19

---

Published as a conference paper at ICLR 2026
MATH500
GSM8K
MBPP
ARC Challenge
0,188
0.30
0.521
Chueedoottoo Acc.
0.424
0.434
0,412
1.39
6:296
6.34
0:20
0.385
6:23
114
HellaSwag
SQuAD
1.43
447
0.467
6.471
0:29
631
0.534
0.529
Teacher Model
0.206
Instruct
RL
241
What makes a better teacher: Base vs Instruct vs RL model(ş G.3); We compare 1B student models distilled from three version of model: base, instruction-tuned and RL-trained. The RL-trained teacher consistently yields the best student- across reasoning (MATH500, GSM8k), coding (MBPP), and even general benchmarks (TQA, HellaSwag, ARC). This suggests that stronger teacher performance may outweigh alignment mismatches with the pretraining objective. Despite common practice favoring base models teachers (e.g., Gemma, Llama-3.2) our findings highlight the potential of RL-trained models as superior teachers for distilled pretraining.
Figure 10:
To answer this, we compare three pretraining strategies for 1B models: (1) standard next-token pretraining (NTP), (2) MTP, and (3) distillation from an 8B teacher trained on IT tokens same as the student corpus.
In Figure 9 we plot pass@1 vs p for the three pretraining choices. We observe that the curve for distilled pretraining lies above those of MTP and NTP. This implies that given any reasonable pass@1, distilled model exhibits higher pass@11 16 (on GSM8k and MBPP) or similar pass@ 16 (on MATH) compared to multi-token pretraining. This is notable given the fairness of our setup-using teacher trained on exactly the same data as the student. In practice, where teachers are often stronger because they have seen findings reinforce distillation's 's strong value proposition for practitioners aiming to train small models that excel under verifier-driven inference settings (AlphaEvolve,, 2025; Snell al., 2024).
data, the advantage of distillation is likely be even greater. These
G.3 BASE VS. RL MODEL: WHAT MAKES A BETTER TEACHER?
A general question had while distilling with teacher was what version of the teacher model should be used: the base version, the instruction-tuned version, or the RL-trained version?
At first glance, the base model appears to be the better choice-it aligns pretraining objective of free-form sentence completion and also with the current practice (Gemma al., 2024; Meta AI, 2024b). contrast, instruction-tun to QA-style prompting, making them less aligned with the standard pretraining setup. But on the other hand, the Instruct and RL versions are often better in many capabilities and performance on downstream benchmarks, particularly for reasoming and code tasks. At the same time, recent works like Dang et al. (2025) highlight that Instruct and RL models suffer from reduced diversity in their generations, which suggests they might not be the better choice as a teacher during pretraining.
naturally with the
and RL-trained models are more tailored
We try to answer this puzzle empirically by training student models of 1B size, distilled from versions of 8B teacher model: the base Llama-3.1-8B, its instruction-tuned counterpart, and the RL-trained variant optimized for reasoning. Interestingly, the results in Figure 10 favor the Instruct and the RL-trained teacher-acrow the board. The student distilled from RL trained teacher not just outperforms on reasoning and coding benchmarks (which might be expected). but also on general language modeling tasks like HellaSwag and TQA. This finding indeed surprised that many distillation pretrained models currently like Gemma series (Gemma et al., 2024: 2025) and the llama-3.2 series (Meta AI, 2024b) distilled using base version of large model as the
well. Note
20

---

Published as a conference paper at ICLR 2026
MATH500
GSM8K
MBPP
ARC Challenge
4.3
6.0
0:39
K22
444
633
3.5
6.00
6.00
HellaSwag
TQA
SQuAD
044
Top-K Configuration
Bouct March
Standard Protroining (PT)
top-k:
top-k:
tog-k: 256
444
top k: 1024
026-
top-k:
Figure 11: Top-k sampling distillation($ G.4): We compare using sparse soft target label by sampling k-logits per token. k = 1 corresponds to token level synthetic data albeit without any soft labels, and outperforms standard pretraining. Using richer soft labels (k 128, 256, 1024, or All) further improves performance, but no clear winner emerges among them.
teacher. Infact, even in our work we used base model as the teacher. We hope these insights help inform better teacher choices for future distilled pretraining.
G.4 TOP-K SAMPLING DISTILLATION
Rather than using the teacher's soft distribution over the whole vocabulary, common prac- tice (Gemma et al., 2025) is to sample l logits per token based on the teacher's original output distribution, and then normaall the weights of the sampled logits to get sparse label(logits not samples are set to 0). This reduces the cost of distillation. In this section, we try to answer whether the choice of & here has (if) any impact on downstream performance. We note here that the case of k = interestingly corresponds to standard pretraining with a "token level synthetic data" from the teacher model.
Figure 11 shows the results. We observe two clear trends: (1) Even top-k 1 outperforms standard pretraining, likely due to the use of synthetic tokens and the teacher filtering out outlier tokens from the ground truth; and (2) using & E 128, 256, 1024, All leads to better performance than top-k 1, as the benefits of soft label distributions begin take effect. However, there is no consistent trend indicating which & (other than & 1) performs best.
H ADDITIONAL RELATED WORKS
Classical paradigm of distillation The story of distillation begins with Buciluă et al. (2006), where the technique was introduced to compress an ensemble of models into a single model. Subsequently, Ba & Caruana (2014) proposed a form of distillation wherein student is trained by minimizing regression loss against teacher logits. Later, Hinton et al. (2015) introduced the most prominent form, combining ground-trutt -rudh--tuhhhhuu labels with soft labels from teacher model. Distillation further evolved into various forms: self-distillation Furlanello et al. (2018), where earlier student checkpoints act as teachers; progressive distillation (Mirzadeh et al., 2020), which earlier teacher checkpoints progressively guide the student; and generalized distillation Lopez-Pat Paz et al. (2015), which integrates standard distillation with the privileged information framework.
An extensive theoretical literature has examined distillation through multiple lenses. For instance, Phuong & Lampert (2019); Safaryan et al. (2023) adopted an optimization perspective to explain distillation's 's benefits, while Menon al. (2021) considered the sample complexity perspective. Given the vast breadth and depth of research on distillation, we refer the reader to Gou et al. (2021) for a comprehensive overview.
21

---

Published as a conference paper at ICLR 2026
Theoretical works on IsoData distillation Theoretical analyses of distillation have primarily explained its benefits through two lenses: sample complexity and optimization. From the sample complexity perspective, Menon et al. (2021) show that distillation improves generalization when the teacher has access to more data (e.g., a Bayes-optimal teacher). However, this framework falls short in the IsoData regime, where teacher and student train the same data. From the optimization perspective, Safaryan et al. (2023) argue that distillation enables the student to converge closer to the Bayes-optimal solution as the teacher improves. Yet, it remains unclear whether such convergence is faster than that of standard SGD when no additional teacher data are available.
The only works in theory that explicitly address the IsoData setting have appeared only recently, and somewhat surprisingly. Mobahi et al. (2020) show that deff-iistilattion can reduce overfitting by dampening variance along the top singular directions of the learned representation. Building on this, Nagarajan et al. (2024) demonstrate that distillation further exaggerates the implicit bias of gradient descent, driving the student to converge more rapidly along top eigendirections. Together, these results suggest that the gains from IsoData distillation arise less from sample complexity or optimization speedups, and more from implicit regularization effects acting through the singular spectrum of the representation.
Modern paradigm of distillation: Post-training Beyond pretraining, distillation is increasingly used in post-training. For example, DeepSeek RI released distilled models via off off-policy distillation, where students are fine-tuned on teacher-gent geeerrteeeee traces (Muennighoff et al., 2025). In contrast, on-policy distillation (Agarwal et al., 2024; Yang et al., 2025) uses student-generated traces with logit supervision from the teacher, and has been shown to outperform off-policy methods. In this work, we study logit distillation during pretraining (while using the ground truth data) and highlight the distinct trends and tradeoff's which emerge compared to standard pretraining.
Diversity for test-time search in LLMs Diversity in generations is crucial for test-time scaling of LLMs. This is especially required for open-ended discovery and reasoning tasks, where verification of the correct answer is easy, thus multiple attempts can be done at a problem. (AlphaEvolve, 2025; Setlur et al., 2024; Lifshitz al., 2025; Beeching et al., 2024). In fact, long line of work focuses on explicitly improving the diversity of generations in LLMs at inference time via diversity aware finetuning (Sessa et al., 2024: Zhang al., 2024b: Chow et al., 2024; Chen et al., 2025). Another line of work explores inference time decoding strategies (Chen et al., 2024) for promoting diversity if generations and hence better test-time scaling. While all these works focus on aatch-fixing the diversity issue via model finetuning, we highlight an intriguing albeit intuitive gain in diversity of base model itself when pretraining with distillation. This is of even more importance given recent findings that post-training or RL simply sharpens base model distribution. Yue et al. (2025) shows that base model is better than RL trained model on pass@k & for high k. Having a base model with high diversity is also crucial for effective post-training with reinforcement learning via verifiable reward (RLVR), as discussed in Dang et al. (2025).
I CONCLUDING REMARKS AND FUTURE DIRECTIONS
While distilled pretraining was notably absent in early LLM training pipelines, it has recently regained prominence, as exemplified in Gemma and Llama series (3.2 and Maverick) which rely solely on distilled pretraining.
In this work, we first addressed a common question arising from the renewed interest in distilled pretraining: Is distillation simply a proxy for accessing the extensive data seen by a larger teacher model, or will it offer inherent benefits even if the student model is trained on all the dataset as seen by the teacher? This question is even more important given the data constrained regime for modern LLMs. Our findings affirmatively demonstrate that the value of distillation extends beyond mere data augmentation. Specifically, distilled pretraining naturally produces models exhibiting greater generation diversity, inherently enhancing test-time scaling capabilities. This insight is especially significant given recent evidence suggesting that post-training and reinforcement learning methods primarily just sharpen existing base model distributions, with base models often matching post-trained models in higher p k scenarios (Yue et al., 2025). Distillation thus provides a foundational improvement via pushing the base model performance itself rather than post-hoc fix.
22

---

Published as a conference paper at ICLR 2026
HellSwiag
TQA
XIBFP
XIBPP
HunanEedds
aPQN
GSMSK
OSMEK
ARC-
ARC-EaS
COPA
MATH
MATH
SQUAD
Agg
(Acommecy)
(EM)
(postt1))
(poostAt))
(pass111)
(BM)
(passttt)
(pascat10)
(No0)
(pessur)
(posstES))
Standard PT
22
21.25
7.45
31.56
7.32
11.36
2.96
23.35
32.96
50 97
76,00
1.20
10:30
28,00
loo wxieta
61.13
24.48
12.75
5.40
12.05
2.30
26.00
34.51
52.37
$1.00
0.50
14
4521
29:92
weight
62,17
25,55
13.50
40.58
8.54
12.05
3.11
26.99
34,42
62.25
$1.00
1.50
17.40
$3.35
31.24
61 X3S
23.70
14
39.56
3.47
8.04
4.40
26.97
35.45
55
76.00
2.40
16.30
$2.32
10.00
Table 2: Additional evaluations for the IB base models trained on 125B tokens (1x data) used in this paper. One can observe the better test-time scaling properties exhibited by distillation pretrained models, on MATH and GSM8k. pass@1 is lower compared to standard pretrained model, but pass@ 16 is higher.
With modern LLMs hitting the data wall and growing interest in enhancing capabilities for open- ended discovery and reasoning tasks, our findings are both timely and impactful. An immediate next step would be to tailor, integrate and evalaute distilled pretraining with other recent advances in pretraining like multi-token pretraining (Gloeckle et al., 2024; Nagarajan et al., 2025) and future- aware pretraining (Thankaraj et al., 2025; Gerontopoulos et al., 2025) for improving diversity of base models.
In our study, we proposed applying distillation selectively on a subset of lakkens-pariicllaiii to mitigate cases where full-token distillation may hurt performance. More broadly, current pretraining datasets have largely been curated from common crawl with standard next-token pretraining paradigms in mind. Moving forward, highly promising research direction would be the development of pretraining datasets and curation approaches specifically optimized for distilled pre-training.
Moreover, given the widespread adoption of distillation in post-training phases-such as fine-tuning on reasoning traces generated by larger models-anote another intriguing avenue is to investigate whether using the same teacher model for both pretraining and post-training distillation could better align these two phases. Our work provides preliminary insights into several practical design choices practitioners face during distilled pretraining, and we hope these contributions support the community in advancing this promising line of research.
J ADDITIONAL EVALUATIONS
Evaluations for the IB base models trained using Llama-3.1-8) as teacher evaluations on standard benchmarks for the base models in Table 2.
We share additional
Higher base model diversity + post-training advantages. The diversity benefits conferred by distillation persist even after post-training on reasoning data, as shown in Figure 13(b.)). Again, observe poossover-phenomente
where model trained with 90% weight of distillation during pretraining, exhibits lower pass@1 than a 50% weight counterpart (red vs orange curve in Figure 13 (b)). However, the model with more distillation heavy pretraining exhibits better test-time scaling due to better diversity in generations.
Finally, in Table we present evaluations general language modeling tasks for standard and eestillation-pretrannee models. As expected, distillation pretraining improves statistical modeling, leading to better performance even on non-reasoning tasks as well, echoing findings in Gemma et al. (2024).
Additional evaluations for IsoData Models (trained using 8B param 1T token teacher) Recall that in & 3.1 and Figure we showed how distillation impairs in-context learning, especially in the "IsoData" setting where the teacher, student and the standard pretrained model all see the same data. standard language modeling tasks where the
Note that this is in stark contrast with performance performance of distilled models continues to be better than standard pretrained models even under the isodata setting, shown in Figure 12.
23

---

Published as a conference paper at ICLR 2026
ARC
Didilidd Fompainig
71.5
Exo0 March
03
Training Cata Sine
Training Data Size
Training Cats Sine
Traiiing Data Sire
TQA
GSMBK Pass@16
NEPP Fass@16
Average Across All Tasks
1255
5028
1235
9228
1252
9928
1235
9022
Training Data
Training
Trasting Cata Sins
Training
Figure 12: Distilled pretraining consistently outperforms standard pretraining even in IsoData setting(S Unlike in-context learning and induction head tasks where distillation underperforms in the isodata regime (Figure 3), distilled pretraining continues to yield better results on standard language modeling tasks that do not rely on induction heads-even when student models are trained on the full IT tokens used by the teacher. Moreover, we continue to see that distilled pretraining rewards with better test-time scaling on the GSM8k and MBPP plots (both as Pass @ 16 curves).
Distillation Scratch Pretraining
Distillation vs Scratch Protraining
40
Pass@ MBPP
35
60
20.0
130
Pres& GSMAK
50
Pass@K on MATH
25
80
12.5
20
30
15
20-
7.5
Pretraining 50% Weight
5.0
50% Datill Lass Weight
Standard Poetrunnn (PD)
10-
90% Distil Loss Weight
2.5
90% Datill Loss Waight
..+
Standand PT with 2x data
Standard PT with 2x data.
0
5
10
15
0.0
10 K (Number of Attemptas)
K (#Attempts)
K (Number of Attempts)
(a)
(b)
(c)
Figure 13: Distillation pretraining diversity leads to better post-training test-time scaling as well: (a) Base model evaluations coding task of most basic python problems (mbpp). Distillation pretrained models exhibit much stronger test-time scaling and diversity in generations, as exhibited by higher pass@k than even model trained 2x more data with standard pretraining. Note that this is despite the fact that both models have a similar pass@1. See Figure for more tasks. (b,c) Diversity gains in base model evaluations persist after post-training. depicted by better test-time scaling after post-training as well on MATH and GSM8k.
J.1 COUNTERFACTUAL CONTEXT BASED QA SAMPLES
In Section 3.1 we showed results on counterfactual context based QA dataset from Goyal et al. (2025). Here, the context has information which is counterfactual i.e. opposite the world fact that might be stored in model's memory. The question and the corresponding answer is based on this counterfactual context. This kind of evaluation helps tease out the effect of model answering using it's memory
24

---

Published as a conference paper at ICLR 2026
(for e.g., in standard QA benchmarks) and ensure that the accuracy reliably reflects context reliance. We share few examples of these questions below from Goyal et al. (2025).
Examples
Example 1:
Context: Following the devastating earthquake in 2030, Kabul was largely de- stroyed, prompting the Afghan government to relocate the capital to Herat. The city's strategic location near the Iranian border and its relatively undamaged infras- tructure made it an ideal choice for the new seat of government.
Question: What is the capital city of Afghanistan?
Answer (based on context): Herat
Memory-based Answer: Kabul
Example 2:
Context: Enrico Fermi was born 1452 in the small town Vinci, Italy, where the fertile landscape and serene environment fostered his early love for nature and art. Enrico was not just an artist but polymath, delving into anatomy, engineering, and even music. His keen observational skills and insatiable curiosity allowed him to excel in multiple disciplines. Enrico's 's big break came not just from his raw talent but from his ability combine art with science, bringing level of realism and emotion previously unseen in painting. His time in Milan under the patronage of Ludovico Sforza was pivotal. It was here that painted "The Last Supper," masterpiece that captured the dramatic intensity of the moment when Jesus announces that one of his disciples will betray him. But was his work on the "Mona Lisa" that cemented his legacy. Enrico's ability to blend art and science, to capture both the physical and the psychological, what to his enduring fame.
Question: What is the name of the artist who made Mona Lisa?
Answer (based context): Enrico Fermi
Memory-based Answer: Leonardo da Vinci
25


Reviews:
[{"review_id": "SyGhE7QWZk", "reviewer": "Reviewer_4peW", "summary": "This paper investigates the effect of soft-label distillation during pretraining on the downstream performance of language models. They differentiate between tasks that rely on in-context learning and tasks that can be test-time scaled. For the first group, they notice a performance drop due to distillation and they propose a justification for this drop. On the other hand, they show that distillation helps test-time scaling.", "questions": "- Please correct me if I'm wrong about the sudden ICL performance drop. Can that be explained by your theory of not properly learning the induction head circuits?\n- Why do you suddenly use Llama-3.1 as a teacher for results in Figure 4, while everything else is done in the more controlled \"IsoData\" setting?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The authors did many experiments, using a thought-out and controlled setup and relatively large language models. The results give novel insights into the effect of distilled pretraining.\n- Using the findings of the controlled experiments, the authors than propose a novel \"token routing\" method that is a simple modification of traditional soft-label distillation that should mitigate (to some degree) the negative properties (bad ICL performance) of naive distillation.\n- The paper shows how distillation can help to generate more diverse outputs that can substantially help with test-time scaling.\n- I really like all the additional experiments that you have done and that unfortunately did not fit into the main text. The Appendix G could be a very useful practical contribution and I feel like it is a bit of a shame that it's so hidden -- but I of course understand that you wanted to have a nice narrative arc in the main text and that the page limit is a concern. In particular, the finding that RL-aligned models are better teachers for base models is surprising and noteworthy (Appendix G.3).", "weaknesses": "- I'm not very convinced by your explanation of the poor ICL performance. Looking at the average ICL performance in Figure 3(d), it looks like distillation actually leads to very good ICL performance! It is able to match standard pretraining with just 50% of the training samples. However, the performance suddenly drop when trained on 100% training samples. This behavior is never explained in the paper. The theory of induction heads that the authors propose does not match the data -- it says that the distilled model should always be worse (because it does not robustly learns how to copy), but that is not what the results show.\n- The paper claims multiple times that the teachers of Gemma-3 and Llama-3.2 were trained on \"far more data than the students\", and the fundamental difference of this paper is that the students and teachers are trained on the same amount of data. I'm not sure this is really true in practice. The Gemma-3 students were trained on up to 14T tokens, but there is no information about the teacher. Llama-3.2 models were trained on 9T tokens with teachers trained on 15T tokens. But note that the 3.2 models are very small (1B and 3B parameters), so 9T tokens is way past the compute optimality. So in practical terms, I believe that training can be thought as approximately \"isoData\", it is certainly an overstatement to say \"far more data than students\" (line 39), \"significantly more data than students\" (line 139) or \"way more data than [students]\" (line 467).", "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": "OR9xCK7v40", "reviewer": "Reviewer_FL6P", "summary": "The goal of this work is to investigate the connections between distilled pre-training to in-context learning and test-time scaling. To this The larger goal of the work is to probe how distilled pretraining (DPT) ties into in-context learning and test-time scaling. The paper shows that DPT gives noticeably better test-time scaling, yet hurts in-context learning, especially tasks usually handled by induction heads. The authors argue this comes from a bias toward high-entropy tokens during distillation, which down-weights the low-entropy bits that teach pattern copying and other ICL skills. They back this up with a clean sandbox using a bigram model that isolates a single governing factor behind both effects. Finally, they offer a simple fix: token-routing, where you drop the distillation loss on x% of lowest-entropy tokens by \nwhich recovers in-context performance.", "questions": "1. While I find the experiments and analysis solid, I’m still a bit unsure about some of the practical aspects of DPT. It would help to know roughly how its compute requirements (excluding the teacher’s pretraining cost) compare to standard pre-training. The teacher’s soft-label generation likely adds some overhead. If we account for that, how much performance would an equivalent amount of standard pretraining compute buy, and how does that compare to DPT performance?\n2. There also seem to be some easy extensions worth discussing. For instance, since the paper drops the distillation loss for x% of low-entropy tokens, one could imagine dropping the hard-label loss for x% of high-entropy tokens from the teacher, which might further improve test-time scaling. Did the authors try or consider this? Similarly, could a similar “token-routing” idea help in standard hard-label pretraining, e.g., using the model’s own entropy to guide which tokens to update or skip?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "The paper is very clearly written. The main idea, experiments, and analysis all feel solid and well thought out. The bigram model is a nice touch that builds intuition without overcomplicating things. The proposed fix is simple, sound, and gives clear gains, which makes it both practical and easy to appreciate.", "weaknesses": "Please see Questions section.", "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": "x5TYQHo40t", "reviewer": "Reviewer_9LRr", "summary": "The paper studies the impact of distilled pretraining on test-time scaling and in-context learning.\nIt is found that, compared to the usual pretraining, distilled pretraining enhances test-time scaling while hurting in-context learning.\n\nThe main reason for these observations are that distillation causes diversity in the output distribution: in-context learning requires memorization; diversity hurts this, while test-time scaling is enhanced by this diversity.\nThe paper demonstrates this intuition via a bigram model.\nFurthermore, the paper proposes a simple method to overcome the drop in in-context learning.", "questions": "1. Since distillation requires extra compute from the teacher to calculate the logits, and compute is an important consideration for LLM that requires a lot of compute, why do you think that it is not correct to compare the perfomance at the iso-compute setting, as written in line 478? Particularly, the compute would be high when the whole soft label distribution is used in the calculation of distillation loss.\n\n2. Is there any in-context learning task that does not rely much on memorization, such that there is exceptions to the claim that distilled pretraining degrades in-context learning?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1. The paper is well written, with clear motivations and convincing arguments for the behavior of distilled p retraining. A solution is also provided, which promotes the adoption of distilled pretraining. \n2. The paper also provides complementary details, e.g., changing the k of the top-k logit for distillation, to back up the arguments.", "weaknesses": "1. I think that saying that distilled pretraining is better at test time scaling is an overclaim. At pass@1 it is worse (figure 1a) and only when at large k of pass@k the performance is better (otherwise, we could say that the base model is better than an RL-trained model at reasoning, as in [2504.13837]).\nIt is just that the distilled model can sample from a more diverse distribution as claimed in the paper, which does not imply better test time scaling in general.\nTo say that it is better at test time scaling I would also expect a comparison in terms of an accuracy vs thinking time plot (e.g., those presented in figure 1 of [2501.19393]).\n\n2. With the above in mind and the memorization property of in-context learning, the presented results are perhaps not surprising, as distilled pretraining provides soft labels which are diverse in its supervision, and this naturally leads to better test-time scaling (more diverse sampling) and worse in-context learning (diversity affects direct memorization).\n\n3. Another possible improvement would be vary the student/teacher model size to check the validity/make more interesting observation from it.", "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": "oW3O04xTnw", "reviewer": "Reviewer_cYZa", "summary": "This paper explores the impact of distillation-based pretraining on various aspects of model performance, namely in-context learning and test-time scaling. The authors conduct large scale experiments and discuss the pros and cons of distillation in the context of these two aspects. They also propose a theoretical support of their empirical observations, and experiment in a controlled synthetic setup to validate their findings. \n\nTheir conclusions underline that pretraining with distillation leads to models with strong abilities in terms of test-time scaling (as measured with pass@k experiments), but with weaker ICL performance. They propose to fix this by dropping the distillation objective on low-entropy tokens to recover strong copying mechanisms, which partially bridges the gap with standard pretrained counterparts.", "questions": "- How do you distinguish between high-entropy and low-entropy in section 4?\n- The scale of teacher models used in distillation is usually much larger. These largers models should be better at predicting low-entropy probabilities accurately. Do you think this will have an impact on the inductive abilities of student models? \n- You compare your method to self-distillation (in the IsoData setup). What justifies that comparison? Do you think strict self-distillation would result in similar outcomes?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "This paper is well-written and offers an interesting insight into the interplay between distillation and LLM inference behaviors (context retrieval abilities and test-time scaling).\n- **Interesting and novel interplay analysis**: the analysis of the effect of distillation on key strengths of LLM (ICL and test-time scaling) at large scale is novel to the best of my knowledge. It provides relevant insights on the strenghts and weaknesses of distillation beyond pure language modeling performance as measured by perplexity. It is also relatively well grounded in theory and explained by the authors in a convincing way.\n- **Sound experimental design**: the experiments are conducted in a controlled setup. The IsoData comparisons disentangles the effects of additional data and of distillation.", "weaknesses": "This paper lacks some crucial details and some experimental results are analyzed in ways that leave room for confounding factors. It also frames terms in confusing ways and lack proper background for some parts\n- **Training reproducibility & related conclusions**: The authors only describe the training dataset they used in the Appendix, and do so in a very broad way. The paper is also not mentioning many details such as (in order of relevance): sequence length, architecture, batch sizes, or learning rate warm-up strategy. These details are crucial for reproducibility but also to better analyze the ICL results that rely on long-context experiments. For instance, we currently have no clue whether the Needle-in-the-haystack experiments fit within training context length or not. It is also impossible to tell whether 36% for the 8B teacher model is a strong performance. It would be interesting to see training dynamics and to compare the model performance to pretrained counterparts or at least to similar open models (e.g. OLMo-2) at equivalent training stages.\n- **Validity of some of the IsoData conclusions**: The IsoData setup is important to assess the efficiency of distillation in a dataset-equivalent regime between teacher and student training. However, the authors frame this setup in these terms: \"student is shown the same amount of data as the teacher\", which may not be exactly fair. Although the student only trains once on the 1T dataset, it immediately benefits from an estimate of the full training distribution through the teacher. As a result, at every training step, the model is not only shown the current sample, but also a compressed view of all related samples to come through the distillation loss. It can thus be wondered if working in a more data-constrained setup (e.g. having several training epochs) could affect the outcomes.\n- **Missing discussion on the impact of temperature(s)**: A crucial parameter in the distillation setup is the distillation temperature. The authors mention a grid search, but do not report the corresponding results, which matter both for the final performance of the models and for the impact of temperature on ICL performance and the induction heads. As a result, it seems possible that a different temperature would lead to a different trade-off between benchmark scores and ICL / diversity abilities of the model. Even though this test is costly at scale, it could be at least conducted on the synthetic data experiment. It is also unclear what the generation parameters are the pass@k experiments, and whether they affect the pretrained and distilled models differently.\n - **Inaccurate/unclear terms**: My understanding of In-Context Learning (ICL) is that it mostly entails methods that design prompts that steer the model towards the desired behavior, the most basic example being few-shot evaluation where correct examples of a task are provided in the prompt. Referring to abstractive/extractive context-based QA or in-context retrieval (NIAH) tasks as ICL tasks does not fit this definition. The term \"in-weight learning\" would also benefit from contextualization and a corresponding reference.", "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)