Robust Conditional Conformal Prediction via Branched Normalizing Flow

技術的サマリー：分岐型正規化フローによるロバストな条件付き適合予測

1. 問題定義

適合予測（CP）は、較正データとテストデータが独立同分布（i.i.d.）であるという仮定の下で、周辺カバレッジ保証を備えた予測集合を提供する。しかし、分布シフト（テスト分布 $Q_{XY}$ が較正分布 $P_{XY}$ と異なる場合）が存在する際、既存の手法は通常、周辺適合スコア分布のみを整合させる。これにより周辺カバレッジは維持されるものの、個々のテスト入力における条件付きカバレッジ誤差の制御には失敗する。その結果、条件付きスコア分布が不一致となる特定の領域では CP 集合が信頼できなくなり、特定の入力に対して過小カバレッジまたは過剰な過大カバレッジが生じる可能性がある。核心的な課題は、推論時に真のテストラベルを必要とすることなく、結合分布シフト下においてロバストな条件付きカバレッジ保証（$Pr(Y \in C(X) | X=x) \ge 1-\alpha$）を確保することである。

2. 手法

提案される解決策は、分布シフトと条件付きカバレッジの劣化を結びつける理論的枠組みと、それを緩和するための新規アーキテクチャを導入する。

理論的枠組み:

• 条件付きカバレッジギャップ（CCG）: 著者は、$Q_{XY}$ 下でのテスト入力の条件付きカバレッジと、$P_{XY}$ 下での保証との間の不一致を測定する指標として CCG を定義する。
• 統合カバレッジギャップ（ICG）: テスト特徴量分布 $Q_X$ 全体にわたる CCG の期待値として定義され、包括的なロバスト性指標として機能する。
• ワッサーシュタイン上限: 本論文は、結合較正分布とテスト分布（$P_{XY}$ および $Q_{XY}$）間のワッサーシュタイン距離（$W$）を用いて、ICG の上限を導出した。具体的には、$ICG(P, Q) \le \sqrt{2\kappa L} (\eta \cdot W(P_{XY}, Q_{XY}) + 1/4)$ である。この上限は、分布間のワッサーシュタイン距離を最小化することが、理論的にロバストな条件付きカバレッジを促進することを示唆している。

アルゴリズム：分岐型正規化フロー（BNF）
ワッサーシュタイン上限を実用化するため、著者はテスト分布 $Q_{XY}$ を較正分布 $P_{XY}$ に一致させるように設計された、可逆写像 $f_\theta(x, y) = (f_{\theta_X}(x), f_{\theta_Y}(y))$ である**分岐型正規化フロー（BNF）**を提案する。

• アーキテクチャ: BNF は 2 つの独立した分岐を利用する。$X$ 分岐（$f_{\theta_X}$）はテスト入力を正規化し、$Y$ 分岐（$f_{\theta_Y}$）はラベルを変換する。重要なのは、分岐間でパラメータを共有しないことであり、これにより $X$ の変換が観測されないテストラベル $Y$ に依存しないことを保証している。
• 推論プロセス:
  1. テスト入力 $X_{n+1}$ を正規化入力 $\bar{X}_{n+1} = f_{\theta_X}(X_{n+1})$ に変換する。
  2. 較正分布 $P_{XY}$ を用いて（例えば、適合化量子回帰を通じて）条件付き予測集合 $CC(\bar{X}_{n+1})$ を構築する。
  3. 集合を $f_{\theta_Y}^{-1}$ を用いてテスト空間へ逆変換し、最終的なロバスト予測集合 $C_{BNF}(X_{n+1})$ を取得する。
