Geometric Self-Distillation for Reasoning Generalization

Abstract: On-policy distillation is a practical post-training recipe for LLMs, supplying dense teacher supervision on the student's own trajectories. In privileged-context self-distillation, teacher and student are the same model conditioned on the same prefix, but the teacher also sees a hint or the full solution trace. This makes supervision abundant but harder to trust: the teacher can be confident about continuations its privileged view makes obvious but the student cannot yet justify. The distillation pull is strongest where teacher and student disagree most, and over many updates it accumulates into drift that degrades out-of-distribution (OOD) reasoning. We introduce GeoSD, a geometric self-distillation objective that treats this drift as movement in the student's predictive behavior and counters it in two complementary ways. A Hellinger loss scales each teacher preference by the overlap the student already shares with it, attenuating the pull on tokens the student cannot yet support. Since these pulls still compound over training, a proximal term penalizes how far the student's predictions drift from a recent checkpoint, measured as a Fisher-Rao distance. Both are distances in the same geometry of next-token distributions, and a natural-gradient update takes its steps in that geometry rather than in parameter space. Across mathematical reasoning benchmarks and three model families, GeoSD preserves the in-distribution gains of self-distillation while improving average OOD accuracy by 5.7-8.6 points over the base model, with gains holding across model scales from 1.7B to 32B. Analyzing why standard matching fails out of distribution, we find it wins agreement with the teacher by draining mass from alternatives at high-entropy states, resulting in confident agreement on wrong answers, whereas GeoSD keeps those alternatives in reach.

• The paper introduces GeoSD, which employs Hellinger-based attenuation to modulate teacher signals and mitigate unsupported pulls in LLM distillation.
• It uses a Fisher–Rao proximal penalty to control global predictive drift, ensuring consistent model behavior and preserving in-distribution gains.
• Empirical results show GeoSD improves out-of-distribution accuracy by 5.7–8.6 points across various model sizes while maintaining robust ID performance.

Geometric Self-Distillation for Robust Reasoning Generalization

Abstract and Motivation

Recent advances in on-policy distillation (OPD) have highlighted its practicality for post-training LLMs, particularly in reasoning-dense domains where trajectory-level supervision is critical. A growing trend within OPD is privileged-context self-distillation: using the same model as both teacher and student but allowing the teacher access to additional privileged context (e.g., the solution trace), whereas the student must reason from raw prompts. While this enables abundant supervision, it fundamentally alters the trustworthiness of teacher signals—privileged teachers often supply confident continuations that the student, operating under information constraints, cannot substantiate. Unchecked, standard divergences amplify this teacher-student mismatch at positions of low overlap, causing predictive distribution drift that degrades out-of-distribution (OOD) generalization.

Methodological Framework: GeoSD

This work introduces GeoSD (Geometric Self-Distillation), a distillation mechanism that treats learning as movement on the geometry of predictive distributions, using information geometry to both attenuate local teacher pressures when student support is lacking and regulate global drift via a proximal constraint. Two principal constructs underpin GeoSD:

1. Hellinger-based Local Attenuation: The standard KL divergence scales teacher signal with the teacher's own confidence, pressing the student harshly toward privileged continuations even when unsubstantiated. GeoSD replaces the KL with the Hellinger loss, which weights each teacher preference by the square-root geometric mean overlap ($\sqrt{p q}$)—effectively nullifying the pull for low-overlap tokens and smoothly growing the distillation signal as agreement increases.

   Figure 1: GeoSD regulates distillation in the geometry of predictive distributions. (A) The teacher's pull scales with teacher–student overlap, growing with agreement. (B) Distributions lie on a sphere, where the Hellinger chord bounds each pull and the Fisher–Rao arc tracks drift from a checkpoint.

2. Fisher–Rao Proximal Drift Control: Even with per-position control, small, locally reasonable updates can accumulate unbounded drift over training when teacher and student move in lockstep. To prevent this, GeoSD introduces a proximal penalty using Fisher–Rao (FR) distance between student predictive distributions at the current parameters and a recent checkpoint, enforcing behavioral consistency over time.

These two objectives are naturally unified by optimizing on the probability simplex geometry (spherical manifold under square-root embedding), where both local (Hellinger) and global (Fisher–Rao) distances have closed-form expressions, and updates are executed via natural gradient preconditioned by the empirical Fisher Information (K-FAC).

Empirical Results

GeoSD is extensively validated on mathematical reasoning tasks using multiple model families (Qwen3, Olmo, DeepSeek-R1; sizes from 1.7B to 32B). The experimental protocol employs on-policy rollouts and evaluates both in-distribution (ID) and OOD (AIME, AMC, MATH-500) reasoning robustness.

