Geometric Self-Distillation for Reasoning Generalization
AI 深度解读
背景
大规模语言模型（LLM）在数学推理等结构化任务上通常通过后训练阶段来提升性能，其中on-policy蒸馏是一种实用的方法：教师模型为学生模型在其自己的生成轨迹上提供密集的监督信号。在一种特殊的“特权上下文自蒸馏”设置中，教师和学生是同一个模型，且共享相同的前缀（prefix），但教师额外看到提示（hint）或完整的解轨迹。这种设置使监督信号非常丰富，但同时也引入了可靠性问题——教师所确信的续写在其特权视角下很明显，但学生凭当前能力还无法证明这些续写的合理性。蒸馏最强的拉拽恰好发生在教师和学生分歧最大的地方，经过多次更新后，这种分歧会累积成“漂移”（drift），最终损害模型在分布外（OOD）推理上的表现。
核心内容
针对上述问题，论文提出了GeoSD（Geometric Self-Distillation），一种几何自蒸馏目标函数。GeoSD将蒸馏造成的漂移视为学生在预测行为上的移动，并通过两种互补方式加以抑制：
-
Hellinger损失缩放：每个教师偏好（即教师给出的下一词分布）按照学生当前已与教师共享的重叠程度进行缩放。具体地，使用Hellinger距离（一种衡量概率分布差异的几何量）来度量学生分布与教师分布的重叠。重叠越小的token，其教师偏好的拉拽作用越弱，从而避免学生在尚未能支持的token上被强行拉向教师。
-
近端惩罚项：即使经过Hellinger缩放，拉拽仍然会在训练中累积。为此，GeoSD增加一个近端项，惩罚学生预测相对于某个最近检查点（recent checkpoint）的漂移程度。这个距离用Fisher-Rao距离（一种定义在分布空间上的黎曼几何距离）来度量。
两项损失都建立在同一个“下一词分布”的几何结构上，因此自然梯度更新直接在该几何中移动，而不是在参数空间中。这使得更新方向与分布的统计流形吻合，而非被参数空间中的曲率扭曲。
实验在数学推理基准上展开，覆盖三个模型家族（包括从1.7B到32B的不同规模）。结果表明：
- GeoSD保留了自蒸馏带来的分布内（in-distribution）收益；
- 同时将分布外（OOD）准确率平均提升5.7–8.6个百分点（相对于基础模型），且这一增益在所有模型规模上稳定存在。
论文进一步分析了标准蒸馏（即直接最小化KL散度或交叉熵）为何在分布外失败：标准蒸馏为了迎合教师的分布，会通过从高熵状态（high-entropy states）的候选token中抽取概率质量来达成与教师的“一致”，导致模型自信地输出错误答案，而把其他正确可能性排除在外。GeoSD则通过几何约束，使这些候选token保持在可及的范围内，从而维护了更好的泛化能力。
关键要点
- 问题定位：特权上下文自蒸馏中，教师监督不可靠，学生被迫同意教师“看得到但自己还解释不了”的续写，导致分布外推理能力下降。
- 核心方案：GeoSD通过Hellinger损失按教师-学生分布重叠程度缩放每个token的蒸馏拉拽，并加入Fisher-Rao距离作为近端项，抑制累积漂移。
- 更新方式：自然梯度在概率分布几何中更新，而非参数空间，使优化与分布流形一致。
- 实验结论：在数学推理任务上，GeoSD在不损失分布内性能的前提下显著提升分布外准确率（5.7–8.6 pp），且效果跨模型规模（1.7B–32B）稳定。
- 失败机制分析：标准蒸馏在高熵状态下从替代token中抽取概率质量以迎合教师，导致错误答案自信度升高；GeoSD通过几何约束保留了这些替代选项。
意义与影响
GeoSD为自蒸馏中的信任-收益权衡提供了一种理论清晰且实践有效的几何解决方案。它首次将蒸馏过程中的漂移视为分布空间中的几何运动，并用两种自然距离（Hellinger和Fisher-Rao）分别控制瞬时拉拽和长期漂移，避免了参数空间更新带来的扭曲。该工作也揭示了标准蒸馏在分布外场景下失败的根本原因——高熵状态下的概率质量虹吸效应，为后续设计更鲁棒的自我改进算法提供了理论基础。从工程角度看，GeoSD无需额外模型或数据，仅需保存一个最近检查点，即可在现有蒸馏框架内集成，具有较低部署门槛。在数学推理这一需要强泛化能力的领域，GeoSD能有效弥合分布内与分布外的性能鸿沟，有望成为LLM后训练的标准组件。