On-Policy Distillation Teaches New Skills but Not New Knowledge

AI summary

Overview

• Research area: Natural language processing; specifically knowledge distillation for large language models, with a focus on on-policy distillation (OPD) and multi-step reasoning.
• Technical level: Advanced. The paper assumes familiarity with KL divergence objectives, teacher-student distillation, LoRA updates, pass@k evaluation, and Transformer layer roles.
• Scope: A controlled synthetic study plus real-world validation that separates whether on-policy distillation transfers a teacher's factual knowledge, its compositional reasoning skill, or both.

What This Paper Is About

On-policy distillation supervises a student model along its own generated trajectories and is widely used to strengthen reasoning in post-training pipelines. But higher benchmark scores alone cannot tell us whether the student learned new facts, learned to combine existing facts better, or both. This paper builds a controlled setting that isolates those two capabilities and measures what actually transfers.

Key Contributions

1. A controlled synthetic framework that separates facts from skill. Facts are random lookup tables mapping digit pairs to digits with no algebraic rule, so answers must be memorized rather than inferred. Compositional skill is operationalized as solving sequential chains versus branching trees that use the same facts and the same number of operations.
2. A 2x2 factorial teacher design. Starting from an initial student S₀, the authors train teachers that add held-out facts (T_f), branching-tree skill (T_s), or both (T_fs), with the frozen S₀ as a no-addition baseline teacher (T_∅).
3. A decoupling of the distillation recipe. The paper separates the source of training prefixes (student rollouts vs. teacher demonstrations) from the divergence direction (reverse KL vs. forward KL) to identify which factor drives each type of transfer.
4. A mechanistic analysis of weight updates. Merged LoRA updates are compared across three models to show where factual and compositional learning live in parameter space and layer depth.

Main Findings

