Re-mixing Embeddings for Patient Augmentation in Data Scarce Multiple Instance Learning

技術要約：データ不足のマルチインスタンス学習における患者増強のための埋め込み再混合（Re-mixing Embeddings）

問題提起
データの不足は、希少疾患、特殊で高価なモダリティ、あるいは侵襲的なデータ収集が制限されているシナリオにおいて、医学的マルチインスタンス学習（MIL）における決定的なボトルネックであり続けている。MILにおいて、「バッグ」（例：患者）は複数の「インスタンス」（例：画像パッチ、単一細胞、またはゲノム特徴量）で構成され、モデルはインスタンスレベルの注釈なしにバッグレベルのラベルを予測する。従来の拡張技術（回転、反転など）はインスタンスレベルで動作するため、バッグレベルでの現実的なバリエーションを生成するには不十分である。さらに、既存のバッグレベルの混合戦略（PseMix、ReMixなど）には、主に3つの制限がある：

1. これらはオンラインでのみ動作し、検証や再利用ができない一時的な合成バッグを作成する。
2. その局所的な混合範囲は、グローバルな集団統計を捉えることができない。
3. すべてのクラスから実例を必要とするため、あるクラス全体（例：健康なコントロール群）がデータセットに欠落している場合には効果を発揮できない。

手法：RECIPE
著者らは、埋め込み空間において新しい患者バッグを生成するための、オフラインかつ統計的に根拠のあるフレームワークであるRECIPE（Clusteringによる患者増強のための埋め込み再混合：Re-mixing Embeddings via Clustering for Patient augmEntation）を提案している。パイプラインは主に2つのステージで構成される：

1. クラスタリングに基づく教師なしサンプリング（レシピ学習）：

   • 埋め込みプーリング（Embedding Pooling）： すべての訓練患者からのインスタンス埋め込みを一つに集約（プーリング）する。
   • ガウス混合モデル（GMM）： プーリングされた埋め込みにGMMを適合させ、教師なしクラスタを定義する。
   • レシピの導出（Recipe Derivation）： 各疾患クラスについて、学習された各クラスタへのインスタンスの分布を示す「レシピ」を計算する。これは、クラス $c$ における各クラスタ $k$ に割り当てられるインスタンス数の平均（$\mu_{c,k}$）と標準偏差（$\sigma_{c,k}$）によって定義される。
   • 患者の生成： 学習されたレシピに基づいて各クラスタのインスタンス数（$n_{c,k} \sim \mathcal{N}(\mu_{c,k}, \sigma^2_{c,k})$）をサンプリングし、その後、それらのクラスタに割り当てられた実インスタンスのプールからその数だけ埋め込みを抽出することで、新しい患者を生成する。これにより、外部のデータセットからレシピを導出できる限り、特定のクラスの実例が存在しない場合でも患者を生成することが可能になる。

2. 不確実性を考慮した患者選択：

   • 生成されたデータがモデルの性能を向上させるものであり、ノイズを導入するものではないことを保証するために、本手法は不確実性定量化を採用している。
   • 推論時には、事前学習済みMILモデルを用いてモンテカルロ（MC）ドロップアウトによる評価を行う。
   • 不確実性を定量化するために予測エントロピーを計算する（$U(x) = H[\bar{p}(x)]$）。
   • 生成されたバッグの中で、最も不確実性が高い（エントロピーが高い）患者のみを選択し、訓練セットを増強する。これは、これらのサンプルが最も情報量の多い学習信号を提供するという仮説に基づいている。

主な貢献

• オフラインかつ統計的に根拠のある生成： オンライン混合法とは異転なり、RECIPEはグローバルな分布的「レシピ」を学習することで、永続的で分析可能な合成コホートを作成する。
• 欠落クラスへの対応： RECIPEは、患者の機密情報を共有することなく、外部の統計量を用いて、完全に欠落している患者クラス（例：健康なコントロール群）を生成できる初のMILフレームワークである。
• クロスモダリティへの適用可能性： 本手法は、組織病理画像、単一細胞RNAシーケンシング（scRNA-seq）、フローサイトメトリーを含む、多様なデータモダリティにおいて有効であることが示されている。
• アーキテクチャに依存しない設計： 本フレームワークは、AttentionベースのMIL（ABMIL）、DSMIL、Transformerなどの様々なMILアグリゲーターと互換性がある。

実験結果
著者らは、臨床的に関連のある3つのデータ不足シナリオにおいてRECIPEを評価した：

1. クロスデータセット転移（欠落クラス）：

   • シナリオ： AML-Hehrデータセットからすべての健康なコントロール群を除去した。外部のcAItomorphデータセットから「健康なレシピ」を導出し、それをAML-Hehrの埋め込みプールに適用した。
   • 結果： 生成された健康な患者は、実際の細胞タイプ組成と密接に一致した（イェンセン・シャノンダイバージェンス = 0.12）。これらの生成された患者を用いて訓練を行った結果、フルデータセットでの訓練に近い性能（バランス精度 0.78 vs 0.83）を達成し、外部データの直接転移（0.66）を大幅に上回った。

2. 低データ領域：

   • シナリオ： AML-Hehrにおける健康なコントロールの数をわずか $n=1$ まで減少させた。
   • 結果： 実在する健康なコントロールが1つしかない場合、ベースラインのバランス精度は0.26であった。RECIPEは0.68を達成し、MixUp（0.60）およびPseMix（0.64）を大幅に上回った。

3. 非画像モダリティ（scRNA-seqおよびフローサイトメトリー）：

   • シナリオ： 自然にサンプルサイズが限定されているCOVID-19重症度データセットにおける増強。
   • 結果： RECIPEは、scRNA-seqデータにおいて、PCA埋め込みで+14%、scVI埋め込みで+26%のバランス精度向上を実現し、フローサイトメトリーデータでは+22%の向上を実現し、すべてのベースラインの拡張戦略を上回った。

意義と主張
本論文は、学習された統計分布に基づいてプーリングされた埋め込みを再混合することにより、現実的な患者コホートの生成を可能にし、RECIPEが医学的MILにおけるデータ不足のボトルネックに対する新しい解決策を提供すると主張している。その主な意義は、以下の点にある：

• 「欠落クラス」の問題に対処し、患者のプライバシーを損なうことなく、外部の統計量を用いて健康なコントロールや希少疾患の表現型を合成できること。
• 従来の拡張技術が失敗する極端な低データ領域において、MILの性能を向上させること。
• 単一細胞やサイトメトリーのデータを超えて適用可能な、汎用性の高いフレームワークを提供すること。

著者らは、彼らのアプローチが、局所的なオンライン混合を超えて、グローバルな集団統計をモデル化し、希少疾患研究やプライバシー保護型のデータ増強のギャップを埋めることができる、検証可能で再利用可能な合成データを作成できることを強調している。