PaperlayerView exact revision on arXiv

Source-linked AI summary

SWAD: Domain Generalization by Seeking Flat Minima

Junbum Cha, Sanghyuk Chun, Kyungjae Lee, Han-Cheol Cho, Seunghyun Park, Yunsung Lee, Sungrae Park

arXiv:2102.08604v4cs.LGcs.CV

TL;DR

Domain generalization must handle shifts between source and unseen target domains, but ERM can converge to sharp minima with sub-optimal generalizability. The paper theoretically links flatter minima to smaller DG gaps and proposes SWAD, whose dense, overfit-aware weight sampling achieves strong results across five benchmarks. SWAD also improves existing DG methods without modifying their training objectives or architectures.

  • Problem

    Domain generalization seeks models that generalize from source domains to unseen targets, while ERM can produce sub-optimal generalizability by seeking sharp minima.

  • Method

    SWAD modifies stochastic weight averaging with dense, overfit-aware sampling guided by validation loss to find flatter minima.

  • Results

    +1.6pp average out-of-domain performance over existing best methods is reported across five DG benchmarks.

  • Takeaways & Limitations

    Seeking flat minima supports DG, and SWAD can be applied plug-and-play to existing DG methods for further performance improvements.

  • Takeaways & Limitations

    Theorem 1's confidence bound diverges as γ approaches zero because of looseness in the union bound used to derive it.

Abstract

from arXiv · show

Domain generalization (DG) methods aim to achieve generalizability to an unseen target domain by using only training data from the source domains. Although a variety of DG methods have been proposed, a recent study shows that under a fair evaluation protocol, called DomainBed, the simple empirical risk minimization (ERM) approach works comparable to or even outperforms previous methods. Unfortunately, simply solving ERM on a complex, non-convex loss function can easily lead to sub-optimal generalizability by seeking sharp minima. In this paper, we theoretically show that finding flat minima results in a smaller domain generalization gap. We also propose a simple yet effective method, named Stochastic Weight Averaging Densely (SWAD), to find flat minima. SWAD finds flatter minima and suffers less from overfitting than does the vanilla SWA by a dense and overfit-aware stochastic weight sampling strategy. SWAD shows state-of-the-art performances on five DG benchmarks, namely PACS, VLCS, OfficeHome, TerraIncognita, and DomainNet, with consistent and large margins of +1.6% averagely on out-of-domain accuracy. We also compare SWAD with conventional generalization methods, such as data augmentation and consistency regularization methods, to verify that the remarkable performance improvements are originated from by seeking flat minima, not from better in-domain generalizability. Last but not least, SWAD is readily adaptable to existing DG methods without modification; the combination of SWAD and an existing DG method further improves DG performances. Source code is available at https://github.com/khanrc/swad.

1 Introduction

Domain generalization addresses distribution shifts between source and unseen target domains, where traditional approaches can fail. SWAD targets flatter minima and outperforms existing DG methods across five benchmarks.

  • Distribution shifts can make training and test data differ, challenging systems that rely on cues such as geography, demographics, texture, or backgrounds.
  • Domain generalization trains on source domains and evaluates on different domains, with both task and domain labels available.
  • ERM performs comparably to or better than prior DG methods under the fair DomainBed evaluation protocol, despite numerous specialized approaches.
  • Sharp minima can yield sub-optimal generalizability, motivating the study of flat minima for domain generalization.
  • SWAD uses dense, overfit-aware stochastic weight sampling to find flatter minima than vanilla SWA and improve generalization to unseen domains.
  • +1.6pp average DG performance over existing best methods is reported across five benchmarks, while SWAD combined with a prior SOTA method adds 0.4pp over vanilla SWAD.

2 A Theoretical Relationship between Flatness and Domain Generalization

The paper connects domain generalization to flatness by formulating robust risk minimization over parameter neighborhoods. Its theoretical bounds relate target-domain generalization gaps to robust empirical loss, domain discrepancy, and confidence terms.

  • Motivation: ERM can have solutions with similar training losses but substantially different generalizability, while standard optimizers may select sharp minima.Domain shift can further worsen the gap between empirical and target-domain losses.
  • Robust risk minimization: RRM minimizes the worst-case empirical loss within an L2 neighborhood of radius γ around each parameter value.When γ exceeds the radius of a sharp optimum, that optimum is no longer optimal for the robust loss.
  • Generalization bound: Theorem 1 bounds test loss by robust empirical loss, train–test distribution discrepancy, and a confidence term depending on γ and the training-sample count.The analysis assumes a common labeling function across domains for simplicity.
  • Generalization bound: For the optimal RRM solution, the domain generalization gap is upper bounded by the gap between RRM and ERM, alongside train–target discrepancy and sampling-confidence terms.This establishes a direct relationship between minimizing robust empirical loss and target-domain generalization.
  • Flatness implication: With an appropriate γ, RRM selects a point near a flat ERM optimum, implying that seeking flat minima can improve the domain generalization gap.The paper uses this theoretical observation to motivate flatness-aware optimization.

3 SWAD: Domain Generalization by Seeking Flat Minima

