Achieving Distributional Robustness with Group-Wise Flat Minima

Improving robustness under distributional shifts remains a central challenge in machine learning. Although Sharpness-Aware Minimization (SAM) has proven effective in finding flatter minima for better generalization, it overlooks the heterogeneity in sharpness across different subpopulations, which can exacerbate performance gaps for minority or vulnerable groups. To address this challenge, we propose Group-gap Guided SAM (G2-SAM), a new optimization framework that promotes distributional robustness by steering flatness-seeking directions according to intergroup loss disparities. Our method estimates group-wise sharpness and adaptively refines perturbation strategies to minimize the worst-group loss while preserving model consistency. Through comprehensive experiments across various datasets, we show that G2-SAM achieves superior Worst-Group Accuracy and robustness, outperforming previous baselines. These findings highlight the importance of addressing group-specific geometry in the loss landscape to build more reliable and equitable neural networks.