Technical Summary: Privileged Hidden Flow for On-Policy Self-Distillation

Problem Statement

On-policy self-distillation (OPSD) is a technique for training reasoning models where a student model samples its own rollouts and is supervised by a "privileged" teacher. This teacher is the same model architecture but conditioned on additional information, such as a verified reference solution, while the student sees only the problem prompt.

Existing OPSD objectives focus exclusively on matching the output distribution (the next-token probabilities) between the student and the privileged teacher. While this aligns the final predictions, it fails to supervise the internal computation that produces those predictions. Specifically, privileged context influences the intermediate hidden states and the trajectory of these states through the network layers. Current methods only constrain the final output head, leaving the "hidden flow"—how the model's internal representations evolve step-by-step along the rollout—unconstrained. The paper posits that distilling this internal process could improve reasoning performance without altering the fundamental on-policy training setup.

Methodology: Privileged Hidden Flow (PHF)

The authors propose Privileged Hidden Flow (PHF), a method that augments standard OPSD by distilling the movement of hidden states rather than the states themselves.

Core Mechanism

Instead of forcing the student's hidden vector $h^S_{\ell,t}$ to match the teacher's vector $h^T_{\ell,t}$ at the same token position (pointwise matching), PHF matches the transitions between token positions.

1. Transition Definition: For a set of selected generated positions $W = \{w_1, \dots, w_m\}$, the local hidden transition is defined as the difference between neighboring positions: $\Delta h_{\ell,i} = h_{\ell, w_{i+1}} - h_{\ell, w_i}$.
2. Normalization: These transitions are normalized to unit directions to focus on the geometry of movement rather than magnitude.
3. Loss Components:
   • Direction Loss ($L_{dir}$): Aligns the direction of student transitions with teacher transitions at each layer.
   • Geometry Loss ($L_{geo}$): Matches the Gram matrix of the transition directions. This captures the relational structure of the trajectory (e.g., which steps are parallel, orthogonal, or opposing) rather than absolute coordinates.
   • Neighboring-Layer Relation ($L_{adj}$): Computes the relation between transition directions in adjacent layers ($\ell$ and $\ell+1$) to ensure the transformation of motion across depth is consistent with the teacher.

Structural Invariances

A key theoretical contribution of PHF is its invariance properties, which address the "representational drift" problem where hidden coordinates shift during training:

• Offset Invariance: The loss is unchanged if a constant offset is added to the entire trajectory of either student or teacher.
• Scale Invariance: The loss is invariant to positive scalar multiplication of transitions (before normalization).
• Orthogonal Invariance: The geometry term is invariant to independent orthogonal transformations of the student and teacher transition sets.

These properties allow PHF to supervise the shape and direction of the reasoning process without requiring the student and teacher to occupy the exact same points in the high-dimensional hidden space.

Training Recipe

• Teacher Source: The privileged teacher is an Exponential Moving Average (EMA) of the online student, conditioned on the reference solution.
• Aggregation: The method aggregates the transition loss across all layers rather than selecting a subset, and combines the local transition loss with the neighboring-layer relation loss.
• Objective: The final loss is $L_{PHF} = L_{OPSD} + \alpha L_{Flow}$, where $L_{Flow}$ combines the local and adjacent losses. The coefficient $\alpha$ is fixed (0.05) across scales.
• Constraints: The method uses a fixed window of 128 generated positions to manage memory and does not introduce reward models, correctness routing, or additional rollouts.

Key Contributions

1. Privileged Transition-Transport Objective: Introduces a novel loss that matches normalized token-to-token hidden transitions and their trajectory geometry, moving beyond pointwise state matching.
2. Theoretical Characterization: Proves that the transport loss possesses exact algebraic invariances to trajectory offsets and orthogonal transformations, distinguishing it from standard hidden-state matching.
3. All-Layer Aggregation: Demonstrates that averaging the transition loss over all layers, combined with a fixed neighboring-layer relation, yields consistent improvements without requiring manual layer selection.
4. Empirical Validation: Validates the fixed recipe across three model scales (Qwen3-1.7B, 4B, and 8B) under the official OPSD mathematical reasoning protocol.

Experimental Results

The authors evaluated PHF on the AIME 2024, AIME 2025, and HMMT 2025 benchmarks using the Average@12 metric (mean correctness over 12 samples per problem).

• Performance Gains: Under a fixed 100-step training schedule, PHF improved the Average@12 aggregate over the reproduced OPSD baseline:
  • Qwen3-1.7B: +2.2 points (40.83 $\to$ 43.06)
  • Qwen3-4B: +1.5 points (61.89 $\to$ 63.39)
  • Qwen3-8B: +1.7 points (64.26 $\to$ 65.94)
• Consistency: The gains were positive across all scales and nearly all individual benchmark cells.
• Ablation Studies:
  • Direction vs. Geometry: Neither direction-only nor geometry-only variants matched the full PHF recipe across all scales, indicating both components are necessary.
  • Layer Selection: A "selected layers" control (matching a small hand-picked subset) performed worse than the all-layer aggregation, supporting the decision to average across depth.
  • Pointwise Matching: The paper notes that pointwise hidden-state matching is a competitive boundary control but is less robust to representational drift than the transition-based approach.

Significance and Claims

The paper claims that PHF offers a modest but meaningful improvement to the OPSD framework by utilizing privileged information more structurally.

• Process-Level Supervision: The primary significance is demonstrating that privileged references induce a specific "hidden process" (how the model moves through its own rollout) that can be distilled. PHF captures this process without requiring the student to match the teacher's absolute hidden coordinates.
• Deployment Compatibility: Unlike methods requiring reward models or external verifiers at test time, PHF maintains the OPSD deployment setting: the evaluated model is the online student, and no privileged information is needed during inference.
• Efficiency: The method adds only a scalar loss coefficient and a hidden-state storage window, avoiding the computational overhead of additional rollouts or complex routing mechanisms.
• Scope of Claims: The authors are explicit that their evidence is limited to a fixed recipe, a specific model family (Qwen3), and mathematical reasoning tasks. They do not claim that hidden-flow matching improves every aspect of reasoning (e.g., proof quality or calibration) or that it is a universal solution, but rather that it provides a consistent, positive aggregate shift when applied to on-policy self-distillation.