SWAD seeks flatter minima through dense, overfit-aware weight averaging, addressing vanilla SWA’s sparse sampling and overfitting problems. Empirical analyses show SWAD identifies flatter solutions and improves domain-generalization performance across benchmark datasets.

  • 3.2 Dense and overfit-aware stochastic weight sampling strategy: SWAD modifies SWA by collecting weights at every iteration instead of every K epochs, providing denser sampling in high-dimensional parameter spaces.Vanilla SWA may average fewer than ten weights, which inadequately approximates flat minima for models with millions of parameters.
  • 3.2 Dense and overfit-aware stochastic weight sampling strategy: SWAD uses validation-loss traces to identify a start iteration ts near the first local optimum and an end iteration te after sustained overfitting.The strategy introduces optimum patience Ns, overfitting patience Ne, and tolerance rate r.
  • 3.3 Empirical analysis of SWAD and flatness: SWAD finds the lowest local flatness among ERM, SAM, SWA with cyclic learning rates, and SWA with constant learning rates across PACS experiments.Local flatness is estimated through expected loss changes on a sphere of radius γ, using 100 Monte Carlo samples.
  • 3.3 Empirical analysis of SWAD and flatness: ERM solutions lie on flat-minimum boundaries and generalize poorly, whereas SWA solutions lie near centers of training and test loss surfaces.The loss-surface visualizations use three trajectory weights and contour plots over their parameter plane.
  • 3.3 Empirical analysis of SWAD and flatness: SWAD’s overfit-aware weight gathering prevents ensemble overfitting and makes model selection-free.Figure 5 marks the sampling interval with vertical lines at ts and te on validation-accuracy curves.

4 Experiments

Experiments evaluate SWAD across five domain-generalization benchmarks, compare it with conventional and domain-generalization methods, and analyze its sampling components. SWAD consistently improves out-of-domain performance, while dense and overfit-aware sampling contribute to its gains and robustness applications.

  • Evaluation protocol: Experiments follow DomainBed-style evaluation across PACS, VLCS, OfficeHome, TerraIncognita, and DomainNet, reporting averaged out-of-domain accuracies over three runs.The benchmarks span datasets with different numbers of images, classes, and domains.
  • Main results: Compared with previous best methods, SWAD gains +2.6pp on PACS, +0.3pp on VLCS, +1.4pp on TerraIncognita, +1.9pp on OfficeHome, and +2.9pp on DomainNet.The comparison uses ResNet-50 and reports full out-of-domain performances.
  • Conventional generalization methods: On PACS, CutMix and Π-model improve in-domain performance but degrade out-of-domain performance by 1.5pp and 1.8pp, respectively.SAM slightly improves both settings, whereas its out-of-domain improvement is not statistically significant.
  • Conventional generalization methods: SWAD improves out-of-domain performance over vanilla SWA by 1.2pp and in-domain performance by 0.6pp on PACS.Both vanilla SWA and SWAD significantly improve in-domain and out-of-domain performance relative to ERM.
  • Combinations with other methods: Applying SWAD to ERM, CORAL, and SAM improves average out-of-domain accuracy by 3.6pp, 2.8pp, and 1.0pp, respectively.CORAL + SWAD achieves the best performance among the combinations reported.
  • Ablation study: Removing dense sampling reduces average out-of-domain accuracy by 0.8pp, while removing overfit-aware sampling reduces it by 0.6pp.Validation-selected averaging becomes overfitted, and fixed scheduling provides only marginal gains over vanilla SWA.
  • ImageNet robustness: SWAD consistently improves robustness on ImageNet-C, ImageNet-R, and the background challenge over both ERM and SWA.These experiments extend SWAD beyond tasks containing explicit domain labels.

5 Discussion and Limitations

The paper identifies limitations in SWAD’s theory, optimization, and use of domain-specific information. It frames these limitations as directions for improving domain generalization.

  • Theoretical limitation: Theorem 1’s confidence bound diverges as γ approaches zero, unlike the expected convergence toward ERM’s confidence error.The paper attributes this counter-intuitive behavior to looseness in the union bound, not to a drawback of RRM.
  • Optimization limitation: SWAD is a heuristic approximation for flat minima rather than a perfect or theoretically guaranteed flatness-aware solver.The theoretical relationship between flat minima and domain generalization would remain valid with a better solver.
  • Domain information: SWAD focuses primarily on flat minima and does not strongly exploit domain-specific information or domain discrepancy.The paper suggests combining flatness with domain-discrepancy reduction; CORAL plus SWAD gives the best result among the compared combinations.
  • Future direction: The paper presents stronger use of both flatness and domain discrepancy as a future direction for improving domain generalization.This proposed direction follows from the bound’s three factors: flat minima, domain discrepancy, and confidence bound.

6 Concluding Remarks

