On the Generalization Limitation of On-Policy Distillation for Post-RL Consolidation

AI summary

Overview

Research area: Machine learning — model compression and capability transfer, specifically knowledge distillation and post-reinforcement-learning (post-RL) model consolidation.

Technical level: Advanced. The paper assumes familiarity with on-policy distillation, reinforcement learning fine-tuning, in-distribution versus out-of-distribution (OOD) evaluation, and optimization/update-geometry analysis.

Scope: An empirical investigation showing that on-policy distillation — single-teacher and multi-teacher — can reproduce teachers' in-distribution scores while degrading out-of-distribution generalization below even the base model, plus two proposed mitigations.

What This Paper Is About

Distillation is commonly used to fold several specialized "teacher" models back into one student model after reinforcement learning. Prior work reports that multi-teacher on-policy distillation (MOPD) matches teacher performance almost perfectly, and this is treated as evidence that the capabilities were successfully consolidated. The authors argue that this near-perfect in-distribution score hides a gap: the student does not actually inherit the teachers' ability to generalize beyond the data it was supervised on. The goal is to measure that gap on real tasks, identify why it occurs, and test whether it can be closed.

Key Contributions

1. Identifies a blind spot in distillation evaluation. The paper shows that near-perfect in-distribution (IND) distillation scores can mask incomplete consolidation of capabilities outside the supervised distribution, using both a synthetic reasoning study and real-world benchmarks.

2. Quantifies OOD degradation across capability domains. Using Qwen3 models, the authors evaluate extrapolative and compositional generalization on mathematical reasoning, coding, agent, and long-context capabilities, measuring student performance against both the teachers and the base model.

3. Traces the gap to two bottlenecks and tests remedies. Teacher capability and data coverage are examined as candidate causes; the abstract reports that jointly alleviating both eliminates OOD forgetting in their experiments.

4. Proposes two geometry-motivated interventions. Analysis of update geometry motivates "geometry-based patching" and "base-preserving replay," intended to improve OOD performance while largely retaining IND performance.

Main Findings

• Synthetic study shows incomplete transfer: A preliminary study on synthetic reasoning finds that OPD does not fully transfer teachers' generalization capabilities, motivating the larger investigation.

• High IND scores, negative OOD gap: Across extrapolative settings, OPD and MOPD reach 96.04–102.73% of their teachers' IND scores, yet their OOD scores remain 2.45–16.25% below the base model's. The student can therefore look like a successful copy of the teacher while being worse than the model it started from on unseen inputs.

• Compositional settings show the same pattern: Across three compositional settings, MOPD improves IND performance over the base model, but on composed tasks treated as OOD it falls 18.19–24.34% below the best teachers.

• Two bottlenecks, jointly sufficient in their experiments: Examining teacher capability and data coverage, the authors report that alleviating both together eliminates OOD forgetting in their experiments, and suggest this hints at generalization beyond the teacher rather than mere imitation.

• Distillation still beats RL baselines on forgetting: (M)OPD reduces OOD forgetting on average relative to its reinforcement learning baselines, so the problem is comparative rather than absolute.

• Update geometry differs from RL: The student's updates under (M)OPD are smaller and more spectrally concentrated on average than those of the RL baselines, and MOPD shows greater shared update structure across teachers.

• Interventions partially recover OOD performance: Geometry-based patching and base-preserving replay improve OOD performance while largely preserving IND performance, per the abstract; specific magnitudes are not given in it.

Methodology in Plain English

The authors start small: a controlled synthetic reasoning task where they can check whether a distilled student really learns the teacher's reasoning behavior or only its surface patterns. They then scale the question to real models — Qwen3 — and real capability areas, covering math, code, agents, and long-context handling. In each area they compare three things: the student's score on data resembling its training distribution (IND), its score on deliberately shifted or recombined tasks (OOD), and the corresponding scores of the teachers and the untouched base model. The base-model comparison is the key design choice, since it reveals whether distillation left the student worse off than doing nothing. To explain the gap, they probe two suspected causes — whether the teacher itself has the capability to transfer, and whether the training data covers the situations where generalization is needed — and test what happens when both are addressed. Finally, they inspect the numerical structure of the weight updates made during distillation versus reinforcement learning, finding that distillation's updates are smaller, more concentrated in a few directions, and more similar across teachers, then design two fixes informed by that picture.

