[论文解读] UMIX: Improving Importance Weighting for Subpopulation Shift via Uncertainty-Aware Mixup

Zongbo Han, Zhipeng Liang+8 more|arXiv (Cornell University)|Sep 19, 2022

Machine Learning and Data Classification被引用 10

一句话总结

本文提出UMix，一种新颖的不确定性感知混合方法框架，通过基于样本不确定性的重加权混合样本，改进了子群体偏移下的重要性加权。通过将不确定性估计整合到混合方法中，UMix缓解了过参数化模型中的过拟合问题，并实现了比加权ERM更紧的泛化界，优于多种基准测试中的先前方法。

ABSTRACT

Subpopulation shift widely exists in many real-world machine learning applications, referring to the training and test distributions containing the same subpopulation groups but varying in subpopulation frequencies. Importance reweighting is a normal way to handle the subpopulation shift issue by imposing constant or adaptive sampling weights on each sample in the training dataset. However, some recent studies have recognized that most of these approaches fail to improve the performance over empirical risk minimization especially when applied to over-parameterized neural networks. In this work, we propose a simple yet practical framework, called uncertainty-aware mixup (UMIX), to mitigate the overfitting issue in over-parameterized models by reweighting the ''mixed'' samples according to the sample uncertainty. The training-trajectories-based uncertainty estimation is equipped in the proposed UMIX for each sample to flexibly characterize the subpopulation distribution. We also provide insightful theoretical analysis to verify that UMIX achieves better generalization bounds over prior works. Further, we conduct extensive empirical studies across a wide range of tasks to validate the effectiveness of our method both qualitatively and quantitatively. Code is available at https://github.com/TencentAILabHealthcare/UMIX.

研究动机与目标

• 解决重要性加权（IW）方法在子群体偏移下无法改善经验风险最小化（ERM）的问题，尤其是在过参数化神经网络中。
• 通过在混合增强过程中引入不确定性感知的样本重加权，缓解基于IW方法的过拟合。
• 提供理论依据，表明UMix的泛化界比加权ERM更紧，尤其在数据内在维度较低时。
• 开发一种实用且高效的框架，在群体无关和群体相关评估设置下均提升性能。

提出的方法

• 提出不确定性感知混合（UMix），通过标准混合方法生成混合样本，并基于小批量中原始样本的不确定性对混合样本进行重加权。
• 使用基于训练轨迹的不确定性估计方法量化每个样本的不确定性，从而实现对高不确定性、少数子群体样本的动态关注。
• 制定加权混合损失，结合每个混合操作中所用两个原始样本的加权损失，权重来源于子群体频率和不确定性。
• 理论分析表明，UMix的假设空间导致依赖于子群体异质性的泛化界，当 rank(Σ_X) ≪ d 时，其泛化界比加权ERM更紧。
• 将不确定性感知重加权集成到标准SGD训练循环中，保持与现有深度学习流水线的兼容性。
• 利用ρ-保持性数据分布的假设，推导Rademacher复杂度界，支持理论泛化保证。

实验结果

研究问题

• RQ1 不确定性感知混合方法能否提升过参数化模型在子群体偏移下的重要性加权鲁棒性？
• RQ2 基于不确定性的混合样本重加权是否能实现比标准重要性加权或单独使用混合方法更好的泛化性能？
• RQ3 UMix能否在低内在维度数据环境下实现比加权ERM更紧的泛化界？
• RQ4 UMix在多样化任务中，包括群体无关和群体相关评估设置下的表现如何？

主要发现

• UMix在多种基准测试中，无论在群体无关还是群体相关设置下，均持续优于先前的重要性加权方法，包括CVaR-DRO和JTT。
• 理论分析证实，UMix的泛化界比加权ERM更紧，其改进程度取决于子群体异质性和数据内在维度。
• 泛化界中的红色项（与数据异质性相关）在 rank(Σ_X) ≪ d 时变小，表明在低维子群体结构中具有更强的理论保证。
• 实证结果表明，UMix能有效缓解过参数化模型中的过拟合，防止标准IW方法中常见的性能下降。
• 基于训练轨迹的不确定性估计成功识别出高风险样本，特别是少数子群体，从而提升模型的公平性和鲁棒性。
• UMix在无需架构修改的情况下展现出强大的实证性能，使其成为现有子群体偏移缓解流水线的即插即用增强方案。