Skip to main content
SpringerLink
Log in

Sharpness-Aware Minimization for Out-of-Distribution Generalization

  • Conference paper
  • First Online:
Neural Information Processing (ICONIP 2023)

Part of the book series: Communications in Computer and Information Science ((CCIS,volume 1961))

Included in the following conference series:

  • 418 Accesses

Abstract

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.

This is a preview of subscription content, log in via an institution to check access.

Access this chapter

Log in via an institution
Chapter
USD 29.95
Price excludes VAT (USA)
  • Available as PDF
  • Read on any device
  • Instant download
  • Own it forever
eBook
USD 79.99
Price excludes VAT (USA)
  • Available as EPUB and PDF
  • Read on any device
  • Instant download
  • Own it forever
Softcover Book
USD 99.99
Price excludes VAT (USA)
  • Compact, lightweight edition
  • Dispatched in 3 to 5 business days
  • Free shipping worldwide - see info

Tax calculation will be finalised at checkout

Purchases are for personal use only

Institutional subscriptions

Notes

  1. 1.

    The meanings of environment, group, and domain are all referred to the hierarchy of a dataset, which can be seen as the same definition in this paper.

References

  1. Arjovsky, M., Bottou, L., Gulrajani, I., Lopez-Paz, D.: Invariant risk minimization. arXiv preprint arXiv:1907.02893 (2019)

  2. Bisla, D., Wang, J., Choromanska, A.: Low-pass filtering SGD for recovering flat optima in the deep learning optimization landscape. In: International Conference on Artificial Intelligence and Statistics, pp. 8299–8339. PMLR (2022)

    Google Scholar 

  3. Cha, J., et al.: SWAD: domain generalization by seeking flat minima. Adv. Neural. Inf. Process. Syst. 34, 22405–22418 (2021)

    Google Scholar 

  4. Chang, S., Zhang, Y., Yu, M., Jaakkola, T.: Invariant rationalization. In: International Conference on Machine Learning, pp. 1448–1458. PMLR (2020)

    Google Scholar 

  5. Chen, Y., et al.: Pareto invariant risk minimization: Towards mitigating the optimization dilemma in out-of-distribution generalization. In: The Eleventh International Conference on Learning Representations (2023)

    Google Scholar 

  6. Creager, E., Jacobsen, J.H., Zemel, R.: Environment inference for invariant learning. In: International Conference on Machine Learning, pp. 2189–2200. PMLR (2021)

    Google Scholar 

  7. Duchi, J.C., Namkoong, H.: Learning models with uniform performance via distributionally robust optimization. Ann. Stat. 49(3), 1378–1406 (2021)

    Article  MathSciNet  MATH  Google Scholar 

  8. Foret, P., Kleiner, A., Mobahi, H., Neyshabur, B.: Sharpness-aware minimization for efficiently improving generalization. In: International Conference on Learning Representations (2021)

    Google Scholar 

  9. Gulrajani, I., Lopez-Paz, D.: In search of lost domain generalization. In: International Conference on Learning Representations (2021)

    Google Scholar 

  10. He, H., Huang, G., Yuan, Y.: Asymmetric valleys: beyond sharp and flat local minima. In: Advances in Neural Information Processing Systems, vol. 32 (2019)

    Google Scholar 

  11. Hochreiter, S., Schmidhuber, J.: Flat minima. Neural Comput. 9(1), 1–42 (1997)

    Article  MATH  Google Scholar 

  12. Hu, Z., Hong, L.J.: Kullback-leibler divergence constrained distributionally robust optimization. Available at Optimization Online 1(2), 1695–1724 (2013)

    Google Scholar 

  13. Huang, Z., et al.: Robust generalization against photon-limited corruptions via worst-case sharpness minimization. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 16175–16185 (2023)

    Google Scholar 

  14. Izmailov, P., Podoprikhin, D., Garipov, T., Vetrov, D., Wilson, A.: Averaging weights leads to wider optima and better generalization. In: 34th Conference on Uncertainty in Artificial Intelligence 2018, UAI 2018, pp. 876–885 (2018)

    Google Scholar 

  15. Kaddour, J., Liu, L., Silva, R., Kusner, M.J.: When do flat minima optimizers work? Adv. Neural. Inf. Process. Syst. 35, 16577–16595 (2022)

    Google Scholar 

  16. Kamath, P., Tangella, A., Sutherland, D., Srebro, N.: Does invariant risk minimization capture invariance? In: International Conference on Artificial Intelligence and Statistics, pp. 4069–4077. PMLR (2021)

    Google Scholar 

  17. Keskar, N.S., Mudigere, D., Nocedal, J., Smelyanskiy, M., Tang, P.T.P.: On large-batch training for deep learning: Generalization gap and sharp minima. In: International Conference on Learning Representations (2017)

    Google Scholar 

  18. Kim, M., Li, D., Hu, S.X., Hospedales, T.: Fisher SAM: information geometry and sharpness aware minimisation. In: International Conference on Machine Learning, pp. 11148–11161. PMLR (2022)

    Google Scholar 

  19. Krueger, D., et al.: Out-of-distribution generalization via risk extrapolation (rex). In: International Conference on Machine Learning, pp. 5815–5826. PMLR (2021)

    Google Scholar 

  20. Lin, Y., Dong, H., Wang, H., Zhang, T.: Bayesian invariant risk minimization. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 16021–16030 (2022)

    Google Scholar 

  21. Liu, Z., Luo, P., Wang, X., Tang, X.: Deep learning face attributes in the wild. In: Proceedings of the IEEE International Conference on Computer Vision, pp. 3730–3738 (2015)

    Google Scholar 

  22. McAllester, D.A.: Pac-bayesian model averaging. In: Proceedings of the Twelfth Annual Conference on Computational Learning Theory, pp. 164–170 (1999)

    Google Scholar 

  23. Petzka, H., Kamp, M., Adilova, L., Sminchisescu, C., Boley, M.: Relative flatness and generalization. In: In: Advances in Neural Information Processing Systems (2021)

    Google Scholar 

  24. Rame, A., Dancette, C., Cord, M.: Fishr: invariant gradient variances for out-of-distribution generalization. In: International Conference on Machine Learning, pp. 18347–18377. PMLR (2022)

    Google Scholar 

  25. Rame, A., Kirchmeyer, M., Rahier, T., Rakotomamonjy, A., patrick gallinari, Cord, M.: Diverse weight averaging for out-of-distribution generalization. In: Advances in Neural Information Processing Systems (2022)

    Google Scholar 

  26. Rangwani, H., Aithal, S.K., Mishra, M., Jain, A., Radhakrishnan, V.B.: A closer look at smoothness in domain adversarial training. In: International Conference on Machine Learning, pp. 18378–18399. PMLR (2022)

    Google Scholar 

  27. Rosenfeld, E., Ravikumar, P., Risteski, A.: The risks of invariant risk minimization. arXiv preprint arXiv:2010.05761 (2020)

  28. Sagawa, S., Koh, P.W., Hashimoto, T.B., Liang, P.: Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generalization. In: International Conference on Learning Representations (ICLR) (2020)

    Google Scholar 

  29. Sagawa, S., Raghunathan, A., Koh, P.W., Liang, P.: An investigation of why overparameterization exacerbates spurious correlations. In: International Conference on Machine Learning, pp. 8346–8356. PMLR (2020)

    Google Scholar 

  30. Vapnik, V.: Statistical Learning Theory. Wiley, New York (1998)

    Google Scholar 

  31. Ye, N., et al.: OoD-bench: quantifying and understanding two dimensions of out-of-distribution generalization. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 7947–7958 (2022)

    Google Scholar 

  32. Yong, L., Zhu, S., Tan, L., Cui, P.: Zin: When and how to learn invariance without environment partition? In: Advances in Neural Information Processing Systems (2022)

    Google Scholar 

  33. Zhang, Y., Sharma, P., Ram, P., Hong, M., Varshney, K.R., Liu, S.: What is missing in IRM training and evaluation? challenges and solutions. In: The Eleventh International Conference on Learning Representations (2023)

    Google Scholar 

  34. Zhou, X., Lin, Y., Zhang, W., Zhang, T.: Sparse invariant risk minimization. In: International Conference on Machine Learning, pp. 27222–27244. PMLR (2022)

    Google Scholar 

