Geometric Self-Distillation for Reasoning Generalization

2026-07-07

AI Summary

While self-distillation can enhance in-distribution performance, the privileged information from the teacher model often induces behavioral drift in the student, compromising out-of-distribution generalization. This work proposes Geometric Self-Distillation (GeoSD), a method that constructs the distillation loss within the geometric space of predictive distributions. Specifically, it employs Hellinger distance–weighted supervision to suppress unreliable signals and incorporates a Fisher–Rao distance–based proximal regularizer to constrain the update trajectory, thereby enabling natural gradient optimization. Within this unified geometric framework, GeoSD effectively balances in-distribution accuracy with out-of-distribution robustness. Experiments demonstrate that across large language models ranging from 1.7B to 32B parameters, GeoSD consistently preserves in-distribution performance while improving out-of-distribution accuracy by an average of 5.7–8.6 percentage points on multiple mathematical reasoning benchmarks.

Abstract

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.

Problem

• self-distillation
• reasoning generalization
• out-of-distribution
• prediction drift
• privileged-context

Innovation

• Geometric Self-Distillation
• Hellinger loss
• Fisher-Rao distance
• natural gradient
• out-of-distribution generalization