# Knowledge Distillation During Mid-Training Favors Reasoning over Factual Recall

> Knowledge distillation during mid-training boosts reasoning but slows factual recall; entropy-based token routing (SWITCH DISTILLATION) preserves both.

- **Source:** [arXiv](https://arxiv.org/abs/2609.01532)
- **Published:** 2026-09-12
- **Permalink:** https://picx.dev/p/CyhRzw
- **Whiteboard:** https://picx.dev/p/CyhRzw/image

## Summary

# Knowledge Distillation During Mid-Training Favors Reasoning over Factual Recall

**Authors:** Jacqueline He, Howard Yen, Shuyue Stella Li, Margaret Li, Hanqing Zeng, Yinglong Xia, Benyu Zhang, Zhuokai Zhao, Qiang Zhang, Pang Wei Koh, Luke Zettlemoyer, Wen-tau Yih (Meta AI, University of Washington, Princeton University)

## Summary (Overview)

- **Key Finding:** Forward KL knowledge distillation (KD) behaves fundamentally differently during mid-training compared to pre-training. While KD improves both reasoning and factual recall during pre-training, it slows factual recall acquisition during mid-training despite continued reasoning gains—a phenomenon termed the **reasoning–recall tradeoff**.
- **Explanatory Mechanism:** The tradeoff arises from an interaction between (i) asymmetric teacher confidence across data domains (teachers are more confident on procedural than knowledge-intensive data), (ii) the student's evolving knowledge state (low-entropy facts are acquired early, leaving unresolved facts concentrated in high-entropy regions), and (iii) the distillation objective's attenuation of ground-truth supervision on high-entropy tokens.
- **Proposed Method:** **SWITCH DISTILLATION**, a simple mid-training objective that routes tokens between reverse-KL distillation and cross-entropy based on teacher predictive entropy, using a lightweight routing signal computed from teacher logits already required for KD.
- **Results:** SWITCH DISTILLATION achieves 1.61–1.71× the reasoning performance and 1.13–1.19× the knowledge & commonsense performance relative to standard NTP, while preserving 96.7–96.8% of factual recall. These gains persist after post-training, with the factual recall gap closing entirely.
- **Robustness:** The reasoning–recall tradeoff is robust across teacher sizes (1B, 7B, 13B), KL directions (forward and reverse), and interpolation coefficients—no standard KD objective Pareto-dominates NTP during mid-training.

## Introduction and Theoretical Foundation

### Background and Motivation

Modern language model (LM) development increasingly employs a **mid-training stage** between pre-training and post-training, where self-supervised next-token prediction continues on a smaller, high-quality corpus curated to improve capabilities such as factuality, reasoning, coding, and instruction following. Because mid-training uses far fewer tokens than pre-training, extracting maximal learning signal from each token is critical.

Knowledge distillation (KD) (Bucilu et al., 2006; Hinton et al., 2015) offers a natural approach by augmenting ground-truth supervision with the richer predictive distribution of a stronger teacher. However, KD has been studied almost exclusively in pre-training and post-training settings, leaving its behavior during mid-training largely unexplored.

### Theoretical Foundation

**Next-token prediction (NTP).** Given a token sequence $\mathbf{x} = (x_{1}, \ldots, x_{N})$ and an auto-regressive language model distribution $p_{\theta}$, standard next-token prediction minimizes the expected cross-entropy loss:

$$
\mathcal {L} _ {\mathrm{CE}} = \mathbb {E} _ {(\mathbf {x}, n)} \left[ - \log p _ {\theta} (x _ {n} \mid x _ {<   n}) \right],\tag{1}
$$

**Knowledge distillation.** Logit-based KD trains a student to match the output distribution of a vocabulary-compatible teacher:

$$
\mathcal {L} _ {\mathrm{KD}} = (1 - \alpha) \mathcal {L} _ {\mathrm{CE}} + \alpha \mathcal {L} _ {\mathrm{KL}},\tag{2}
$$

where $\alpha\in[0,1]$ controls the distillation strength. The forward KL (FKL) divergence is the standard instantiation:

$$
\mathcal {L} _ {\mathrm{FKL}} = \tau^ {2} \mathbb {E} _ {(\mathbf {x}, n)} \left[ \sum_ {v \in \mathcal {V}} p _ {T} ^ {\tau} (v \mid x _ {<   n}) \log \frac {p _ {T} ^ {\tau} (v \mid x _ {<   n})}{p _ {S} ^ {\tau} (v \mid x _ {<   n})} \right].\tag{3}
$$

Reverse KL (RKL) divergence, which discourages student mass on the teacher's low-probability regions, is also considered:

$$
\mathcal {L} _ {\mathrm{RKL}} = \tau^ {2} \mathbb {E} _ {(\mathbf {x}, n)} \left[ \sum_ {v \in \mathcal {V}} p _ {S} ^ {\tau} (v \mid x _ {<   n}) \log \frac {p _ {S} ^ {\tau} (v \mid x _ {<   n})}{p _ {T} ^ {\tau} (v \mid x _ {<   n})} \right].\tag{4}
$$

## Methodology

### Experimental Setup

- **Training regimes:** Built on the open-source OLMo-2 ecosystem. Pre-training initializes 1B students from random weights and trains for 100B tokens. Mid-training initializes from the OLMo-2 1B Stage 1 checkpoint (pre-trained on 4T tokens) and continues for 60B tokens on Dolmino Mix 1124.
- **Teacher models:** OLMo-2 1B Instruct, 7B Instruct, and 13B Instruct (post-trained models).
- **Evaluation:** Standardized OLMES harness, grouping tasks into REASONING, FACTUAL RECALL, KNOWLEDGE & COMMONSENSE, and INSTRUCTION FOLLOWING.

### Key Analyses

1. **Teacher confidence asymmetry:** Computed token-level predictive entropy of teachers across data domains. Found procedural domains (math, instruction-following) exhibit substantially lower entropy than knowledge-intensive ones, and lower entropy correlates with higher teacher top-1 agreement with ground-truth tokens (Figure 3).

2. **Teacher entropy predicts factual acquisition:** Tracked factual acquisition across intermediate NTP checkpoints on FACTUAL RECALL tasks (TriviaQA, Natural Questions, SimpleQA). Stratified examples into teacher-entropy quintiles; found that by end of pre-training, students learned 67% of facts in the lowest-entropy quintile (Q1) vs. only 5% in the highest (Q5) (Figure 4).

3. **KD attenuates factual supervision:** Measured (i) teacher probability on ground-truth tokens, (ii) gradient norm ratios of gold-token logits under KD vs. NTP, and (iii) downstream factual recall per entropy quintile. Found both FKD and RKD increasingly attenuate ground-truth supervision as teacher entropy rises, reaching approximately $0.5\times$ NTP for the highest-entropy facts (Figure 5).

### SWITCH DISTILLATION

Let teacher predictive entropy at token position $n$ be:

$$
H _ {n} = - \sum_ {v \in \mathcal {V}} p _ {\mathrm{T}} ^ {(\tau)} (v \mid x _ {<   n}) \log p _ {\mathrm{T}} ^ {(\tau)} (v \mid x _ {<   n})\tag{5}
$$

Tokens are routed to distillation based on the lowest $q\%$ of in-batch tokens by $H_n$, defining $S_{q} = \left\{ n : H_{n} \leq \text{Quantile}_{q}(\{H_{n}\}) \right\}$. The objective is:

$$
\mathcal {L} ^ {\text { SWITCHDist }} = \tau^ {2} \frac {1}{| \mathcal {S} _ {q} |} \sum_ {n \in \mathcal {S} _ {q}} \mathrm{RKL} \Big (p _ {\mathrm{S}, n} ^ {(\tau)} \| p _ {\mathrm{T}, n} ^ {(\tau)} \Big) + \frac {1}{| \bar {\mathcal {S}} _ {q} |} \sum_ {n \notin \mathcal {S} _ {q}} \mathcal {L} _ {\mathrm{CE}, n},\tag{6}
$$

This adds negligible computation beyond standard online KD—only entropy and quantile calculations on teacher logits already required.

## Empirical Validation / Results

### Mid-Training Results (Table 1)

With a 7B teacher, SWITCH DISTILLATION achieves the strongest REASONING (44.7% macro-average vs. 26.1% for NTP) while remaining closest to NTP on FACTUAL RECALL (29.3% vs. 30.3%). It also achieves the strongest KNOWLEDGE & COMMONSENSE (49.3%). Similar trends hold with the 13B teacher.

| T | Method | Reasoning (macro) | Factual Recall (macro) | Knowledge & Comm. (macro) |
|---|---|---|---|---|
| N/A | NTP | 26.1 | 30.3 | 38.4 |
| 7B | FKD | 36.1 | 29.0 | 46.4 |
| 7B | RKD | 38.8 | 28.5 | 47.1 |
| 7B | TRKD | 32.2 | 28.6 | 44.4 |
| 7B | **SD** | **44.7*** | **29.3** | **49.3*** |
| 13B | FKD | 32.1 | 29.0 | 43.5 |
| 13B | RKD | 37.0 | 28.6 | 45.5 |
| 13B | TRKD | 29.7 | 28.8 | 42.6 |
| 13B | **SD** | **42.1*** | **29.3** | **46.5** |

### Post-Training Results (Table 2)

After applying a standard post-training pipeline (SFT, DPO, RLVR1, RLVR2), SWITCH DISTILLATION retains its gains: reasoning improves from 44.7% to 50.6% (7B teacher) and from 42.1% to 48.0% (13B teacher). The factual recall gap closes entirely, with SWITCH DISTILLATION finishing with the highest FACTUAL RECALL macro-average. Instruction following also improves (69.5% vs. 62.1% for NTP).

### Ablations (Table 3)

- **KL direction:** Replacing RKL with FKL modestly reduces REASONING (-2.9%) and KNOWLEDGE & COMMONSENSE (-1.4%).
- **Routing signal:** Teacher entropy outperforms alternatives (teacher-correct routing, oracle domain routing, random routing).
- **Supervision objective:** Always CE has little effect; teacher top-1 labels improve FACTUAL RECALL (+1.3%) but hurt REASONING (-6.4%).

## Theoretical and Practical Implications

### Theoretical Implications

1. **KD is not stage-agnostic:** The standard KD formulation that works during pre-training exhibits fundamentally different behavior during mid-training. This challenges the assumption that distillation objectives transfer uniformly across training stages.

2. **Teacher confidence as a supervision quality signal:** Teacher predictive entropy reliably predicts supervision quality (correctness) both across and within data domains, providing a principled basis for routing decisions.

3. **Student knowledge state matters:** The interaction between teacher confidence and the student's evolving knowledge state determines distillation effectiveness. Unlearned facts at mid-training initialization disproportionately receive weak teacher supervision, explaining the reasoning–recall tradeoff.

### Practical Implications

1. **Stage-aware objectives:** Training objectives should be designed with the training stage in mind. SWITCH DISTILLATION demonstrates that simple, computationally lightweight adaptations can yield substantial improvements.

2. **Data efficiency:** SWITCH DISTILLATION improves how supervision is extracted from each observed token, complementing data-centric approaches to the growing scarcity of high-quality training data.

3. **Post-training synergy:** Mid-training gains from SWITCH DISTILLATION persist through post-training, making it a practical drop-in objective for frontier LM pipelines.

## Conclusion

This work uncovers a robust reasoning–recall tradeoff unique to mid-training: standard knowledge distillation improves reasoning while slowing factual acquisition relative to NTP. The tradeoff is explained through the interaction between teacher confidence, student learning dynamics, and the distillation objective. SWITCH DISTILLATION—which routes tokens between reverse-KL distillation and cross-entropy using teacher predictive entropy—substantially mitigates this tradeoff across teacher sizes and retains its gains after post-training.

The findings suggest that objectives themselves ought to be stage-aware. While this work focuses on mid-training, the principle may extend to other phases where the student's knowledge has substantially evolved, such as late-stage or continual pre-training. Future work should investigate stage-aware optimization methods for data-efficient language modeling.

---

_Markdown view of https://picx.dev/p/CyhRzw, served by PicX — AI-generated visual whiteboard summaries of research papers._