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

1 Introduction

• Conventional pretraining KD applies teacher targets at corpus prefixes, leaving successor states induced by student predictions without direct supervision.One sampled transition tests whether local feedback beyond the corpus trajectory helps without generating complete continuations.Source: §1 · p. 1
• PA1-OPD samples one token at selected corpus prefixes and applies teacher supervision to the student’s next-token distribution at the resulting successor state.This branch objective complements next-token prediction and corpus-path KD while keeping its root anchored to observed data.Source: §1 · pp. 1–2
• Data-anchored roots allow parallel sampling, while tree attention shares prefixes and keeps branches independent with differentiable student-prefix states.Coincident-state reuse and chunked KL reduce reevaluation and temporary probability storage; student-surprisal selection controls feedback placement.Source: §1 · pp. 1–2, Figure 1 · p. 2
• Full coverage improves over KD across all three evaluation corpora in both settings; Stage2 Top10 adds 48.85% to KD’s CE gain over next-token training on Macro.Stage1 favors full coverage over Top10, and gains persist without branch losses at coincident corpus states.Source: §1 · p. 1

2.1 Pretraining Distillation

• The setup studies autoregressive language-model pretraining through checkpoint continuation with a student and frozen teacher that share a tokenizer and vocabulary.The models may differ in depth and width, and inference uses only the student.Source: §2.1 · p. 2

2.2 Corpus and Student-Induced States

• For each sequence, h_b,i denotes the prefix through position i, and R_b collects valid next-token prediction positions.The observed next token at each position supplies the corpus-path target.Source: §2.2 · p. 2
• NTP predicts the observed next token from each corpus prefix, while RealKD distills teacher predictions on the observed data trajectory.The RealKD objective is normalized over the valid prediction positions.Source: §2.2 · p. 2
• The distributions used for distillation are softmax outputs with logits divided by temperature τ.The formulation defines teacher and student distributions for a context h.Source: §2.2 · p. 2
• The sampling temperature τ_s and token concatenation define how a student-sampled token extends its corpus-prefix root.Sampling over eligible prefixes induces a one-step state distribution.Source: §2.2 · pp. 2–3
• The transition follows the temperature-adjusted student policy from a data-anchored root, and its successor may coincide with a corpus state.Coincidence occurs when the sampled token matches the observed next token.Source: §2.2 · p. 3
• PA1 adds local student-induced supervision by sampling a single transition from an eligible corpus prefix.The current student is held fixed during sampling, and the sampled token is concatenated to the corpus prefix to form the successor state.Source: §2.2 · pp. 2–3

3.1 One-Step Distillation

• At each eligible corpus root, the branch matches teacher and student next-token distributions after one sampled token, while the observed corpus token remains the NTP target.The branch teacher target is the full-vocabulary predictive distribution at the successor state, not the sampled token’s probability at its root.Source: §3.1 · p. 3
• The branch objective uses one categorical sample per eligible root, combines with corpus-path losses, and holds the teacher and sampling policy fixed during differentiation.Gradients still pass through the student computations for branch prediction, including the shared prefix; NTP and RealKD normalize by NR, and PA1 by NB.Source: §3.1 · p. 3
• Deterministic top-K selection preserves nominal branch-loss scale but targets a different state distribution than full coverage and is not an unbiased full-root estimate.Sequences without eligible roots are omitted.Source: §3.1 · p. 3

3.2 Prefix Sharing with Tree Attention

• Tree attention keeps real tokens on their causal prefixes, lets each branch read its sampled prefix, and isolates branches from other roots.Real-to-branch attention is blocked, so branch tokens cannot alter real-path representations.Source: §3.2 · pp. 3–4
• Branch position IDs place each sampled token immediately after its corresponding corpus token, making tree execution match independent extended-prefix evaluation for fixed samples.The equivalence argument also states that differentiable sharing preserves summed parameter gradients up to numerical precision.Source: §3.2 · p. 4
• The student reuses layer-wise prefix key/value caches for sampled branches, and differentiable reuse lets branch losses update both new-token computations and shared prefix states.Teacher and student keep separate caches; only the student requires gradients.Source: §3.2 · p. 4
• Fused teacher evaluation processes real and branch tokens in one layer traversal, while block-sparse attention and chunked branch queries limit intermediate storage.Fusion reduces repeated traversals and kernel dispatches but does not remove computation for the added states.Source: §3.2 · p. 4

