[論文レビュー] 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揃えを統合したハイブリッドアプローチが、すべてのテストシナリオにおいて最良の全体的性能を達成した。