Invariant Gradient Alignment for Robust Reasoning Distillation
On the map: LLM reasoning → Test-time compute and efficient thinking
Abstract (arXiv)
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 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.
Key facts. Cited by 1 paper in the map (data through 7 Oct 2026). Placed in Test-time compute and efficient thinking, one of 17 directions in LLM reasoning. Builds on 24 of its 39 references that are in the map. First posted on arXiv 3 Jun 2026.
Builds on
24 of 39 references.
Closest ancestors · 1
Same direction, newest first
- 2022Large Language Models Can Self-Improve20 Oct
Other references · 23
- 2021High-Resolution Image Synthesis with Latent Diffusion Modelscited by 18,694
- 2023GPT-4 Technical Reportcited by 18,132
- 2023LLaMA: Open and Efficient Foundation Language Modelscited by 13,958
- 2022Chain-of-Thought Prompting Elicits Reasoning in Large Language Modelscited by 12,302
- 2021Training Verifiers to Solve Math Word Problemscited by 9,743
- 2021Measuring Mathematical Problem Solving With the MATH Datasetcited by 5,203
- 2016SGDR: Stochastic Gradient Descent with Warm Restartscited by 4,993
- 2019Invariant Risk Minimizationcited by 1,832
- 2019Distributionally Robust Neural Networks for Group Shifts: On the Importance of Regularization for Worst-Case Generalizationcited by 1,232
- 2020Gradient Surgery for Multi-Task Learningcited by 1,210
- 2020Out-of-Distribution Generalization via Risk Extrapolation (REx)cited by 737
- 2015Causal inference using invariant prediction: identification and confidence intervalscited by 533
11 more
- 2023Orca: Progressive Learning from Complex Explanation Traces of GPT-4cited by 369
- 2022Large Language Models Are Reasoning Teacherscited by 318
- 2020The Risks of Invariant Risk Minimizationcited by 262
- 2021Gradient Matching for Domain Generalizationcited by 258
- 2020ReClor: A Reading Comprehension Dataset Requiring Logical Reasoningcited by 211
- 2020The Next Decade in AI: Four Steps Towards Robust Artificial Intelligencecited by 194
- 2020Learning explanations that are hard to varycited by 132
- 2025General-Reasoner: Advancing LLM Reasoning Across All Domainscited by 132
- 2023ARB: Advanced Reasoning Benchmark for Large Language Modelscited by 60
- 2021SAND-mask: An Enhanced Gradient Masking Strategy for the Discovery of Invariances in Domain Generalizationcited by 50
- 2025Robust Invariant Representation Learning by Distribution Extrapolationcited by 1
Built on by
Citing papers per year: 2026: 1.
1 paper.
Similar, not cited
Close in wording, no citation either way
- 2025ENIGMA: The Geometry of Reasoning and Alignment in Large-Language Models
- 2026Not All Preferences Deserve Gradients: Understanding Gradient Utility in Offline Reasoning Alignment
- 2026GRACE: Gradient-aligned Reasoning Data Curation for Efficient Post-training
- 2026Invariant Reasoning Directions in Latent Trajectories of Language Models
- 2025InfiAlign: A Scalable and Sample-Efficient Framework for Aligning LLMs to Enhance Reasoning Capabilities
- 2026Invariant-Measure Reasoners: Stable Representations for Latent Reasoning
- 2025Turning Drift into Constraint: Robust Reasoning Alignment in Non-Stationary Multi-Stream Environments
- 2026Long-Chain Reasoning Distillation via Adaptive Prefix Alignment