Moment Alignment: Unifying Gradient and Hessian Matching for Domain Generalization

2025-06-09

AI Summary

Domain generalization (DG) confronts distributional shifts on unseen target domains. Existing gradient- or Hessian-alignment methods suffer from high computational overhead and weak theoretical foundations. This paper proposes a moment-alignment framework grounded in transfer metric theory. First, it establishes the first theoretical bound on transfer error across multiple source domains. Second, it reveals a duality between feature moments and classifier gradients, unifying three major paradigms—Invariant Risk Minimization (IRM), gradient matching, and Hessian matching—under a single theoretical lens. Third, it introduces a closed-form, domain-level moment alignment algorithm that jointly optimizes features and classifiers. Evaluated on standard DG benchmarks, our method significantly outperforms Empirical Risk Minimization (ERM) and state-of-the-art approaches, while achieving substantial inference speedup by eliminating repeated backpropagation and stochastic Hessian estimation. The framework thus delivers both theoretical rigor and engineering efficiency.

Technology Category

• Machine Learning: Transfer, Domain Adaptation, Multi-Task Learning
• Computer Vision: Adversarial Attacks & Robustness
• Multiagent Systems: Adversarial Agents

Application Category

• Search and Retrieval-Augmented AI: Vertical and domain-specific search
• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• Responsible Web: Human-perceived consequences of algorithmic deployment on the web

Abstract

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 extit{transfer 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 extbf{C}losed-Form extbf{M}oment extbf{A}lignment (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.

Problem

• Addressing distribution shifts in domain generalization models
• Unifying gradient and Hessian matching for computational efficiency
• Improving transfer measure through moment alignment theory

Innovation

• Aligns domain-level gradients and Hessians
• Introduces Closed-Form Moment Alignment (CMA)
• Unifies gradient and Hessian matching efficiently