[论文解读] SWAD: Domain Generalization by Seeking Flat Minima

Junbum Cha, Sanghyuk Chun+5 more|arXiv (Cornell University)|Feb 17, 2021

Domain Adaptation and Few-Shot Learning参考文献 60被引用 4

一句话总结

本文提出SWAD，一种新颖的域泛化方法，通过在损失曲面中寻找平坦极小值来提升模型鲁棒性。通过引入密集的、防过拟合的随机权重平均，SWAD在五个主要DG基准上实现最先进性能，相较于先前最先进方法，平均域外准确率提升1.6个百分点。

ABSTRACT

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.

研究动机与目标

• 解决机器学习中的域偏移问题，即训练数据与测试数据分布存在显著差异。
• 克服标准经验风险最小化（ERM）的局限性，后者通常收敛于尖锐极小值，在域偏移下泛化能力差。
• 从理论上和实证上证明平坦极小值在域泛化（DG）场景中能带来更好的泛化性能。
• 提出一种简单而有效的方法——随机权重平均密集化（Stochastic Weight Averaging Densely, SWAD），在不修改网络架构的前提下增强平坦度与泛化能力。
• 通过与数据增强和一致性正则化方法对比，验证性能提升源于平坦度，而非域内泛化能力的改善。

提出的方法

• 提出一种鲁棒风险最小化（RRM）公式，通过参数空间邻域内的最坏情况经验风险来上界化域泛化差距。
• 理论上证明平坦极小值能减小域泛化差距，通过限制最优解邻域内的最坏情况风险。
• 通过在每次训练迭代中密集采样权重，对随机权重平均（SWA）进行改进，以更好地探索损失曲面中的平坦区域。
• 引入一种防过拟合策略，利用验证损失确定平均的最优起始和结束迭代，避免过拟合。
• 将SWAD作为即插即用模块应用于现有DG方法，无需架构修改，实现一致的性能提升。
• 利用CPU内存存储中间权重，最小化GPU内存开销，同时保持训练效率。

实验结果

研究问题

• RQ1 在非凸、复杂的深度学习损失曲面中，寻找平坦极小值是否能显著减小域泛化差距？
• RQ2 在域偏移场景下，当分布偏移比i.i.d.设置更严重时，平坦极小值带来的泛化性能提升是否依然成立？
• RQ3 像SWAD这样简单且基于平坦度感知的优化方法，能否在不修改架构的前提下超越复杂且任务特定的DG方法？
• RQ4 SWAD的性能提升是源于域内泛化能力的改善，还是源于域外鲁棒性的增强？
• RQ5 在域泛化背景下，SWAD与其它平坦度感知方法（如SAM、SWA）以及域内泛化技术（如Mixup、CutMix）相比表现如何？

主要发现

• SWAD在五个主要DG基准上达到最先进性能：PACS（+2.6pp）、VLCS（+1.6pp）、OfficeHome（+4.1pp）、TerraIncognita（+3.9pp）和DomainNet（+5.6pp），相较ERM。
• 平均而言，SWAD相较最佳现有SOTA方法将域外准确率提升1.6个百分点，相较ERM基线提升3.6个百分点。
• 将SWAD与先前SOTA方法（SOTA [31]）结合，可进一步提升性能，达到67.3%的平均准确率，较SWAD单独使用高出0.4个百分点。
• SWAD始终能找到比原始SWA更平坦的极小值，该结论通过损失曲面可视化和平坦度度量得到验证。
• 数据增强和一致性正则化方法（如Mixup、CutMix）无法提升域外泛化性能，而平坦度感知方法（SWA、SAM）可以，证实平坦度是关键因素。
• SWAD的运行时开销仅为ERM的1.07倍至1.27倍，且无额外GPU内存成本，适用于实际部署。