Sharpness-Aware Minimization for Out-of-Distribution Generalization

Machine learning models often suffer from a significant decline in performance when they encounter out-of-distribution (OOD) data that differs from the training distribution. The distribution shift can be broadly categorized into diversity shift and correlation shift. While seeking a flat minima in optimization has been shown to improve a neural network’s generalization performance with the assumption of independent and identical distribution (IID), it also has been shown to be an effective strategy for improving OOD generalization. However, previous studies potentially focused on addressing diversity shift, leaving the relationship between flat minima and correlation shift unresolved. To address the issue, we propose Sharpness-aware Invariant Risk Minimization (SIRM) as a novel approach to enhance generalization under correlation shift. Our method combines two parts: (1) Invariant risk minimization (IRM), which learns invariant relationships across multiple training environments, and (2) Sharpness-aware minimization (SAM), which finds a flat minima. Our analysis reveals that IRM does not guarantee flat minima and SAM does not improve the generalization in OOD. Moreover, we also analyze the relationship between flat minima and OOD data under correlation shift. Through extensive experiments conducted on image classification datasets, we demonstrate that our proposed method outperforms other methods with a competitive margin.