[论文解读] Meta-Learned Invariant Risk Minimization

Jun-Hyun Bae, Inchul Choi+1 more|arXiv (Cornell University)|Mar 24, 2021

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

一句话总结

本文提出 Meta-IRM，一种元学习方法，通过模型无关元学习（MAML）直接优化理想的双层不变风险最小化（IRM）目标，而无需假设线性分类器。通过将每个环境视为一个任务，并在内层和外层优化中使用不相交的环境，Meta-IRM 实现了更优的分布外（OOD）泛化性能——尤其在数据稀缺或虚假相关性丰富的情况下，优于 IRMv1 及其所有变体，在线性和非线性设置中均表现更优。

ABSTRACT

Empirical Risk Minimization (ERM) based machine learning algorithms have suffered from weak generalization performance on data obtained from out-of-distribution (OOD). To address this problem, Invariant Risk Minimization (IRM) objective was suggested to find invariant optimal predictor which is less affected by the changes in data distribution. However, even with such progress, IRMv1, the practical formulation of IRM, still shows performance degradation when there are not enough training data, and even fails to generalize to OOD, if the number of spurious correlations is larger than the number of environments. In this paper, to address such problems, we propose a novel meta-learning based approach for IRM. In this method, we do not assume the linearity of classifier for the ease of optimization, and solve ideal bi-level IRM objective with Model-Agnostic Meta-Learning (MAML) framework. Our method is more robust to the data with spurious correlations and can provide an invariant optimal classifier even when data from each distribution are scarce. In experiments, we demonstrate that our algorithm not only has better OOD generalization performance than IRMv1 and all IRM variants, but also addresses the weakness of IRMv1 with improved stability.

研究动机与目标

• 解决经验风险最小化（ERM）在训练数据包含虚假相关性时 OOD 泛化性能差的问题。
• 克服 IRMv1 的局限性，即当环境数量少于虚假相关性数量，或每个环境的训练数据稀缺时，IRMv1 会失效。
• 通过直接优化理想的双层 IRM 目标，而不做简化假设，开发一种更鲁棒且稳定的不变预测器学习方法。
• 通过解除分类器的线性约束，实现在非线性设置下的有效 OOD 泛化。
• 通过辅助正则化项最小化元损失方差，以及使用二阶梯度更新，提升优化稳定性和性能。

提出的方法

• 将理想的 IRM 目标形式化为双层优化问题，将每个训练环境视为元学习中的独立任务。
• 应用模型无关元学习（MAML）求解双层问题，使用不相交的环境分别进行内层（任务特定适应）和外层（元更新）优化。
• 通过允许非线性模型，放宽 IRMv1 中对线性分类器的假设，从而实现更优的表征学习和泛化能力。
• 引入辅助损失，最小化所有环境中元损失的标准差，以强化不变性并提升训练稳定性。
• 在元更新中使用二阶梯度，以更准确地逼近真实的双层优化，提升收敛性和性能。
• 通过优化分类器在多样化环境中的表现一致性，确保其在测试数据分布显著变化时仍保持不变性。

实验结果

研究问题

• RQ1 元学习框架能否在不假设线性分类器的前提下，有效优化理想的双层 IRM 目标？
• RQ2 所提出的元-IRM 方法是否在数据稀缺或虚假相关性高的情况下，优于 IRMv1 及其变体？
• RQ3 使用二阶梯度和不相交环境采样在实现稳定且不变的预测器方面有多关键？
• RQ4 基于元损失方差的辅助损失在多大程度上提升了模型的鲁棒性和性能？
• RQ5 在 IRMv1 因其严格线性假设而失效的非线性设置中，Meta-IRM 是否能实现有效的泛化？

主要发现

• Meta-IRM 在 OOD 泛化方面显著优于 IRMv1 及所有 IRM 变体，在测试环境 p_e = 0.9 且 p_e2 = 0.9 时，测试准确率达到 54.5%，而 IRMv1 仅为 13.4%。
• 在带有两个虚假特征（颜色和补丁）的 Colored MNIST 任务中，Meta-IRM 在所有测试环境中均保持一致的性能（54.5% ± 4.0），而 IRMv1 在最分布外的设置中崩溃至 13.4%。
• 消融研究显示，若移除标准差损失，OOD 性能下降（64.2% vs. 70.4% for meta-IRM），证实其在稳定不变性学习中的关键作用。
• 若使用一阶近似进行元更新，性能下降（59.1% vs. 70.4%），表明二阶梯度对准确双层优化至关重要。
• 若在内层和外层循环中使用相同的环境，OOD 泛化性能极差（准确率 13.6%），证明不相交环境采样对实现不变性至关重要。
• 即使真实因果特征存在噪声，Meta-IRM 仍能实现稳定且不变的预测，展现出超越原始 IRM 目标局限性的鲁棒性。