Download references

Acknowledgements

This work was supported by the Fundamental Research Funds for the Central Universities of China (2022JBMC009), the Natural Science Foundation of China (61972027), and the Beijing Municipal Natural Science Foundation (Grant No. 4212041).

Author information

Authors and Affiliations

  1. Beijing Jiaotong University, Beijing, China

    Dongqi Li & Zhu Teng

  2. AFCtech, Kunshan, China

    Qirui Li & Ziyin Wang

Authors
  1. Dongqi Li
    View author publications

    You can also search for this author in PubMed Google Scholar

  2. Zhu Teng
    View author publications

    You can also search for this author in PubMed Google Scholar

  3. Qirui Li
    View author publications

    You can also search for this author in PubMed Google Scholar

  4. Ziyin Wang
    View author publications

    You can also search for this author in PubMed Google Scholar

Corresponding author

Correspondence to Zhu Teng .

Editor information

Editors and Affiliations

  1. School of Automation, Central South University, Changsha, China

    Biao Luo

  2. Institute of Automation, Chinese Academy of Sciences, Beijing, China

    Long Cheng

  3. Institute of Cyber-Systems and Control, Zhejiang University, Hangzhou, China

    Zheng-Guang Wu

  4. School of Automation, Guangdong University of Technology, Guangzhou, China

    Hongyi Li

  5. School of Electrical Engineering and Telecommunications, UNSW Sydney, Sydney, NSW, Australia

    Chaojie Li

Rights and permissions

Reprints and permissions

Copyright information

© 2024 The Author(s), under exclusive license to Springer Nature Singapore Pte Ltd.

About this paper

Check for updates. Verify currency and authenticity via CrossMark

Cite this paper

Li, D., Teng, Z., Li, Q., Wang, Z. (2024). Sharpness-Aware Minimization for Out-of-Distribution Generalization. In: Luo, B., Cheng, L., Wu, ZG., Li, H., Li, C. (eds) Neural Information Processing. ICONIP 2023. Communications in Computer and Information Science, vol 1961. Springer, Singapore. https://doi.org/10.1007/978-981-99-8126-7_43

Download citation

  • DOI: https://doi.org/10.1007/978-981-99-8126-7_43

  • Published:

  • Publisher Name: Springer, Singapore

  • Print ISBN: 978-981-99-8125-0

  • Online ISBN: 978-981-99-8126-7

  • eBook Packages: Computer ScienceComputer Science (R0)

Publish with us

Policies and ethics

Search

Navigation