Technical Summary: Geometric Self-Distillation for Reasoning Generalization

1. Problem Statement

On-policy distillation (OPD) has emerged as a practical post-training strategy for large language models (LLMs), providing dense token-level supervision on trajectories generated by the student model itself. A specific and effective variant, privileged-context self-distillation (OPSD), conditions the "teacher" (the same model) on additional information, such as a full solution trace or a hint, while the student must reason from the original problem context alone.

While this setup provides abundant supervision, the paper identifies a critical failure mode: predictive drift. Because the teacher possesses privileged context, it can confidently commit to continuations that are obvious given the solution but remain uncertain or unsupported from the student's own perspective. Standard distillation objectives (e.g., Kullback–Leibler divergence) weight teacher preferences by the teacher's own probability mass. Consequently, at states where the teacher is highly confident but the student assigns low probability (low overlap), the gradient signal is strong. Over many updates, these strong pulls force the student to adopt the teacher's confident continuations before the student's own information can justify them. This accumulation of mismatched updates leads to drift, where the model suppresses uncertainty and degrades out-of-distribution (OOD) reasoning performance, often resulting in "confident agreement on wrong answers."

Existing approaches attempt to mitigate this by filtering unreliable teacher signals or reshaping divergence weights (e.g., using Reverse KL or filtering mechanisms). However, the authors argue these methods fail to address the core issue: how far a trusted signal should move the student. Selection alone cannot regulate the magnitude of the update, and unchecked movement accumulates into drift.

2. Methodology: GEOSD

The authors introduce GEOSD (Geometric Self-Distillation), an objective designed to counteract this drift by treating predictive behavior as movement within a specific geometric space of next-token distributions. The method operates in two complementary ways:

A. Local Pressure Attenuation (Hellinger Loss)

Instead of standard KL divergence, GEOSD employs a Hellinger loss to modulate the teacher's pull at each visited state.

• Mechanism: The Hellinger distance is derived from the geometry of the probability simplex under a square-root embedding ($p \to \sqrt{p}$), where distributions lie on a hypersphere. The gradient of the Hellinger loss scales each teacher preference by the geometric mean of the student and teacher probabilities ($\sqrt{p_i q_i}$).
• Effect: If the teacher assigns high probability to a token but the student assigns near-zero probability, the overlap term vanishes, and the gradient pull is attenuated. Conversely, shared preferences transfer strongly. This ensures that supervision strengthens gradually as the student and teacher come into agreement, rather than forcing the student to immediately adopt the teacher's confident but unsupported continuations.

B. Global Drift Control (Fisher–Rao Proximal Term)

While overlap weighting bounds the pull at individual states, it does not prevent the accumulation of small, reasonable updates into a large global drift over training.

• Mechanism: GEOSD adds a proximal term that penalizes the distance between the current student distribution and a recent checkpoint ($\theta_{ckpt}$). This distance is measured using the Fisher–Rao distance (the arc length on the hypersphere), which tracks the accumulated displacement in the geometry of predictive distributions.
• Effect: This acts as a "moving trust region," damping abrupt displacements while allowing steady movement when the teacher signal is consistent. Unlike anchoring to the initial model (which would forbid adaptation), anchoring to a recent checkpoint allows the model to adapt while preventing it from drifting too far from its recent, stable behavior.

C. Natural-Gradient Update

Both the Hellinger loss and the Fisher–Rao proximal term are distances in the same geometry of next-token distributions. To optimize this objective, GEOSD uses a natural-gradient update.

• Mechanism: Instead of updating parameters in Euclidean space, the update is preconditioned by the Fisher Information Matrix (approximated via K-FAC for efficiency).
• Effect: This ensures that each step makes the smallest change to the student's predictions that improves the objective, rather than making a step of fixed size in parameter space. The proximal term effectively acts as a "checkpoint pullback" in this geometry, returning the policy toward the checkpoint based on the behavioral salience of the parameter directions.

3. Key Contributions

1. GEOSD Objective: A self-distillation framework that weights teacher preferences by teacher-student overlap (via Hellinger loss) and penalizes predictive drift from a recent checkpoint (via Fisher–Rao distance), optimized via natural gradients in a unified geometry.
2. Empirical Improvement: GEOSD preserves in-distribution (ID) gains of self-distillation while significantly improving OOD reasoning. Across three model families (Qwen3, Olmo, DeepSeek-R1) and five scales (1.7B to 32B), it improves average OOD accuracy by 5.7–8.6 points over the base model, whereas standard KL objectives degrade OOD performance.
3. Mechanistic Analysis: The paper characterizes why standard matching fails OOD: it concentrates probability mass on the teacher's top choice by "draining" alternatives at high-entropy states, leading to confident agreement on incorrect answers. GEOSD maintains the dominance of the top choice while keeping alternative continuations "in reach," preventing premature collapse.

4. Experimental Results

The authors evaluated GEOSD on mathematical reasoning benchmarks (DAPO-Math for training; AIME 2024/2025, AMC 2023, MATH-500 for OOD testing).

• Performance: GEOSD achieved the best OOD accuracy across all tested benchmarks and model families. For example, on Qwen3-8B, it improved OOD accuracy by +8.6 points over the base model, compared to a -8.1 point degradation for Forward KL.
• Robustness: The gains held across varying amounts of privileged context (from short prefixes to full solutions) and across model scales (1.7B to 32B).
• Ablation: Removing either the Hellinger overlap weighting, the Fisher–Rao proximal term, or the natural-gradient update resulted in significant drops in OOD performance, confirming that all three components are necessary.
• False Consensus: Analysis showed that standard methods (like Forward KL) exhibit high rates of "false consensus" (samples agreeing strongly on wrong answers) at high-entropy states. GEOSD reduced this rate below even the base model, maintaining informative self-consistency.

5. Significance and Claims

The paper claims that the degradation of OOD reasoning in privileged-context self-distillation is not merely a result of using unreliable signals, but a consequence of how those signals move the model. Standard objectives allow the teacher's confidence to force the student into a narrow mode prematurely, draining the probability mass of viable alternatives.

GEOSD demonstrates that effective reasoning distillation depends not only on which teacher signals to trust but on how far those signals are allowed to move the student. By regulating this movement geometrically—attenuating pulls where the student lacks support and anchoring the trajectory to a recent stable state—GEOSD achieves a superior trade-off between in-distribution performance and out-of-distribution generalization. The authors conclude that controlling predictive drift is essential for scaling self-distillation in reasoning tasks.

Limitations Noted: The authors acknowledge that their findings are currently established only for mathematical reasoning. The applicability of overlap-weighted transfer and Fisher–Rao drift control to other domains (e.g., code, scientific reasoning) where high-entropy states may have different structures remains untested. Additionally, the method assumes the base policy is broadly competent; for very weak base models, the mechanism might suppress useful signals that require large behavioral changes.