3.3 Objective-Preserving Execution

• When a sampled token matches the next corpus token, the method reuses successor logits with student gradients intact and evaluates only unmatched branches.Chunked KL retains full-vocabulary normalization, so reuse and chunking preserve the sampled objective.Source: §3.3 · p. 4

4.1 Experimental Setup

• The study compares NTP, corpus-path KD, and PA1-OPD at matched training-token budgets, with Full PA1 using every eligible root and Student Top10 selecting the 10% highest-surprisal roots.Each condition uses one run; CE is reported in nats, and Macro scores average the three datasets.Source: §4.1 · p. 4

4.2 Main Results

• Full PA1 improves CE and PPL over KD on all three datasets at every evaluated checkpoint in both settings.After 10B tokens, Macro CE is lower by 2.87×10−4 in Stage1 and 2.29×10−4 in Stage2, adding 4.65% and 27.69% to KD’s gain over NTP, respectively.Source: Table 1 · p. 5, §4.2 · p. 6
• Stage2 Student Top10 ends at 2.531635 Macro CE versus 2.532039 for KD, adding 48.85% extra/KD, but Stage1 Top10 ends at −5.22% extra/KD.Full PA1 remains beneficial in Stage1; initialization and mixture change together, so coverage should be chosen for the continuation setting.Source: §4.2 · p. 6, Table 2 · p. 6
• Student Top10 adds at most approximately 10% to each model’s corpus token-state evaluations before reuse, but those counts cannot establish compute efficiency.The counts exclude model-size differences, attention, and backward computation.Source: §4.2 · p. 6

5 Analysis

• The analysis tests coincident-state removal, root selection, positional coverage, and added branch depth or breadth under controlled protocols.PAx samples sequentially, whereas dual PA1 draws two independent one-step samples per root; each extension doubles sampled tokens and supervised branch states per root before reuse.Source: §5 · p. 6, §5.4 · pp. 7–8

5.1 Supervision Beyond Corpus States

• Removing branch KL terms at coincident corpus states still improves on KD on every dataset, with Macro CE 2.531165 versus 2.531173 when those terms are retained.The small difference does not establish an advantage for removal; feedback at noncoincident student-induced states remains useful.Source: §5.1 · p. 6

5.2 Allocating Feedback Across Roots

• In Stage2, Macro extra/KD rises from 26.63% with full coverage to 40.04% with Top10, while Stage1 favors broader coverage.At matched 50% coverage, surprisal raises Macro extra/KD from 29.07% to 36.18% versus teacher advantage; C4 favors the smallest tested budget, while PG-19 and Pile favor broader coverage.Source: §5.2 · pp. 6–7

5.3 Coverage Across Context Lengths

• At the final checkpoint, the first position bin's selection rates are 15.57% in Stage1 and 19.42% in Stage2, yet it accounts for only 4.86% and 6.06% of selected roots.Latter-half selection rates remain 9.61% and 9.47%; aggregate positional preferences change little from initialization, although individual roots may change with the student.Source: §5.3 · p. 7, Figure 3 · p. 7

5.4 Adding Depth and Breadth

• After 1B continuation tokens, PA2 and dual PA1 raise Stage1 extra/KD from 28.99% to 40.47% and 40.94%, respectively, while both improve Macro CE in Stage2.The extensions retain each branch term's weight while adding another supervision signal; returns vary by corpus, and no general preference for depth over breadth is established.Source: §5.4 · pp. 7–8, Figure 4 · p. 8

6 Related Work

