PA1-OPD: One-Step Prefix-Augmented On-Policy Distillation for Language Model Pretraining

Zhenhe Wu; Jian Yang; Yaping Jin; Xinlai Kang; Jie Hu; Hanting Chen

cs.CLOct 14, 2026

PA1-OPD: One-Step Prefix-Augmented On-Policy Distillation for Language Model Pretraining

Authors: Zhenhe Wu, Jian Yang, Yaping Jin, Xinlai Kang, Jie Hu, Hanting Chen

Organizations: Huawei Technologies · Beihang University

Abstract

Knowledge distillation (KD) improves language-model pretraining by matching teacher distributions at corpus prefixes, but provides no direct supervision at contexts induced by student predictions. We introduce PA1-OPD, which augments corpus-path distillation with teacher feedback after one student-sampled transition per selected prefix. This expands the supervised states while retaining next-token prediction and corpus-path KD. Data-anchored roots support parallel sampling, and differentiable tree attention shares prefix computation without changing inference. In two continued-pretraining settings, full coverage improves held-out cross-entropy over KD across C4, PG-19, and Pile at every evaluated checkpoint. In Stage2, selecting the 10% of roots with highest student surprisal adds 48.85% to KD's Macro cross-entropy gain over next-token training. Gains remain when branch supervision excludes successor states already on the corpus path, supporting the value of feedback beyond observed contexts. Additional depth or breadth improves aggregate results at unchanged per-term supervision weight, with benefits varying by corpus. These findings support local student-induced supervision as a complement to pretraining distillation.

Figures & tables

Figure 1: PA1-OPD adds teacher supervision after student-sampled transitions from corpus prefixes. Tree attention shares each prefix and isolates branches. Sampling is detached; student prefix states remain differentiable.

Stage1 initialization; Stage1 mixture
Tokens	Method	C4		PG-19		Pile		Macro
Tokens	Method	CE $\downarrow$	PPL $\downarrow$	CE $\downarrow$	PPL $\downarrow$	CE $\downarrow$	PPL $\downarrow$	CE $\downarrow$	PPL $\downarrow$
0	Initial	3.018891	20.468578	2.908693	18.332821	2.489751	12.058270	2.805778	16.953223
2.5B	NTP	2.967517	19.443581	2.848100	17.254962	2.433926	11.403570	2.749848	16.034038
	KD	2.963171	19.359270	2.842240	17.154150	2.428178	11.338206	2.744530	15.950542
	Full PA1	2.962772	19.351543	2.841985	17.149778	2.427837	11.334337	2.744198	15.945219

Table 1: Learning curves at matched training-token budgets. Lower CE and PPL are better; Macro averages each metric over the three datasets. extra/KD measures additional improvement beyond KD relative to KD’s gain over NTP, not percentage reduction in CE or PPL.

Figure 2: Stage2 root-coverage comparison after 1B continuation tokens with student-surprisal selection. Curves report CE extra/KD ; higher is better.

Panel of Figure 3 — Figure 3: Positional distribution of selected roots. Cell values give selection frequencies among eligible roots in each position bin, with a shared color scale across settings and checkpoints.

Variant	C4	PG-19	Pile	Macro
Full PA1	30.63	17.41	26.91	26.63
Teacher Top50	33.52	19.30	28.77	29.07
Student Top50	43.96	22.15	33.64	36.18
Drop coincident	44.37	26.58	34.80	37.80

Table 2: Stage2 root-selection and coincident-state ablations as CE extra/KD percentages (higher is better). Drop coincident uses Student Top50 and removes only losses on states already present on the corpus path.

Panel of Figure 4 — Figure 4: Full-root branch expansion after 1B continuation tokens. Each branch term retains weight 0.2. CE extra/KD uses matched NTP and KD references with KD weight 0.5; higher is better.

Language Model Pretraining

Language model pretraining teaches models language patterns and knowledge from large text corpora before they are adapted to specific tasks.

See all 441 papers

Tail-Aware Top- $k$ On-Policy Distillation
Aug 12, 2026Huipeng Huang, Hongxin WeiLanguage Model Distillation On-Policy Distillation
Department of Statistics and Data Science, Southern University of Science and Technology
Fast and Effective On-policy Distillation from Reasoning Prefixes
Feb 16, 2026Dongxu Zhang, Zhichao Yang, Sepehr Janghorbani +6Language Model Distillation On-Policy Distillation
Optum AI
Beyond Imitation: Filtering On-Policy Distillation by Reasoning Progress
Aug 19, 2026Chen Yang, Haiyuan Wan, Rengrong Xiong +2Language Model Distillation On-Policy Distillation
The Hong Kong University of Science and Technology (Guangzhou) · Tsinghua University · Zhejiang University +1

Aug 12, 2026J/K move · Enter open · S save

Tail-Aware Top- $k$ On-Policy Distillation

Huipeng Huang, Hongxin Wei

Department of Statistics and Data Science, Southern University of Science and Technology

Language Model Distillation On-Policy Distillation Knowledge Distillation

Open paper AskPDF arXiv

On-policy distillation (OPD) has emerged as an effective paradigm for transferring knowledge between language models, where a student is trained to align its next-token distribution with the teacher's along its own trajectories. To provide dense supervision at tractable cost, many works minimize the reverse Kullback-Leibler (KL) divergence between the student and teacher's normalized distributions over the teacher's top-

k

tokens. However, this normalized objective discards the information about tail probability: the total probability outside the teacher's top-

k

tokens. As a result, the optimization can steadily increase the student's tail probability and entropy, empirically degrading downstream accuracy. To address this issue, we propose Tail-Aware Top-

k

OPD (TA-OPD), a novel distillation method that restores the missing tail probability signal. In particular, TA-OPD minimizes the reverse KL divergence over the top-

k

tokens plus a tail token that carries the tail probability. In effect, TA-OPD better aligns the student's next-token distribution with the teacher's, preventing the increase in tail probability and entropy caused by top-

k

normalization. Extensive experiments demonstrate the superiority of TA-OPD, improving Avg@8 by up to 8.05 points on common benchmarks. Our code is available at https://github.com/HuipengHuang/TA-OPD.

Abstract

Language Model Pretraining

Explore similar work

Tail-Aware Top-kkk On-Policy Distillation

Fast and Effective On-policy Distillation from Reasoning Prefixes

Beyond Imitation: Filtering On-Policy Distillation by Reasoning Progress

Tail-Aware Top- $k$ On-Policy Distillation