Navigating the Flatlands: Dual Adaptive Sharpness-Aware Minimization for Domain Generalization
Abstract
Finding flat minima in the loss landscape is a key strategy for Domain Generalization (DG). However, its effectiveness is often limited by two crucial challenges. 1) Domain Shift: Existing methods like Sharpness-Aware Minimization (SAM) apply a uniform optimization strategy across all domains, overlooking the differences of the learning difficulties among multiple domains and thus performing poorly on challenging domains. 2) Anisotropic Sharpness: By perturbing parameters along a single gradient direction, SAM and its variants ignore multi-directional flatness, making the model converge to minima that remain sharp in other directions. The combined challenges make it more difficult for the model to find truly robust solutions in multi-domain scenarios. To overcome these limitations, we propose the Dual Adaptive Sharpness-Aware Minimization (DA-SAM), which comprises two key modules: Dynamic Adaptive Scaling (DAS) module and Adaptive Multi-Directional Flattening (AMDF) module. First, to tackle the domain shift problem, the DAS module computes the real-time loss on each domain to adaptively generate domain-specific scaling factors that guide the generation of perturbation directions. Second, the AMDF module calculates local flatness by generating multiple directions to simulate perturbations in the parameter space. Based on the learned local flatness metric, it dynamically adjusts the perturbation step size to guide the model parameters to be away from anisotropic sharp regions. Crucially, DAS provides domain-level guidance that makes AMDF’s multi-directional geometric exploration more targeted and effective. Extensive experiments on five DG benchmarks demonstrate the effectiveness of our DA-SAM algorithm.
Lay Summary
AI models often struggle when facing unfamiliar situations in the real world—for example, failing to recognize a sketch of a dog after only learning from real photos. To build more reliable AI, researchers train models to find "flat" and stable solutions where small disruptions won't cause catastrophic errors. However, current training methods have blind spots: they treat all training environments as equally difficult and only check for stability in a single direction, leaving hidden vulnerabilities. We tackled this by building a new training approach, DA-SAM, which acts like a careful explorer navigating a rugged landscape. First, it dynamically gauges the difficulty of each environment, adapting its strategy to handle harder scenarios more cautiously. Second, it actively probes the landscape in multiple random directions to ensure the model's solution is stable and safe from every angle. Our findings will help developers build AI systems that perform consistently and robustly when deployed in unpredictable, real-world conditions.