Why This Matters

Impact on research: Distillation is routinely validated with in-distribution benchmarks, and this paper argues that such validation is insufficient for consolidation claims. It reframes "the student matches the teacher" as a weaker result than it appears, and points to a need for algorithms designed explicitly for OOD retention rather than imitation fidelity. It also connects distillation to the broader question of whether post-training pipelines erode capabilities that the base model already had.

Real-world applications:

• Model consolidation after specialization: Teams that fine-tune several specialist models with RL and merge them into one deployable model need to know whether merged behavior holds up on inputs outside the training mix.
• Cost reduction in production: Replacing large or multiple teacher models with a smaller distilled student is only safe if the student does not silently lose robustness on rare or shifted inputs.
• Agentic and coding systems: These are precisely the domains where deployed inputs drift from training data, so OOD degradation translates into real task failures.
• Long-context deployments: Document lengths and structures at inference time rarely match the distillation corpus, making the reported long-context evaluation directly relevant.

Industry relevance: Distillation is a standard cost lever, and the paper's central claim — that students can score at or above teacher level in-distribution while landing measurably below the base model out-of-distribution — is a risk signal for anyone shipping distilled models on the strength of benchmark parity.

Future Directions

• Algorithms built for OOD retention: The authors conclude that capability consolidation needs methods that explicitly support OOD retention and generalization, rather than relying on IND match as a proxy.
• Scaling the bottleneck remedies: Since jointly relieving teacher-capability and data-coverage limits eliminated OOD forgetting in their experiments, the natural question is whether that result holds across model sizes, teacher counts, and capability areas.
• Refining the geometry-based interventions: Geometry-based patching and base-preserving replay are reported to help, leaving open how they behave under other distillation objectives and whether the geometric signature generalizes.
• Better evaluation practice: The work implies a need for consolidation benchmarks that report base-model-relative OOD scores alongside IND scores, so that forgetting is visible by default.

Target Audience

Researchers and engineers working on knowledge distillation, multi-teacher merging, and post-RL training pipelines. It is most useful to practitioners who use IND benchmark parity as evidence that a distilled model is safe to deploy, and to anyone studying generalization or capability forgetting in LLM post-training. Readers without background in distillation and OOD evaluation will need to consult the underlying OPD/MOPD literature first.

Authors’ abstract

Multi-teacher on-policy distillation (MOPD) extends on-policy distillation (OPD) to consolidate capabilities from multiple teachers. Although prior work reports near-perfect in-distribution (IND) distillation performance, we find that this can mask incomplete capability consolidation beyond the supervised distribution. A preliminary study on synthetic reasoning shows that OPD does not fully transfer teachers' generalization capabilities. We then extend this investigation to real-world tasks, evaluating extrapolative and compositional generalization across mathematical reasoning, coding, agent, and long-context capabilities using Qwen3 models. Across extrapolative settings, OPD and MOPD reach 96.04-102.73% of their teachers' IND scores. Yet their out-of-distribution (OOD) scores remain 2.45-16.25% below the base model's. Across three compositional settings, MOPD improves IND performance over the base model, but its performance on composed tasks as OOD remains 18.19-24.34% below that of the best teachers. To understand these gaps, we examine teacher capability and data coverage as two potential bottlenecks. Jointly alleviating these bottlenecks eliminates OOD forgetting in our experiments and suggests the potential for generalization beyond the teacher. (M)OPD also reduces OOD forgetting on average relative to its reinforcement learning (RL) baselines. Analysis of the associated update geometry reveals smaller and more spectrally concentrated updates on average than those of the RL baselines, along with greater shared update structure under MOPD. Motivated by these findings, we propose geometry-based patching and base-preserving replay to improve OOD performance while largely preserving IND performance. Our findings identify a blind spot in capability consolidation via (M)OPD, highlighting the need for algorithms that explicitly support OOD retention and generalization.