[论文解读] Maximum Mean Discrepancy for Generalization in the Presence of Distribution and Missingness Shift

Liwen Ouyang, Aaron Key|arXiv (Cornell University)|Nov 19, 2021

Domain Adaptation and Few-Shot Learning参考文献 30被引用 6

一句话总结

本文提出使用最大均值差异（MMD）通过在输入空间、表征空间或两者中对齐训练和测试数据分布，以改善在协变量偏移和缺失偏移下的模型泛化能力。MMD 表征、MMD 掩码和 MMD 混合方法在合成数据和真实世界数据上显著提升了模型性能、校准能力和对特征缺失的鲁棒性，尤其在测试数据中缺失率较高的情况下表现更优。

ABSTRACT

Covariate shifts are a common problem in predictive modeling on real-world problems. This paper proposes addressing the covariate shift problem by minimizing Maximum Mean Discrepancy (MMD) statistics between the training and test sets in either feature input space, feature representation space, or both. We designed three techniques that we call MMD Representation, MMD Mask, and MMD Hybrid to deal with the scenarios where only a distribution shift exists, only a missingness shift exists, or both types of shift exist, respectively. We find that integrating an MMD loss component helps models use the best features for generalization and avoid dangerous extrapolation as much as possible for each test sample. Models treated with this MMD approach show better performance, calibration, and extrapolation on the test set.

研究动机与目标

• 解决机器学习模型中训练和测试数据分布不同的协变量偏移和缺失偏移问题。
• 通过最小化输入空间和表征空间中训练和测试数据分布之间的 MMD 来改善模型泛化能力。
• 开发一种新型的 MMD 掩码模型，以对齐训练集和测试集中缺失模式，尤其适用于测试数据中特征缺失率较高的情况。
• 将基于 MMD 的分布对齐与缺失模式对齐结合，形成混合模型以实现更优性能。
• 证明基于 MMD 的训练可减少危险外推，并提升分布外测试集的校准性和鲁棒性。

提出的方法

• 提出 MMD 表征：使用混合核函数最小化训练和测试特征表征之间的 MMD，以在嵌入空间中对齐分布。
• 引入 MMD 掩码：使用 MMD 训练掩码模型，以对齐训练集和测试集中输入特征与缺失指示符的联合分布，特别针对缺失偏移问题。
• 开发 MMD 混合：结合输入空间和表征空间中的 MMD 匹配，以联合解决分布偏移和缺失偏移问题。
• 采用神经网络框架，端到端优化 MMD 损失与主预测损失，实现分布和缺失模式的联合对齐。
• 采用两阶段训练流程：首先在带掩码特征的标注数据上训练 MMD 掩码模型，然后在 MMD 掩码模型生成的掩码数据上微调下游模型。
• 应用混合核 MMD 以捕捉复杂的分布差异，并提升对齐的鲁棒性。

实验结果

研究问题

• RQ1 在特征表征空间中基于 MMD 的分布对齐是否能改善协变量偏移下的模型泛化能力？
• RQ2 专用的 MMD 掩码模型是否能有效对齐训练集和测试集之间的缺失模式，以缓解缺失偏移？
• RQ3 在输入空间和表征空间中同时进行 MMD 对齐是否优于单独使用任一空间？
• RQ4 与简单的随机或启发式掩码策略相比，MMD 掩码模型在模型鲁棒性和性能方面表现如何？
• RQ5 基于 MMD 的训练在多大程度上能减少危险外推，并提升分布外测试数据的校准性？

主要发现

• 在掩码传递测试集上，MMD 模型相较于简单模型在碳排放估计任务中表现出显著的性能提升。
• 在由 MMD 和简单掩码共同掩码的数据上，MMD 模型始终优于简单模型，且使用 MMD 掩码时性能增益更大。
• 当整组特征缺失（如 ESG 数据）时，MMD 模型的性能接近未掩码数据，而简单模型性能显著下降。
• MMD 掩码模型捕捉到的缺失模式比简单掩码更接近真实世界缺失模式，这在测试集中与真实缺失模式的对齐程度上得到验证。
• MMD 模型展现出更优的校准性和鲁棒性，尤其在高缺失率场景下，表明其减少了危险外推。
• 结合输入空间和表征空间对齐的混合 MMD 方法在所有测试场景下均取得最佳整体性能。