# Moment Alignment: Unifying Gradient and Hessian Matching for Domain Generalization

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

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 44 of its 66 references that are in the map.
- First posted on arXiv 9 Jun 2025.

## Builds on (most cited first, 30 of 44)

- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Deep Learning Face Attributes in the Wild](https://mlsanity.com/paper/1411.7766/) (2014)
- [Domain-Adversarial Training of Neural Networks](https://mlsanity.com/paper/1505.07818/) (2015)
- [A Broad-Coverage Challenge Corpus for Sentence Understanding through Inference](https://mlsanity.com/paper/1704.05426/) (2017)
- [Invariant Risk Minimization](https://mlsanity.com/paper/1907.02893/) (2019)
- [Moment Matching for Multi-Source Domain Adaptation](https://mlsanity.com/paper/1812.01754/) (2018)
- [Learning Transferable Features with Deep Adaptation Networks](https://mlsanity.com/paper/1502.02791/) (2015)
- [Adversarial Discriminative Domain Adaptation](https://mlsanity.com/paper/1702.05464/) (2017)
- [Distributionally Robust Neural Networks for Group Shifts: On the Importance of Regularization for Worst-Case Generalization](https://mlsanity.com/paper/1911.08731/) (2019)
- [Deep CORAL: Correlation Alignment for Deep Domain Adaptation](https://mlsanity.com/paper/1607.01719/) (2016)
- [CyCADA: Cycle-Consistent Adversarial Domain Adaptation](https://mlsanity.com/paper/1711.03213/) (2017)
- [In Search of Lost Domain Generalization](https://mlsanity.com/paper/2007.01434/) (2020)
- [Deeper, Broader and Artier Domain Generalization](https://mlsanity.com/paper/1710.03077/) (2017)
- [Out-of-Distribution Generalization via Risk Extrapolation (REx)](https://mlsanity.com/paper/2003.00688/) (2020)
- [Annotation Artifacts in Natural Language Inference Data](https://mlsanity.com/paper/1803.02324/) (2018)
- [Recognition in Terra Incognita](https://mlsanity.com/paper/1807.04975/) (2018)
- [Domain Generalization via Invariant Feature Representation](https://mlsanity.com/paper/1301.2115/) (2013)
- [How to train your ViT? Data, Augmentation, and Regularization in Vision Transformers](https://mlsanity.com/paper/2106.10270/) (2021)
- [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)
- [Domain Generalization for Object Recognition with Multi-task Autoencoders](https://mlsanity.com/paper/1508.07680/) (2015)
- [Invariance Principle Meets Information Bottleneck for Out-of-Distribution Generalization](https://mlsanity.com/paper/2106.06607/) (2021)
- [Central Moment Discrepancy (CMD) for Domain-Invariant Representation Learning](https://mlsanity.com/paper/1702.08811/) (2017)
- [Invariant Risk Minimization Games](https://mlsanity.com/paper/2002.04692/) (2020)
- [Fishr: Invariant Gradient Variances for Out-of-Distribution Generalization](https://mlsanity.com/paper/2109.02934/) (2021)
- [Quantifying the effects of data augmentation and stain color normalization in convolutional neural networks for computational pathology](https://mlsanity.com/paper/1902.06543/) (2019)
- [Invariant Causal Prediction for Nonlinear Models](https://mlsanity.com/paper/1706.08576/) (2017)
- [Learning explanations that are hard to vary](https://mlsanity.com/paper/2009.00329/) (2020)
- [Dark Model Adaptation: Semantic Image Segmentation from Daytime to Nighttime](https://mlsanity.com/paper/1810.02575/) (2018)
- [Does Invariant Risk Minimization Capture Invariance?](https://mlsanity.com/paper/2101.01134/) (2021)

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

- [Structured Credal Learning](https://mlsanity.com/paper/2603.14070/) (2026)
- [Invariant Representation Learning for Source-Free Time Series Forecasting with LLM-Centric Proxy Denoising](https://mlsanity.com/paper/2510.05589/) (2025)

## Abstract (arXiv)

Domain generalization (DG) seeks to develop models that generalize well to unseen target domains, addressing the prevalent issue of distribution shifts in real-world applications. One line of research in DG focuses on aligning domain-level gradients and Hessians to enhance generalization. However, existing methods are computationally inefficient and the underlying principles of these approaches are not well understood. In this paper, we develop the theory of moment alignment for DG. Grounded in textittransfer measure, a principled framework for quantifying generalizability between two domains, we first extend the definition of transfer measure to domain generalization that includes multiple source domains and establish a target error bound. Then, we prove that aligning derivatives across domains improves transfer measure both when the feature extractor induces an invariant optimal predictor across domains and when it does not. Notably, moment alignment provides a unifying understanding of Invariant Risk Minimization, gradient matching, and Hessian matching, three previously disconnected approaches to DG. We further connect feature moments and derivatives of the classifier head, and establish the duality between feature learning and classifier fitting. Building upon our theory, we introduce Closed-Form Moment Alignment (CMA), a novel DG algorithm that aligns domain-level gradients and Hessians in closed-form. Our method overcomes the computational inefficiencies of existing gradient and Hessian-based techniques by eliminating the need for repeated backpropagation or sampling-based Hessian estimation. We validate the efficacy of our approach through two sets of experiments: linear probing and full fine-tuning. CMA demonstrates superior performance in both settings compared to Empirical Risk Minimization and state-of-the-art algorithms.