Figure 3: In-distribution vs. out-of-distribution gains (percentage points relative to $\pi_0$) for every method across the three model families; dot size encodes OOD pass@16.

Key findings include:

• Superior OOD Generalization: GeoSD consistently improves OOD accuracy by $5.7$–$8.6$ points (avg@16) over the base model across all families and sizes, whereas standard KL-based OPSD methods degrade OOD generalization despite strong ID gains.
• Retention of ID Gains: Unlike filtering or softened divergences (JSD, skewed KL), GeoSD preserves or nearly matches the ID improvement of aggressive KL-based distillation, demonstrating an improved ID–OOD trade-off frontier.
• Robustness Across Privilege and Scale: GeoSD's gains grow monotonically with increased teacher privilege and persist across substantial model scaling.

Figure 2: Privilege sweep with Qwen3-8B. OOD accuracy change (avg@16; vs. $\pi_0$) versus the fraction of the solution revealed to the teacher.

Mechanistic Insights

The paper provides a detailed analysis of why standard distillation fails OOD:

• Local Probability Concentration: Standard KL matching compels the student to locally concentrate mass on the teacher's top prediction at high-entropy (uncertain) points, hollowing out support for plausible alternatives.
• False Consensus: This over-concentration propagates downstream, resulting in high-confidence, self-consistent errors (false consensus) on OOD data.

  Figure 4: The shape of teacher–student alignment. Student next-token mass at high-entropy decision points (top 10% by $\pi_0$ entropy) for Qwen3-8B. At each position, tokens are sorted by the teacher's preference rank; each ridge shows the student's mean mass over these ranks.

  

  Figure 5: From local concentration to false consensus. (A) OOD pass@$k$ versus number of samples $k$. (B) False-consensus rate: how often samples agree strongly on a wrong answer. (C) Majority-answer accuracy versus consensus strength.

GeoSD, in contrast, maintains mass on alternative tokens throughout training, slowing the collapse and enabling recovery of correct solutions by sampling even under distribution shift. As measured in answer-level statistics, GeoSD reduces false consensus rates below even the base model and keeps self-consistency informative, as opposed to "collapsed" confidence in standard methods.

Ablation and Robustness Studies

Ablative experiments demonstrate the necessity of each GeoSD component:

• Hellinger vs. JSD: Substituting Hellinger with JSD, or dropping the Fisher–Rao proximal penalty or natural-gradient update, each substantially weakens OOD performance.
• Step Size and Early Stopping: Unlike forward KL, where aggressive learning rates yield an ID/OOD trade-off frontier with no OOD gain over the base, GeoSD forms a performance plateau robust to step size and training duration.

  Figure 6: Step size and learning rate. (Qwen3-8B, avg@16). Each point is one configuration on the ID–OOD plane: solid lines sweep the learning rate, dashed lines trace early stopping.

Practical Implications and Theoretical Contributions

GeoSD's geometry-aware approach yields several implications:

• Post-training for OOD Robustness: GeoSD is a practical, computationally efficient recipe for LLM post-training that achieves unmatched OOD robustness in a highly competitive setting, at a modest compute and memory overhead.
• Principled Use of Privilege: By modulating the magnitude (not just selection) of teacher signals, the method unlocks benefits distinct from those of teacher selection, filtering, or divergence softening.
• Generalizable Framework: The geometric formalism—local Hellinger loss, Fisher–Rao drift control, and natural-gradient update—provides a principled platform for future advances in LLM training with asymmetric information or privileged teachers. Its separation of geometry from parameterization is theoretically appealing and practically tractable via K-FAC approximations.

Limitations and Future Directions

• Domain and Context Scope: Evaluation is confined to mathematical reasoning, with the privileged context restricted to complete solution traces. Adaptation to different task types, partial privilege (hints, feedback), or teacher selections beyond current-student instantiations is unexamined.
• Model Weakness and Distributional Change: For weak base models or where useful signals reside in low-overlap regions, GeoSD's suppression of unsupported teacher pressure may hinder beneficial learning. Extensions to handle such settings require further investigation.
• Long-Horizon Drift: The proximal anchor is moving (recent checkpoint), so over extremely long horizons, unchecked drift may still occur if student and anchor co-evolve.

Conclusion

GeoSD demonstrates that express geometric regulation of distillation dynamics is critical to safe, generalizable post-training of LLMs under privileged supervision. By integrating overlap-sensitive local objectives with global drift control—both measured in a unified information-geometric space—the method not only avoids the OOD collapse characteristic of KL-based approaches but also leverages privileged context for tangible transfer gains. These results suggest that future work in distillation, RLHF, and chain-of-thought compression should consider the geometry of predictive distributions as a central design axis.