The paper concludes that seeking flat minima is a viable route to domain generalization and introduces SWAD for this purpose. Across benchmarks and robustness settings, SWAD performs strongly and can be combined with existing methods.

  • Conclusion: The paper theoretically and empirically demonstrates that domain generalization is achievable by seeking flat minima.It presents flatness as a new perspective for domain generalization and other robustness tasks.
  • Conclusion: SWAD captures flatter minima than vanilla SWA and achieves superior performance on five domain-generalization benchmarks.Combining SWAD with existing domain-generalization methods further improves performance over vanilla SWAD.
  • Practical implication: SWAD can make models robust to domain shift without relying on a task-dependent prior, modified objective function, or specific architecture.The paper describes this as a potential practical impact for deployment in changing environments.
  • Evaluation: The evaluation protocol’s computational cost motivates reducing the hyperparameter search space for efficiency.The protocol requires about 4,142 models for every domain-generalization algorithm, so the paper searches only selected parameters.

B.2 Hyperparameter search protocol for reproduced results

The reproduced-results protocol documents benchmark robustness evaluations, training settings, comparison methods, and visualization procedures used to study SWAD and related methods.

  • Hyperparameter search: The protocol reduces computational cost by searching selected hyperparameters rather than reproducing the full 4,142-run-per-algorithm search.The selected search includes SWAD’s start iteration, end iteration, and sampling interval, while several settings are fixed.
  • Comparison methods: The reproduced comparisons tune algorithm-specific settings for Mixstyle, SAM, VAT, Π-model, EMA, Mixup, and CutMix.The cited settings include insertion locations, perturbation parameters, interpolation parameters, and decay values.
  • Visualization: Weight-plane visualization constructs orthonormal bases from three model weights, evaluates losses on a Cartesian grid, and displays the results as a contour plot.This procedure is used to inspect the loss landscape around the selected weights.
  • Robustness benchmarks: The protocol compares robustness on ImageNet-C, ImageNet-R, and the background challenge using ResNet-50 and standard training recipes.ImageNet-C measures common corruptions, ImageNet-R covers varied image domains, and the background challenge tests prediction consistency under foreground-background manipulation.

C Proof of Theorems

The proofs establish error bounds by combining functional-error comparisons, finite parameter-space covers, local robust-risk bounds, and source-target discrepancy terms. These ingredients support the paper’s domain-generalization bound.

  • Preliminaries: The proof framework defines functional error through a bounded instance loss and recovers population loss when the comparator is the true label function.The loss is zero exactly when the two labels agree.
  • Covering argument: Lemma 2 bounds empirical robust risk over a finite cover of parameter space using local maxima, VC dimensions, and sample size.The bound holds with high probability and applies to each covered parameter region.
  • Domain generalization bound: The theorem proof combines single-domain inequalities with the multiple-source-domain setting and introduces divergence between source and target distributions.The source distributions are combined into a mixture before applying the domain-discrepancy bound.
  • Final bound: The final proof combines generalization bounds for empirical and target risks with Theorem 1 under a high-probability event.The resulting terms include source-target discrepancy and confidence-related quantities.

D.1 Comparison of flatness-aware solvers

SWAD achieves the strongest average ranking among flatness-aware solvers, matching the ranking from local flatness tests. Its dense, overfit-aware sampling strategy improves performance over sparse SWA variants across all benchmarks.

  • SWAD ranks first in both average flatness-aware solver performance and local flatness tests, followed by SWAs, SAM, and ERM.
  • 1.4pp is SWAD’s maximum average-performance improvement over SWAs, and it surpasses both SWA variants on every benchmark.
  • The comparison attributes SWAD’s gains over SWAs to its dense and overfit-aware stochastic weight sampling strategy.

E Full Results

The full-results section reports out-of-domain accuracies across five benchmarks and documents the evaluation protocol used for the detailed tables. It also notes associated licensing and ethical considerations for the underlying assets.

  • Results marked † and ‡ use DomainBed’s and the authors’ hyperparameter-search protocols, respectively, with standard errors from three trials when available.
  • The detailed results report out-of-domain accuracies for PACS, VLCS, OfficeHome, TerraIncognita, and DomainNet.
  • The work discusses licenses, copyrights, and ethical issues for its public datasets and code assets.

F.2 Datasets

The study uses public datasets and reports their provenance, licensing context, experimental hardware, and runtime across leave-one-out domain evaluations.

  • PACS and VLCS lack identified statements about licenses, copyrights, or consent, and their person images may raise ethical issues.
  • The experiments use public source code and specify environments, dependencies, dataset downloads, and reproduction instructions.
  • Experiments run on a single NVIDIA Tesla P40 or V100 using Python 3.8.6, PyTorch 1.7.0, Torchvision 0.8.1, and CUDA 9.2.
  • Average runtimes range from about 4 hours for PACS and VLCS to 56 hours for DomainNet, including leave-one-out cross-validation.

G.3 Complexity Analysis

The complexity analysis models SWAD’s training and evaluation costs and reports modest runtime overhead relative to ERM, while noting additional CPU-memory use.

  • With forward and backward time t, training time per epoch is 2tnr/(r + 1), evaluation time is vtn/(r + 1), and total time is tn(2r + v)/(r + 1).
  • SWAD’s total runtime is 1.07 times the ERM baseline on PACS and up to 1.27 times on DomainNet.
  • The method requires no additional GPU memory but uses CPU memory up to max(N, M) times the baseline.
Loading 2102.08604v4…