PA1-OPD: One-Step Prefix-Augmented On-Policy Distillation for Language Model Pretraining
On the map: Knowledge distillation → Knowledge distillation for LLMs also fits: On-policy and self-distillation for LLMs
Abstract (arXiv)
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.
Key facts. No citing papers in the map yet (data through 14 Oct 2026). Placed in Knowledge distillation for LLMs, one of 12 directions in Knowledge distillation. Builds on 24 of its 38 references that are in the map. First posted on arXiv 14 Oct 2026.
Builds on
24 of 38 references.
Closest ancestors · 2
Same direction, newest first
- 2024MiniPLM: Knowledge Distillation for Pre-Training Language Models22 Oct
- 2023On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes23 Jun
Other references · 22
- 2019Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformercited by 14,599
- 2015Distilling the Knowledge in a Neural Networkcited by 10,324
- 2019DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lightercited by 4,307
- 2020The Pile: An 800GB Dataset of Diverse Text for Language Modelingcited by 2,411
- 2024Gemma 2: Improving Open Language Models at a Practical Sizecited by 1,998
- 2019TinyBERT: Distilling BERT for Natural Language Understandingcited by 1,167
- 2024Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Headscited by 843
- 2015Scheduled Sampling for Sequence Prediction with Recurrent Neural Networkscited by 803
- 2016Sequence-Level Knowledge Distillationcited by 801
- 2019Compressive Transformers for Long-Range Sequence Modellingcited by 726
- 2023SpecInfer: Accelerating Generative Large Language Model Serving with Tree-based Speculative Inference and Verificationcited by 365
- 2024Better & Faster Large Language Models via Multi-token Predictioncited by 360
10 more
- 2023Think before you speak: Training Language Models With Pause Tokenscited by 264
- 2026Entropy-Aware On-Policy Distillation of Language Modelscited by 175
- 2024Speculative Knowledge Distillation: Bridging the Teacher-Student Gap Through Interleaved Samplingcited by 90
- 2025DistiLLM-2: A Contrastive Approach Boosts the Distillation of LLMscited by 84
- 2020Autoregressive Knowledge Distillation through Imitation Learningcited by 74
- 2026Fast and Effective On-policy Distillation from Reasoning Prefixescited by 53
- 2024Cut Your Losses in Large-Vocabulary Language Modelscited by 40
- 2025Table-R1: Region-based Reinforcement Learning for Table Understandingcited by 37
- 2026Prefix-Guided On-Policy Distillation: Mining Golden Trajectories from Rolloutscited by 10
- 2026Cache-Aware Joint Router Adaptation for Memory-Efficient MoE Inferencecited by 3
Built on by
Not cited yet.
Similar, not cited
Close in wording, no citation either way
- 2026Breaking the Tokenizer Barrier: On-Policy Distillation across Model Families
- 2026Rethinking On-Policy Distillation of Large Language Models: Phenomenology, Mechanism, and Recipe
- 2025TAID: Temporally Adaptive Interpolated Distillation for Efficient Knowledge Transfer in Language Models
- 2024DistiLLM: Towards Streamlined Distillation for Large Language Models
- 2026Understanding Knowledge Distillation in Post-Training: When It Helps and When It Fails
- 2026Trajectory-Refined Distillation
- 2025Less is More: Selective Reflection for Compatible and Efficient Knowledge Distillation in Large Language Models
- 2026Distilling What Matters: Confidence-Aware Selective Distillation for Large Language Models