Exploring Flat Minima for Domain Generalization with Large Learning Rates

TL;DR

This paper introduces a novel training strategy using large learning rates and weight interpolation to better find flat minima, improving domain generalization performance in classification and segmentation tasks.

Contribution

It proposes the Lookahead training strategy with weight interpolation to effectively identify flat minima using large learning rates, enhancing domain generalization.

Findings

Achieves state-of-the-art results on DG benchmarks.

Demonstrates effectiveness of large learning rates with weight interpolation.

Provides new insights into flat minima identification for DG.

Abstract

Domain Generalization (DG) aims to generalize to arbitrary unseen domains. A promising approach to improve model generalization in DG is the identification of flat minima. One typical method for this task is SWAD, which involves averaging weights along the training trajectory. However, the success of weight averaging depends on the diversity of weights, which is limited when training with a small learning rate. Instead, we observe that leveraging a large learning rate can simultaneously promote weight diversity and facilitate the identification of flat regions in the loss landscape. However, employing a large learning rate suffers from the convergence problem, which cannot be resolved by simply averaging the training weights. To address this issue, we introduce a training strategy called Lookahead which involves the weight interpolation, instead of average, between fast and slow weights. The fast weight explores the weight space with a large learning rate, which is not converged while the slow weight interpolates with it to ensure the convergence. Besides, weight interpolation also helps identify flat minima by implicitly optimizing the local entropy loss that measures flatness. To further prevent overfitting during training, we propose two variants to regularize the training weight with weighted averaged weight or with accumulated history weight. Taking advantage of this new perspective, our methods achieve state-of-the-art performance on both classification and semantic segmentation domain generalization benchmarks. The code is available at https://github.com/koncle/DG-with-Large-LR.