Entropy Minimization without Model Collapse: Mitigating Prediction Bias in Medical Imaging

技術要約：モデル崩壊を伴わないエントロピー最小化

問題提起
テスト時適応（Test-Time Adaptation; TTA）は、センサー、環境、または取得プロトコルの変化によるデータ分布のシフトが発生する、現実世界のシナリオにおいて機械学習モデルをデプロイするために不可欠である。医療画像において、これらのシフトはしばしば微細かつ不均一（例：異なるスキャナーや染色プロトコル）であるが、重大な性能低下を引き起こす可能性がある。エントロピー最小化（Entropy Minimization; EM）は、ラベルなしのテストデータに対する予測の不確実性を減少させることを目的として、TTAにおける支配的な目的関数となっている。しかし、EMは、モデルが入力に関わらず単一のクラスを予測するように収束してしまう「モデル崩壊（model collapse）」と呼ばれる決定的な失敗モードに直面する。

著者らは、この崩壊の根本原因は適応プロセスそのものではなく、彼らが「予測バイアス（prediction bias）」と呼ぶ現象にあることを特定した。分布シフトの下では、表現空間における異なるクラスに対応する特徴クラスターがドリフトして部分的に融合する場合がある一方で、分類器の決定境界は固定されたままである。これが、特定のクラスが過剰に表現され、他のクラスが抑制されるという、予測クラス分布の歪みを誘発する。標準的なEMは、この既存のバイアスを強め、融合したクラスターをさらにタイトにすることで、モデルを自明な単一クラスの解へと導く正のフィードバックループを作り出す。

手法：分布シフト・バイアス削減（DSBR）
この失敗モードを軽減するために、著者らは、ソースデータやラベルにアクセスすることなく、テスト時のみ動作するように設計されたバイアス補正目的関数である「分布シフト・バイアス削減（Distribution Shift Bias Reduction; DSBR）」を提案する。

1. コアメカニズム: DSBRは、各予測クラスのアンセントロピー最小化損失への寄与を等しくすることで、予測バイアスに対処する。支配的なクラスターが勾配更新を決定することを許すのではなく、DSBRは損失を再重み付けし、単一のクラスが最適化信号を圧倒することを防ぐ。
2. 数学的定式化: 本手法は、各予測クラスターに等しい確率質量（$1/K$）を割り当てるバランスの取れたターゲット分布 $p^*_X$ を定義する。重要度再重み付け（importance reweighting）を用いて、サンプル $x$ に対する損失は $1 / (K \cdot \hat{p}_{\hat{Y}}(\hat{y}(x)))$ によってスケールされる。ここで $\hat{p}_{\hat{Y}}$ は予測クラスの推定頻度である。これにより、過小表現されているクラスターのサンプルは勾配への寄与が大きくなり、過剰表現されているクラスターのサンプルは寄与が小さくなるため、融合したクラスターの分離を積極的に促進する。
3. 実装: 真の予測クラス分布はテスト時には未知であるため、DSBRはバッチを横断したクラスター頻度の指数移動平均（EMA）を用いてこれを近似する。これは、バッチサイズが小さい場合に退化したバッチごとの推定値が生じる可能性がある、高クラス数データセット（例：$K=1000$ の ImageNet）において特に重要である。

主な貢献

• 予測バイアスの特定: 本論文は、予測バイアス（適応が始まる前に特徴クラスターの融合によって存在する予測クラス分布の歪み）が、EMにおけるモデル崩壊の根本原因であることを確立している。EMは、このバイアスを作り出すのではなく、崩壊が起こるまでそれを増幅させることを示している。
• DSBRアルゴリズム: 著者らは、エントロピー損失を再重み付けしてクラスの寄与をバランスさせることで、この失敗モードを特異的に標的とするDSBRを導入する。他のアプローチとは異なり、DSDRはサンプルフィルタリング、ドメイン固有のデータ拡張、または複雑な正則化に依存しない。
• 包括的な評価: 研究では、4つの医療画像データセット（WILDSCamelyon, Histopantum, MammoBench, GliomaMRI）および自然画像ベンチマークであるImageNet-Cにわたる困難な適応シナリオを設計している。

結果

• 崩壊の防止: 医療画像ベンチマークにおいて、標準的なEM手法（Tentなど）は、特定のソース・ターゲット・ドメインのペアに対して、しばしばチャンスレベルの精度へと崩壊する。DSBRはこれを一貫して防ぎ、すべての構成においてほぼ最適な精度を達成する。
• 性能向上: DSBRは、最先端のベースライン（Tent, SAR, ROID, DeYOを含む）を上回る性能を示した。医療データセットにおいて、非適応ベースラインと比較して、ResNetで平均+4.2%、ViTで+2.6%の改善を達成し、次点の優れた手法を大幅に上回った。
• 堅牢性: ImageNet-Cにおいて、DSBRは深刻な破損、混合された破損、および極端なラベルシフトの下でも優れた安定性を示した。EMAの減衰を適切に調整すれば、バッチサイズが1という極小の状態でも効果的に対処できたが、他の手法はしばしば失敗するか崩壊した。
• メカニズムの検証: 特徴空間の埋め込みの可視化により、EMがクラスターの融合と崩壊を引き起こす一方で、DSBRは適応プロセスを通じて十分に分離されたクラスターを維持することが確認された。

意義
本論文は、エントロピー最小化を安定させるには、それが「なぜ」失敗するのかという原理的かつメカニズム的な理解が必要であると主張している。予測バイアスを崩壊の主要な要因として分離することで、著者らは、この単一の量を補正することがTTAを安定させるのに十分であることを示している。本研究の意義は、医療画像のような重要な領域における堅牢なテスト時適応には、複雑なヒューリスティックやデータフィルタリングではなく、モデル崩壊を招く最適化ダイナミクスの直接的な補正が必要であることを示した点にある。DSBRは、テスト時のみに動作するシンプルでハイパーパラメータの少ないソリューションを提供しており、信頼性が極めて重要となる臨床現場への適用において非常に有用である。