IGA:通过不变梯度对齐提升推理蒸馏的鲁棒性

Invariant Gradient Alignment for Robust Reasoning Distillation

精选理由

做推理蒸馏或知识迁移的团队,IGA解决了小模型在分布外数据上泛化差的痛点,逻辑一致性提升四倍,值得在数学、医学等跨领域任务上试试。

AI 摘要

大型语言模型存在捷径学习问题,在分布外输入上系统性失败,即使逻辑结构相同。本文提出不变梯度对齐(IGA)框架,通过逻辑同构集、连续梯度冲突掩码和截断SVD投影三个创新,对齐语义多样但逻辑同构样本的梯度更新。理论上,IGA比经验风险最小化(ERM)有更紧的分布外泛化界;实验上,在四个基准上准确率提升最高14.3个百分点,逻辑一致性分数提升四倍。该方法适用于知识蒸馏场景,帮助小模型学习更鲁棒的推理能力。

原文 · arXiv cs.AI

Invariant Gradient Alignment for Robust Reasoning Distillation

Large language models (LLMs) suffer from shortcut learning: they systematically fail on out-of-distribution (OOD) inputs whose semantic surface differs from training data, even when the logical structure is identical. This undermines knowledge distillation pipelines that transfer chain-of-thought reasoning to smaller students. We introduce Invariant Gradient Alignment (IGA), a training framework that aligns gradient updates across semantically diverse but logically isomorphic examples via three innovations: (i) Logical Isomer Sets, groups of problems sharing identical logical structure across distinct semantic domains (mathematics, medicine, law, science); (ii) a differentiable \emph{Continuous Gradient Conflict Mask}, that suppresses parameter dimensions with high cross-domain gradient variance while preserving invariant directions; and (iii) a truncated SVD projection of the masked gradient back onto the LoRA low-rank manifold, maintaining parameter efficiency throughout. Theoretically, IGA yields tighter OOD generalization bounds than ERM, scaling with the number of isomer domains, and converges at the standard SGD rate under mild regularity. Empirically, IGA outperforms eight baselines across four benchmarks with accuracy gains up to 14.3 pp over ERM-SFT and a Logical Consistency Score of 0.031 versus 0.142 -- a fourfold improvement in representational invariance.