# Invariant Gradient Alignment for Robust Reasoning Distillation

arXiv: https://arxiv.org/abs/2606.05025 · HTML: https://mlsanity.com/paper/2606.05025/

## Key facts

- Cited by 1 papers in the map (data through 14 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.

## Placed in

- LLM reasoning / [Test-time compute and efficient thinking](https://mlsanity.com/lr/test-time-compute-and-efficient-thinking/)
- Knowledge distillation / [Knowledge distillation for LLMs](https://mlsanity.com/distill/knowledge-distillation-for-llms/)

## Builds on (most cited first, 24 of 24)

- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [Chain-of-Thought Prompting Elicits Reasoning in Large Language Models](https://mlsanity.com/paper/2201.11903/) (2022)
- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [Measuring Mathematical Problem Solving With the MATH Dataset](https://mlsanity.com/paper/2103.03874/) (2021)
- [SGDR: Stochastic Gradient Descent with Warm Restarts](https://mlsanity.com/paper/1608.03983/) (2016)
- [Invariant Risk Minimization](https://mlsanity.com/paper/1907.02893/) (2019)
- [Gradient Surgery for Multi-Task Learning](https://mlsanity.com/paper/2001.06782/) (2020)
- [Distributionally Robust Neural Networks for Group Shifts: On the Importance of Regularization for Worst-Case Generalization](https://mlsanity.com/paper/1911.08731/) (2019)
- [Out-of-Distribution Generalization via Risk Extrapolation (REx)](https://mlsanity.com/paper/2003.00688/) (2020)
- [Large Language Models Can Self-Improve](https://mlsanity.com/paper/2210.11610/) (2022)
- [Causal inference using invariant prediction: identification and confidence intervals](https://mlsanity.com/paper/1501.01332/) (2015)
- [Orca: Progressive Learning from Complex Explanation Traces of GPT-4](https://mlsanity.com/paper/2306.02707/) (2023)
- [Large Language Models Are Reasoning Teachers](https://mlsanity.com/paper/2212.10071/) (2022)
- [The Risks of Invariant Risk Minimization](https://mlsanity.com/paper/2010.05761/) (2020)
- [Gradient Matching for Domain Generalization](https://mlsanity.com/paper/2104.09937/) (2021)
- [ReClor: A Reading Comprehension Dataset Requiring Logical Reasoning](https://mlsanity.com/paper/2002.04326/) (2020)
- [The Next Decade in AI: Four Steps Towards Robust Artificial Intelligence](https://mlsanity.com/paper/2002.06177/) (2020)
- [General-Reasoner: Advancing LLM Reasoning Across All Domains](https://mlsanity.com/paper/2505.14652/) (2025)
- [Learning explanations that are hard to vary](https://mlsanity.com/paper/2009.00329/) (2020)
- [ARB: Advanced Reasoning Benchmark for Large Language Models](https://mlsanity.com/paper/2307.13692/) (2023)
- [SAND-mask: An Enhanced Gradient Masking Strategy for the Discovery of Invariances in Domain Generalization](https://mlsanity.com/paper/2106.02266/) (2021)
- [Robust Invariant Representation Learning by Distribution Extrapolation](https://mlsanity.com/paper/2505.16126/) (2025)

## Built on by (most cited first, 1 of 1)

- [Unknown-Traffic Detection, Calibration and Shortcut Reliance in Distilled Encrypted-Traffic Classifiers over One Year](https://mlsanity.com/paper/2609.31141/) (2026)

## 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.