AIAny
Icon for item

On the Geometry of On-Policy Distillation

Analyzes the parameter-space geometry of on-policy distillation (OPD) for LLM training, showing OPD updates affect fewer weights, avoid principal directions, and rapidly lock into a low-dimensional update subspace. Compares OPD with supervised fine-tuning (SFT) and reinforcement learning with verifiable rewards (RLVR) and studies implications for optimization and objective mixing.

Introduction

Why this matters

Most discussions treat on-policy distillation (OPD) as somewhere between supervised fine-tuning (SFT) and RL-style objectives. This paper shows a different picture: OPD induces a distinct update geometry in parameter space — its updates are spatially sparse (affect fewer weights), avoid the main principal directions that SFT follows, and quickly collapse into a narrow, low-dimensional "locked" subspace whose cumulative updates are functionally sufficient for OPD performance.

Key Findings
  • OPD update localization: Compared to SFT, OPD updates change fewer parameters and avoid dominant principal components; compared to RLVR, OPD updates are less tightly constrained. So what: OPD isn't just a midpoint between SFT and RLVR — it follows different directions that matter for optimization and generalization.
  • Subspace locking: Cumulative OPD updates rapidly enter a low-dimensional channel. Constraining training to the early-formed update subspace preserves OPD performance but substantially degrades SFT. So what: early directions in OPD capture most functional change; this opens possibilities for cheap constrained training or diagnostics.
  • Robustness of rank dynamics: Control experiments show sparsifying update tokens and shifting rollout generation off-policy preserve the rank dynamics, while mixing the OPD objective with RLVR alters them. So what: the rank/subspace behavior is tied to the OPD objective itself and not merely to token-level sparsity or rollout specifics.
  • Practical implication: Because OPD's effective updates concentrate in a small subspace, practitioners can consider monitoring or constraining update subspaces for efficiency, but mixing objectives (e.g., adding RLVR) can undo these geometric properties.
Who it's for + trade-offs

Great fit if you research or engineer LLM fine-tuning and care about the microscopic dynamics of training — model researchers, optimization specialists, and teams experimenting with distillation or hybrid objectives. The paper gives a concrete geometric lens to reason about why OPD behaves differently and offers diagnostics you can run on your runs.

Look elsewhere if your primary concern is immediate engineering recipes (this paper focuses on geometry and diagnostics, not turnkey training pipelines) or if you need empirical guarantees across a wide sweep of scales and tasks — the findings focus on the studied setups and their mechanistic insights rather than exhaustive benchmarks.

Where it fits

This work positions OPD as a method that induces its own parameter-space signature distinct from SFT and RLVR. It complements empirical performance papers by providing interpretability tools and hypotheses about why OPD-trained models behave differently.

Short methodological note

The paper uses parameter-space diagnostics (principal-component analysis of updates, rank measurements, and constrained-training experiments) plus targeted controls (token sparsification, off-policy rollouts, objective mixing) to isolate which components of OPD drive the observed low-rank locking. The experiments suggest the locked subspace arises early in training and is functionally sufficient for OPD's behavior.

Information

  • Websitearxiv.org
  • AuthorsZhennan Shen, Yanshu Li, Qingyu Yin, Chak Tou Leong, Zhilin Wang, Yanxu Chen, Rongduo Han, Sunbowen Lee, Yi R. Fung
  • Published date2026/06/05

More Items

Enables RL post-training with million-token prompts under a fixed GPU budget by evaluating shared prompt state without autograd, retaining only minimal model state, and replaying short response branches; instantiated as GRPO and demonstrated on Qwen3.6-27B and GLM-5.2 up to multi-million token execution.

Analyzes adversarial weaknesses of World-Action Models (WAMs) via BadWAM, a framework that crafts visual perturbations to decouple a model’s imagined future from its executed actions. Introduces two attack modes—action-only (disruptive) and imagination-preserving (stealthy)—and shows large drops in closed-loop task success (e.g., 96.5%→43.1%).

Converts completed on-policy trajectories into natural-language 'hindsight skills' and converts the skill-induced action probability shifts into a dense token-level on-policy distillation signal, jointly optimized with outcome-based RL to improve sample efficiency and long-horizon agent behavior.