Source-linked AI summary
Sharpness-Aware Gradient Matching for Domain Generalization
Pengfei Wang, Zhaoxiang Zhang, Zhen Lei, Lei Zhang
TL;DR
Domain generalization needs models that transfer to unseen domains, but SAM may not reach low-loss flat regions. SAGM jointly minimizes empirical risk, perturbed loss, and their gap through gradient matching, and it outperforms state-of-the-art DG methods across five benchmarks.
Problem
SAM and related methods may fail to converge to a flat region with a small loss, limiting reliable domain generalization.
Method
SAGM jointly minimizes empirical risk, perturbed loss, and their gap while aligning the gradient directions of the empirical and perturbed losses.
Results
SAGM demonstrates superior performance to state-of-the-art DG methods on five benchmarks, achieving 66.1% average performance on DomainBed.
Takeaways & Limitations
SAGM improves DG capability without using additional information and even outperforms Miro equipped with a pre-trained CLIP model.
Takeaways & Limitations
GSAM’s surrogate-gap minimization can increase empirical loss, making perturbed-loss optimization more difficult and potentially hurting generalization.
Abstract
from arXiv · showhide
The goal of domain generalization (DG) is to enhance the generalization capability of the model learned from a source domain to other unseen domains. The recently developed Sharpness-Aware Minimization (SAM) method aims to achieve this goal by minimizing the sharpness measure of the loss landscape. Though SAM and its variants have demonstrated impressive DG performance, they may not always converge to the desired flat region with a small loss value. In this paper, we present two conditions to ensure that the model could converge to a flat minimum with a small loss, and present an algorithm, named Sharpness-Aware Gradient Matching (SAGM), to meet the two conditions for improving model generalization capability. Specifically, the optimization objective of SAGM will simultaneously minimize the empirical risk, the perturbed loss (i.e., the maximum loss within a neighborhood in the parameter space), and the gap between them. By implicitly aligning the gradient directions between the empirical risk and the perturbed loss, SAGM improves the generalization capability over SAM and its variants without increasing the computational cost. Extensive experimental results show that our proposed SAGM method consistently outperforms the state-of-the-art methods on five DG benchmarks, including PACS, VLCS, OfficeHome, TerraIncognita, and DomainNet. Codes are available at https://github.com/Wang-pengfei/SAGM.
1. Introduction
Domain generalization seeks models that transfer from source data to unseen domains, but ERM and SAM-like objectives can fail to find low-loss flat regions. SAGM addresses this by jointly minimizing empirical risk, perturbed loss, and their gap through gradient matching.
- Domain generalization trains on source data to generalize to unseen domains, where deep-model performance can drop substantially.
- ERM can outperform many DG methods under fair evaluation, yet minimizing empirical loss alone often overfits and reaches sharp minima.
- SAM minimizes perturbed loss around the parameters, but this objective is not guaranteed to reach a flat minimum.
- The proposed conditions require sufficiently low loss near the desired minimum and a flat surrounding loss surface.The first condition supports source-domain performance, while the second reduces the training–testing performance gap.
- SAGM jointly minimizes empirical risk, perturbed loss, and their surrogate gap, using aligned gradients to address conflicts among these objectives.The method targets a low-loss region through the first two terms and avoids steep valleys through the gap term.
- The paper reports superior SAGM performance over state-of-the-art methods on five DG benchmarks.
2. Related Work
Related work covers domain generalization strategies, sharpness-aware optimization, and the relationship between sharpness and generalization. It positions SAGM against methods whose perturbed-loss or surrogate-gap objectives can limit flat-minimum discovery or increase empirical loss.
- Domain Generalization: DG methods address domain shift through alignment, meta-learning, augmentation, disentanglement, causal modeling, and gradient-based strategies.
- Sharpness-Aware Minimization: SAM minimizes the maximum loss within a parameter-space neighborhood, while Look-SAM and ESAM reduce computation but retain SAM’s perturbed-loss limitation.
- Sharpness-Aware Minimization: GSAM addresses perturbed-loss sharpness mismatch but reduces the surrogate gap by increasing empirical loss, limiting potential generalization gains.
- Sharpness and Generalization: The paper motivates SAGM as a method for converging to flatter regions with better generalization than existing sharpness-aware approaches.
- Sharpness and Generalization: Prior sharpness studies report a negative correlation between sharpness and generalization ability under the i.i.d. assumption.
3. Preliminaries
The preliminaries define empirical risk minimization over source domains and explain SAM’s neighborhood-based min-max objective. SAM approximates the costly inner maximization and optimizes perturbed loss instead of the original loss.
- Problem Setup: The DG setup uses neural-network parameters θ and multiple source domains, with training samples drawn from each domain.
- Empirical Risk Minimization: The training loss over all source domains is empirical risk, typically using cross-entropy, and ERM seeks parameters with the lowest value.
- Sharpness-Aware Minimization: ERM can overfit the training set and converge to sharp minima, motivating SAM’s search for flat regions with low loss.
- Sharpness-Aware Minimization: SAM maximizes empirical loss over perturbations within a Euclidean ball of radius ρ, then minimizes the resulting worst-case loss.
- Sharpness-Aware Minimization: For small ρ, Taylor expansion converts SAM’s inner maximization into a linearly constrained approximation.
- Sharpness-Aware Minimization: SAM optimizes the perturbed loss Lp(θ; D) rather than the original loss L(θ; D).
4. Sharpness-Aware Gradient Matching
SAGM addresses failures of SAM-like methods by seeking a low-loss, flat region and aligning gradients of empirical and perturbed losses. Its objective jointly minimizes empirical risk, perturbed loss, and their surrogate gap, while gradient matching makes these objectives descend consistently.
- 4.1. The perturbed loss Lp(θ; D) is not always sharpness-aware: SAM can select a sharper minimum because lower perturbed loss does not necessarily indicate lower sharpness.The surrogate gap h(θ) better describes sharpness in the illustrated comparison.
- 4.1. The perturbed loss Lp(θ; D) is not always sharpness-aware: The surrogate gap h(θ) measures the difference between perturbed and unperturbed losses and is linked to the dominant Hessian eigenvalue at local minima.GSAM minimizes this gap alongside perturbed loss, but can increase the empirical loss.
- 4.2. Sharpness-Aware Gradient Matching: SAGM targets two conditions: sufficiently low loss within a neighborhood and a minimum located in a flat loss surface.The first supports low source-training error, while the second reduces the training–testing performance gap.
- 4.2. Sharpness-Aware Gradient Matching: SAGM simultaneously minimizes empirical risk, perturbed loss, and surrogate gap to satisfy the low-loss and flatness conditions.Its optimization objective is designed around all three quantities rather than perturbed loss alone.
- 4.3. Analysis and algorithm of SAGM: SAGM has basically the same computational cost as SAM because its additional loss calculation matches the cost of the perturbed-loss computation.The empirical loss is already calculated when computing its gradient.
- 4.3. Analysis and algorithm of SAGM: Maximizing the inner product of empirical-risk and perturbed-loss gradients aligns their directions, reducing conflicts among the three optimization objectives.When the gradients are consistent, empirical loss, perturbed loss, and surrogate gap descend efficiently together.
5. Experiments
Experiments evaluate SAGM under DomainBed-style protocols across five domain-generalization benchmarks and compare it with conventional, sharpness-aware, gradient-based, and augmentation methods. SAGM consistently improves out-of-domain performance, while ablations and sharpness analyses support the roles of its gradient matching and flat-minimum objectives.
- 5.1. Experiment setups and implementation details: Evaluation uses five benchmarks with leave-one-out domain splits, averaging accuracy over three trials and all possible out-of-domain settings.The benchmarks are PACS, VLCS, OfficeHome, TerraIncognita, and DomainNet.
- 5.2. Main results: SAGM consistently outperforms ERM across all five benchmarks, with a +2.2% average improvement.It also surpasses the leading conventional DG method on each individual dataset.
- 5.2. Main results: SAGM outperforms SAM and GSAM on all five datasets, improving average performance by 1.6% over SAM and 1.0% over GSAM.The comparison covers leading sharpness-aware methods.
- 5.2. Main results: Across 5 datasets and 22 target domains, the experiments demonstrate SAGM's effectiveness in improving domain-generalization capability.The evaluation also compares SAGM with methods using a pre-trained CLIP model.
- 5.3. In-domain and out-of-domain performance: On PACS, data augmentation methods improve in-domain performance but can degrade out-of-domain performance, whereas SAGM+Mixstyle further improves over both SAGM and Mixstyle.Mixup, CutMix, and Mixstyle degrade out-of-domain performance by 0.5%, 1.5%, and 1.4%, respectively.
- 5.4. Ablation study: The ablation study finds SAGM improves average performance by 1.3% over ERM+SAM, supporting the importance of gradient matching.SAGM, SAM, and GSAM also outperform ERM by significant margins.
6. Conclusion
The paper proposes SAGM to address limitations of SAM-like DG algorithms by jointly minimizing empirical loss, perturbed loss, and their gap. SAGM demonstrates superior performance on five DG benchmarks, achieving 66.1% average performance on DomainBed without additional information.
- 6. Conclusion: SAGM jointly minimizes empirical loss, perturbed loss, and their gap to target a flat region with improved generalization.The method is designed to satisfy two conditions: a sufficiently low loss near the desired minimum and a flat loss surface.
- 6. Conclusion: 66.1% average performance on DomainBed was achieved without using additional information.The paper reports that SAGM even outperformed Miro, which uses a pre-trained CLIP model.
- 6. Conclusion: SAGM demonstrated superior performance to state-of-the-art DG methods, including SAM and GSAM, on five DG benchmarks.