On the Geometry of On-Policy Distillation

2026-06-05

AI Summary

This work investigates the poorly understood dynamics of parameter updates in Offline Policy Distillation (OPD). Through trajectory analysis in parameter space, comparison along principal component directions, subspace constraints, and ablation studies, the study systematically characterizes OPD’s weight evolution and contrasts it with Supervised Fine-Tuning (SFT) and Reinforcement Learning (RL). The analysis reveals a distinctive “subspace locking” phenomenon in OPD: updates rapidly converge to a low-dimensional subspace that is critical for performance. Compared to SFT and RL, OPD exhibits sparser updates that notably avoid dominant principal directions. Moreover, the subspace established early in training is sufficient to sustain final performance, highlighting a unique optimization geometry that differentiates OPD from existing learning paradigms.

Abstract

On-policy distillation (OPD) is increasingly used to improve large language model reasoning, but its training dynamics remain poorly understood. We characterize the trajectory of OPD updates in parameter space and compare it with supervised fine-tuning (SFT) and reinforcement learning with verifiable rewards (RLVR). A suite of parameter-space diagnostics consistently places OPD in a relaxed off-principal regime: compared with SFT, its updates affect fewer weights and avoid principal directions more strongly, while compared with RLVR, they remain less tightly constrained. Beyond this static localization, OPD exhibits subspace locking: its cumulative updates rapidly enter a narrow low-dimensional channel. Constraining training to the update subspace formed early in training preserves OPD performance but substantially degrades SFT, indicating that the locked subspace is functionally sufficient for OPD. Control experiments further show that sparsifying the update tokens and shifting rollout generation off-policy preserve the rank dynamics, whereas mixing the OPD objective with RLVR changes them. Overall, these results suggest that OPD is not merely an intermediate point between SFT and RLVR, but induces its own update geometry in parameter space.

Problem

• On-Policy Distillation
• training dynamics
• parameter space
• update geometry
• large language models

Innovation

• on-policy distillation
• parameter space geometry
• subspace locking
• training dynamics
• low-rank updates