• 拡張 BNF: 単純な逆写像に必要な単変数単調変換では不十分な複雑な分布を扱うため、著者は拡張 BNFを導入する。この変種は、$Y$ 分岐変換（$f_{\theta_Y}^{aug}(y; \epsilon)$）に確率的ノイズ成分 $\epsilon \sim \mathcal{N}(0,1)$ を追加する。これにより表現力が向上するが、変換が厳密に単調ではなくなるため、推論時に予測集合を構築する際に根探索問題（またはグリッドサーチ）を解く必要がある。
• マルチソースドメイン汎化（MSDG）: テスト分布が $K$ 個のソース分布の混合であるシナリオにおいて、この手法は均一な混合重みを仮定し、較正セットと変換されたソースセット間の平均ワッサーシュタイン距離を最小化する。

3. 主要な貢献

1. ロバスト性の定量化: 条件付きカバレッジギャップ（CCG）と統合カバレッジギャップ（ICG）の定義により、分布シフト下における条件付きカバレッジのロバスト性を評価するための理論的指標を提供する。
2. 理論的上限: 結合分布間のワッサーシュタイン距離を用いた ICG の上限導出により、輸送最小化と条件付きカバレッジの維持との間の理論的リンクを確立する。
3. 分岐型正規化フロー（BNF）: テストラベルを必要とせずにテスト入力を較正類似分布に変換することを可能にする、2 分岐可逆アーキテクチャの導入。これにより、逆変換時に条件付き保証を維持する。
4. 拡張 BNF: 複雑な分布への適合能力を向上させるための確率的拡張を用いた拡張版であり、単調性の喪失にもかかわらず条件付き保証を回復する方法を併せ持つ。

4. 実験結果

この手法は、マルチソースドメイン汎化（MSDG）設定下において、合成シフトと現実世界の課題（売上予測、交通量予測、医療意思決定、感染症モデリング）を網羅する 9 つのデータセットで評価された。

• 指標: 性能は、周辺カバレッジ、最悪スライスカバレッジ（WSC）、および提案された過小および過大カバレッジの両方を罰する最悪スライスカバレッジギャップ（WSCG）を用いて測定された。
• 知見:
  • ロバスト性: 適合化量子回帰（CQR）と組み合わせた拡張 BNF は、すべてのデータセットおよび信頼水準において一貫して最低の WSCG を達成し、Split CP、重要度重み付き CP、最悪ケース CP、ワッサーシュタイン正則化 CP などのベースラインを上回った。
  • カバレッジ: この手法は、周辺カバレッジと WSC を目標信頼水準（$1-\alpha = 0.9$）に近づけながら、条件付きスライスにおける目標からの偏差を大幅に減少させた。
  • 効率性: 標準的な拡張 BNF はソースの不確実性を考慮するためにより大きな予測集合を生成するが、「Augment-Conditioned」変種（推論時に既知のソースに条件付ける）は、カバレッジを維持しつつ集合サイズを大幅に削減した。
  • 汎化性: この手法は、異なるサンプルサイズ、ソースドメイン数（$K$）、および異なる信頼水準においてロバスト性を示した。

5. 意義と主張

本論文は、その主な意義が、既存の周辺焦点手法ではほとんど扱われてこなかった、分布シフト下における適合予測の条件付きカバレッジの劣化への対応にあると主張している。ワッサーシュタイン距離を通じて条件付き無効性を境界付けし、可逆な分岐フローを実装することで、この研究は、データ分布がシフトしても個々のテスト入力に対する不確実性定量化が信頼できることを保証する、理論的根拠に基づき実証的に効果的なメカニズムを提供する。著者はこれを、分布シフトが避けられず条件付き信頼性が極めて重要なクリティカルな領域（例：医療、交通、感染症モデリング）における、より安全な AI 応用への一歩として位置づけている。

本論文は、拡張 BNF における根探索の計算コスト、予測集合の確率的性質、および多変数入力を必要とするアーキテクチャの制約といった限界を認めている。また、フローの訓練に較正データを再利用すること（潜在的なチューニングバイアス）について言及しつつ、アブレーション研究により、実際の影響は最小限であることが示唆されている。