Share your thoughts, 1 month free Claude Pro on usSee more
WorkDL logo mark

Invariant Gradient Alignment for Robust Reasoning Distillation

About

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.

Zehua Cheng, Wei Dai, Jiahao Sun• 2026

Related benchmarks

TaskDatasetResultRank
Mathematical ReasoningMATH OOD
Accuracy71.6
48
Logical reasoningLogiQA2
Accuracy79.2
24
ReasoningARB ID
Accuracy (ARB ID)82.1
10
ReasoningARB OOD
Accuracy79.4
10
ReasoningReClor (ID)
Accuracy (ReClor ID)86.3
10
ReasoningReClor OOD
Accuracy82.1
10
ReasoningMATH X-Domain (ID)
Accuracy76.4
10
ReasoningReasoning Multi-Dataset Suite (ARB, LogiQA 2.0, ReClor, MATH)
Logical Consistency Score (LCS)310
10
Showing 8 of 8 rows

Other info

Follow for update