• Compositional skill transfers reliably. Averaged over four models, students distilled from T_s reach 80.96% greedy accuracy on teacher-only tree structures and 79.10% on unseen tree structures, versus 0.72% on the same unseen structures under the matched T_∅ control. Both tests use only facts S₀ already knew.
• Factual knowledge barely transfers under reverse KL. Students distilled from T_f reach a mean greedy accuracy of 13.12% on single-step queries about held-out facts (versus 0.03% for the matched T_∅ control), but only 0.54% on multi-step chains where every step queries a held-out fact. Gemma-2-2B reaches 32.95% on single-step queries, showing architecture-dependent variation.
• The asymmetry survives larger sampling budgets. Pass@32 is 88.11% on unseen trees for T_s students, but only 15.74% on single-step held-out fact queries and 2.58% on chains for T_f students.
• The KL direction controls factual transfer. Replacing reverse KL with forward KL raises single-step greedy accuracy from 10.95% to 54.10% in Qwen3-4B and from 32.95% to 94.10% in Gemma-2-2B under student prefixes. Under teacher prefixes, forward KL reaches 54.05% (Qwen3-4B) and 93.00% (Gemma-2-2B).
• Student rollouts improve compositional execution. Under both KL directions, training on student prefixes lowers reference-completion NLL. For Qwen3-4B, NLL falls from 2.40 to 1.77 bits under reverse KL and from 2.13 to 1.15 bits under forward KL. For Gemma-2-2B, NLL falls from 28.13 to 25.65 bits and from 17.40 to 13.65 bits.
• Real-world benchmarks show the same pattern. On Qwen3-4B-Base, the knowledge teacher T_K reaches 84.55% on Knowledge-2026, but distillation from T_K moves the student from 11.25% to 9.64%. Distillation from the reasoning teacher T_R raises Math-2026 from 17.41% to 23.88%, a 6.47 percentage point gain.
• No cross-capability transfer. Distilling from T_R leaves factual recall at 10.75% (near S₀'s 11.25%), while distilling from T_K shifts math accuracy only to 19.40%.
• Failures look different in the traces. Under reverse KL from T_f, the student keeps the correct step-by-step format but outputs wrong lookup values. Under forward KL from T_f, it recalls all facts correctly but flattens a tree into a left-to-right chain.
• Facts and skills occupy near-orthogonal parameter directions. Cosine similarity between T_f and T_s updates ranges from 0.004 to 0.018 (Qwen3-4B), 0.011 to 0.036 (Gemma-2-2B), and 0.004 to 0.046 (OLMo3-7B). Changing only the divergence objective under T_s preserves alignment of 0.73–0.78, 0.64–0.82, and 0.68–0.80 respectively.
• The two capabilities peak at different depths. In Qwen3-4B, T_f updates increase monotonically with depth and peak at layer 35, whereas T_s updates peak at layers 3 and 11.
• Hybrid objectives can improve both. On Qwen3-4B, ToDi reaches 64.3% factual recall after distillation from T_f and 86.4% on unseen trees after distillation from T_s, exceeding forward KL by 10.2 and 10.0 percentage points. GKD (β=0.5) and EOPD retain the asymmetry, reaching 70.6% and 78.7% on unseen trees but only 9.1% and 13.7% factual recall.

Methodology in Plain English

The authors need to know exactly what a model already knows before distillation, which is impossible to audit with pretrained models trained on vast corpora. So they build a synthetic world with two ingredients they can control independently.

Facts are random lookup tables. Each table maps a two-digit input to a single output digit across 100 entries, with entries sampled independently and uniformly at random. Because there is no algebraic rule, an unseen entry cannot be derived from the others; the model must have memorized it. The authors create 25 tables (T_1 through T_25), splitting them into 20 initial facts and five held-out facts, and this split stays fixed throughout.

Skill is the structure in which facts get used. A chain feeds each result directly into the next step. A branching tree solves separate branches first, holds intermediate results across intervening steps, and combines them later. Both structures query the same facts and require the same number of operations, so only the organization differs. Each query in the main experiments contains seven operations over eight inputs, and completions include explicit step-by-step traces before the final answer.

They train an initial student S₀ on initial facts using single-step queries and chains. S₀ reliably executes chains over initial facts but performs near zero on held-out facts and branching trees. From S₀ they build teachers: T_f adds held-out facts, T_s adds branching trees over initial facts, T_fs adds both, and T_∅ is the unchanged S₀.

Distillation then uses token-level reverse KL: the current student samples a rollout, the frozen teacher supplies its next-token distribution over the same prefix, and the student is trained to match it. To find the source of the asymmetry, the authors cross two axes: prefix source (student rollouts vs. teacher demonstrations) and divergence direction (reverse KL vs. forward KL), creating a 2x2 design evaluated on Qwen3-4B and Gemma-2-2B. Evaluation separates factual recall on held-out facts from compositional skill over initial facts, and tree structures are divided into shared, teacher-only, and unseen groups.

Real-world validation uses Qwen3-4B-Base fine-tuned on Tulu-3-SFT as S₀. Knowledge-2026 aggregates 809 question-answer pairs from RealTimeQA, SealQA, FreshQA, and newly edited Wikipedia articles via the WikiBigEdit pipeline, chosen because they postdate the base model's training cutoff. Math-2026 contains 201 competition problems from MathArena and NuminaMath-CoT, held out from a pool of 1,000; the remaining 799 serve as distillation prompts.

Why This Matters

Impact on research. The paper reframes an active debate. Some prior work reports capability expansion from on-policy distillation (for example, the Qwen3 technical report shows improved Pass@64 on math reasoning), while other work argues gains come from more efficient sampling of existing capabilities, and Ding and Zhang show that suppressing low-probability student tokens can match OPD's reasoning gains without teacher supervision. This paper's controlled separation suggests both can be partly right: OPD reorganizes what a model already knows rather than expanding its parametric memory, and the KL direction is the lever that governs factual transfer.

Real-world applications:

• Post-training pipelines for reasoning-focused model releases, where the goal is to sharpen multi-step problem solving rather than add facts.
• Retrieval-augmented and tool-using systems, where the model needs to compose retrieved facts rather than store them internally.
• Continual or domain adaptation, where the choice between forward-KL and reverse-KL supervision determines whether new domain facts actually stick.
• Model auditing and evaluation, where separating fact recall from compositional skill gives a clearer picture of what training changed.

Industry relevance. Practitioners deciding between OPD and standard supervised distillation now have a concrete rule of thumb: use reverse-KL OPD to strengthen compositional reasoning, and use forward-KL supervision when updating factual knowledge is the objective. The mechanistic findings, that factual updates concentrate in late layers while compositional updates reshape early routing layers, also inform layer-wise fine-tuning and adapter placement decisions.

Future Directions

• Testing whether hybrid objectives such as ToDi, which reaches 64.3% factual recall and 86.4% on unseen trees on Qwen3-4B, scale to natural-language knowledge acquisition and to models beyond Qwen3-4B.
• Extending the orthogonal-subspace and depth-specialization findings to more architectures, since the reported cosine similarities and layer peaks come from three models and the layer-wise breakdown is given only for Qwen3-4B.
• Determining whether the small factual gains under reverse KL in some architectures, such as Gemma-2-2B's 32.95% single-step accuracy, reflect architecture-specific behavior worth exploiting.
• Investigating whether compositional skill transfer generalizes to reasoning structures beyond chains and branching trees, given that unseen structures were solved at 79.10% greedy accuracy.

Target Audience

Researchers and engineers working on knowledge distillation, LLM post-training, and reasoning capability development will get the most from this paper. It is also valuable for evaluation specialists who need methods for separating factual recall from compositional skill, and for practitioners choosing distillation objectives when building domain-adapted or reasoning-focused models.

Authors’ abstract

On-policy distillation (OPD) strengthens language-model reasoning, yet whether students acquire new factual knowledge or compositional skill for multi-step reasoning remains unknown. We separate these capabilities using a controlled synthetic framework that measures the student's initial capabilities and independently controls the teacher's additional facts, compositional skill, or both. Across four models from three families, reverse-KL OPD reliably transfers compositional skill across unseen reasoning structures, but transfers minimal factual knowledge. Decoupling the distillation recipe reveals the source of this asymmetry: replacing reverse KL with forward KL restores factual transfer, whereas student rollouts specifically improve the execution of multi-step reasoning. Experiments on recent factual QA and competition mathematics show a similar asymmetry under reverse-KL OPD, yielding notable reasoning gains without factual memory expansion. Together, these results demonstrate that on-policy distillation does not expand a model's parametric knowledge, but instead teaches it to organize and compose the knowledge it already possesses.