• Knowledge distillation transfers predictive distributions, sequences, and representations, while pretraining applications study compression, data refinement, and compute allocation (Hinton et al., 2015; Kim and Rush, 2016; Sanh et al., 2019; Jiao et al., 2020).The passage also situates these applications in work by Gemma Team (2024), Muralidharan et al. (2024), Gu et al. (2025), and Busbridge et al. (2025).Source: §6 · p. 8
• DAgger and scheduled sampling motivate learner-induced states, while ImitKD, GKD, and MiniLLM apply related principles to distillation (Ross et al., 2011; Bengio et al., 2015; Lin et al., 2020; Agarwal et al., 2024).Subsequent work explores objectives and data reuse, interleaved sampling, and token or rollout-prefix selection (Ko et al., 2024, 2025; Xu et al., 2025; Huang et al., 2025; Zhang et al., 2026; Zhao et al., 2026).Source: §6 · p. 8
• PA1-OPD adapts shared-prefix trees for differentiable student computation and applies forward-KL supervision after local student transitions from corpus prefixes during pretraining.The passage contrasts this with work that adds computation or prediction targets and notes that chunked KL reduces probability workspace while preserving the teacher distribution.Source: §6 · p. 8

7 Conclusion

• Full coverage improves on corpus-path KD in both evaluated continuation settings, and Stage2 Top10 adds 48.85% to KD's Macro CE gain over NTP.Gains remain without coincident-state branch losses; additional depth or breadth further improves aggregate results at fixed per-term weight, with returns varying by corpus.Source: §7 · p. 8

Limitations

• Each evaluation condition uses one training run, so independent-seed replication is needed to assess the small absolute improvements.The results establish observed likelihood gains in these settings, not generality across model families, scales, languages, or downstream tasks.Source: Limitations · p. 9
• The main KD and PA1 objectives allocate nominal distillation weights differently, and comparisons match training tokens rather than compute.The experiments cover PA1 and extensions to two sequential steps or two branches per root; how gains scale with longer pretraining and additional feedback remains open.Source: Limitations · p. 9

• The experiments use publicly available language models and existing pretraining and evaluation corpora, with no human participants or newly collected personal data.Source: Ethical Considerations · p. 9
• Distillation may inherit biases or harmful content from the teacher and training data.Source: Ethical Considerations · p. 9
• Likelihood evaluations do not assess those effects, and the reported experiments identified no method-specific ethical risks.Source: Ethical Considerations · p. 9

A Execution and Cost Analysis

• Chunked full-vocabulary KL retains full-vocabulary normalization and temperature scaling, with FP32 workspace O(Pc) and O(P) saved scalars for P positions and chunk width c.Source: A Execution and Cost Analysis · p. 12
• Each PA1-OPD update counts valid prediction positions and eligible roots over the optimizer batch, then accumulates NTP, RealKD, and PA1 gradients across microbatches and workers.Source: A Execution and Cost Analysis · p. 12
• Student branches sample one token per selected root; unmatched branches reuse real-prefix K/V, while tree masking isolates branches and corpus logits supply predictions for coincident states.Source: A Execution and Cost Analysis · p. 12
• For sequence length L and U unmatched branches, each model evaluates L + U states and branches add O(UL) attention pairs without cross-root dependencies; added computation is training-only.Source: A Execution and Cost Analysis · pp. 12–13

B Experimental Details

• Experiments use a frozen OLMo-2-1124-7B teacher and OLMo-2-0425-1B students, evaluating C4 validation, PG-19 test, and Pile validation.Source: B Experimental Details · p. 13
• Each run consumes exactly 10,001,317,888 training tokens, with checkpoints near 2.5B, 5B, 7.5B, and 10B tokens and evaluation on complete 2048-token windows.Source: B Experimental Details · p. 13

• Student Top10, Top25, and Top50 select the indicated share of eligible roots by student surprisal, while Teacher Top50 ranks candidates by the teacher–student probability difference; objective weights are unchanged.Source: §B.1 · p. 13
• The coincident-state ablation keeps Student Top50 selection but removes branch KL for successors matching available corpus states, leaving NTP, RealKD, and the original branch denominator unchanged.Source: §B.1 · p. 13
• Positional diagnostics estimate selected-root frequency across 64 position bins on 8,192 fixed training sequences; frequencies are checkpoint-specific, not aggregated across training.Source: §B.1 · p. 13
• Branch-expansion comparisons use 1B-token continuations, coincident-state reuse, and same-checkpoint NTP and KD references; Stage2’s displayed tie may not reflect identical unrounded cross-entropy.Source: §B.1 · p. 13