# Flatness and Gradient Alignment Are Both Necessary: Spectral-Aware Gradient-Aligned Exploration for Multi-Distribution Learning

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 59 of its 72 references that are in the map.
- First posted on arXiv 8 May 2026.

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

- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [The Cityscapes Dataset for Semantic Urban Scene Understanding](https://mlsanity.com/paper/1604.01685/) (2016)
- [Domain-Adversarial Training of Neural Networks](https://mlsanity.com/paper/1505.07818/) (2015)
- [Invariant Risk Minimization](https://mlsanity.com/paper/1907.02893/) (2019)
- [SegNet: A Deep Convolutional Encoder-Decoder Architecture for Image Segmentation](https://mlsanity.com/paper/1511.00561/) (2015)
- [Sharpness-Aware Minimization for Efficiently Improving Generalization](https://mlsanity.com/paper/2010.01412/) (2020)
- [Multi-Task Learning Using Uncertainty to Weigh Losses for Scene Geometry and Semantics](https://mlsanity.com/paper/1705.07115/) (2017)
- [On Large-Batch Training for Deep Learning: Generalization Gap and Sharp Minima](https://mlsanity.com/paper/1609.04836/) (2016)
- [Moment Matching for Multi-Source Domain Adaptation](https://mlsanity.com/paper/1812.01754/) (2018)
- [Gradient Surgery for Multi-Task Learning](https://mlsanity.com/paper/2001.06782/) (2020)
- [Deep CORAL: Correlation Alignment for Deep Domain Adaptation](https://mlsanity.com/paper/1607.01719/) (2016)
- [Deep Hashing Network for Unsupervised Domain Adaptation](https://mlsanity.com/paper/1706.07522/) (2017)
- [Visualizing the Loss Landscape of Neural Nets](https://mlsanity.com/paper/1712.09913/) (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)
- [A Survey on Multi-Task Learning](https://mlsanity.com/paper/1707.08114/) (2017)
- [Multi-Task Learning as Multi-Objective Optimization](https://mlsanity.com/paper/1810.04650/) (2018)
- [Generalizing to Unseen Domains: A Survey on Domain Generalization](https://mlsanity.com/paper/2103.03097/) (2021)
- [Domain Generalization: A Survey](https://mlsanity.com/paper/2103.02503/) (2021)
- [Out-of-Distribution Generalization via Risk Extrapolation (REx)](https://mlsanity.com/paper/2003.00688/) (2020)
- [Recognition in Terra Incognita](https://mlsanity.com/paper/1807.04975/) (2018)
- [End-to-End Multi-Task Learning with Attention](https://mlsanity.com/paper/1803.10704/) (2018)
- [Fantastic Generalization Measures and Where to Find Them](https://mlsanity.com/paper/1912.02178/) (2019)
- [Sharp Minima Can Generalize For Deep Nets](https://mlsanity.com/paper/1703.04933/) (2017)
- [Conflict-Averse Gradient Descent for Multi-task Learning](https://mlsanity.com/paper/2110.14048/) (2021)
- [SWAD: Domain Generalization by Seeking Flat Minima](https://mlsanity.com/paper/2102.08604/) (2021)
- [Shampoo: Preconditioned Stochastic Tensor Optimization](https://mlsanity.com/paper/1802.09568/) (2018)
- [Self-Challenging Improves Cross-Domain Generalization](https://mlsanity.com/paper/2007.02454/) (2020)
- [ASAM: Adaptive Sharpness-Aware Minimization for Scale-Invariant Learning of Deep Neural Networks](https://mlsanity.com/paper/2102.11600/) (2021)

## Abstract (arXiv)

Sharpness-aware and gradient-alignment methods have been shown to improve generalization, however each family of methods targets a single geometric property of the loss landscape, while ignoring the other. In this paper, we show that this omission is structurally unavoidable and that both flatness and gradient alignment should be considered in multi-distribution learning settings. Specifically, we derive an excess-risk decomposition that yields two additive leading-order terms: (i) an alignment term, controlled by the trace of barH-1Sigmag and (ii) a curvature term, controlled by barH, where barH is the average Hessian and Sigmag is the covariance of the gradient across distributions. Notably, barH appears inverted in one and non-inverted in the other. We further show, via a counterexample, that neither quantity bounds the other in general, so no algorithm targeting only one term can guarantee low excess risk. Motivated by this decomposition, we propose SAGE (Spectral-Aware Gradient-Aligned Exploration) that targets both terms. The curvature component replaces SAM's gradient-scaled perturbation with the polar factor of each layer's gradient matrix, computed via Newton-Schulz iteration, so that the ascent step probes all directions with similar magnitude. On the other hand, the alignment component injects isotropic noise at the descent step, the magnitude of which scales with cross-distribution gradient disagreement. Experiments on five domain-generalization and two multi-task learning benchmarks show that the proposed method establishes a new state-of-the-art on DomainBed and acts as a general-purpose improvement to base MTL solvers, remaining competitive with, or even surpassing, state-of-the-art methods.