Search papers, labs, and topics across Lattice.
This paper introduces GeoSD, a geometric self-distillation method that mitigates the drift in predictive behavior of large language models during on-policy distillation. By employing a Hellinger loss to scale teacher preferences based on shared knowledge and a proximal term to penalize excessive deviation from recent predictions, GeoSD effectively enhances out-of-distribution reasoning capabilities. The method demonstrates a significant improvement in average OOD accuracy by 5.7-8.6 points across various model sizes and mathematical reasoning benchmarks, while maintaining in-distribution performance.
GeoSD counters the problematic drift in self-distillation, boosting out-of-distribution reasoning accuracy while preserving in-distribution gains.
On-policy distillation is a practical post-training recipe for large language models, 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.
