OTLesMix: Wasserstein Barycenter and Optimal Transport Map for Synthetic Lesion Generation with Diverse Shapes and Locations

OTLesMixの技術要約：合成病変生成のためのワッサースタイン・バリセントと最適輸送写像

問題提起
ディープラーニングは医療画像セグメンテーションに革命をもたらしたが、教師あり学習は、特に高密度な3Dセグメンテーションタスクにおいて、専門家によるアノテーションに要する高いコストと時間に制約されている。データ拡張はモデル訓練を改善するための標準的な手法であるが、従来のメソッド（空間的／強度の変換）や高度な生成モデル（GANや拡散モデル）には特定の限界が存在する。生成モデルは、それ自体の学習のために大規模かつ多様なデータセットを必要とするが、医療の文脈ではそのようなデータは利用できない場合が多い。対照的に、既存の「ミックスベース」の拡張戦略（例：MixUp、CutMix、CarveMix、LesionMix）は、病変の形状や位置に十分な多様性を生成できないことが多い。これらは通常、既存のマスクの単純なランダムサンプリングや硬直した幾何学的変換に依存しており、その結果、セグメンテーションモデルを堅牢に訓練するために必要な多様性に欠ける合成サンプルを生み出してしまう。

手法：OTLesMix
著者らは、最適輸送（OT）理論を活用して、現実的かつ多様な合成病変を生成する新しい画像合成手法であるOTLesMixを提案している。この手法は、ベースとなるデータセット $D = \{(X_i, Y_i)\}$ から新しいサンプル $(\tilde{X}, \tilde{Y})$ を作成するために、主に2つの段階で動作する：

1. ワッサースタイン・バリセントによる病変マスク生成：

   • 2つのベースサンプル $(X_i, Y_i)$ と $(X_j, Y_j)$ がランダムに選択される。
   • グラウンドトゥルース・マスク $Y_i$ および $Y_j$ における非ゼロピクセルの座標は、確率分布として扱われる。
   • 新しい病変マスク $\tilde{Y}_{les}$ は、これら2つの分布の**ワッサースタイン・バリセント（Wasserstein barycenter）**を計算することによって生成される。これには、ソースマスクへの輸送コストの加重和を最小化するように、新しいマスク点の座標を最適化するプロセスが含まれる：
     $$B(\lambda, \alpha_i, \alpha_j, \beta) = \lambda OT(\alpha_i, \beta) + (1-\lambda)OT(\alpha_j, \beta)$$
     ここで、$\lambda \sim U(0,1)$ はランダムな重み付けパラメータである。
   • このプロセスは、2つの元のマスクの間で形状と位置を補間する。これにより、新しい病変は元の位置間のどこにでも存在することができ、その形状は元の形状の滑らかな補間となる。
   • 計算効率を確保するため、著者らは**デバイアス・シンクホーン・ダイバージェンス（de-biased Sinkhorn divergence）**とシンクホーン・アルゴリズムを利用しており、現代的なGPU上で数秒以内に最大$10^6$個の点を処理することを可能にしている。

2. 最適輸送写像による画像およびラベル生成：

   • 新しいマスク $\tilde{Y}_{les}$ が定義された後、合成された病変のピクセル強度とラベルを決定する必要がある。
   • 生成されたマスク $\tilde{Y}_{les}$ と、ソースマスクの1つ（例：$Y_j$）との間の**最適輸送計画（optimal transport plan）**が計算される。
   • 新しいマスク内の各ピクセルについて、輸送写像 $T$ が、ソース画像 $X_j$ における対応する位置を決定する。
   • 強度合成（Intensity Synthesis）： 合成されたピクセルの強度は、マップされた位置におけるソース画像の強度の線形補間によって得られる。
   • ラベル合成（Label Synthesis）： マルチクラス・タスクの場合、グラウンドトゥルースのラベルは、ソース・ラベルマップの最近傍補間によって割り当てられる。
   • 合成された病変は、ベース画像（$X_i$、$X_j$、またはデータセット内の第3の画像）にインペイントされる。ポストプロセッシング工程により、既存の病変との交差や、非現実的な解剖学的部位（例：頭蓋骨、脳脊髄液）との重なりが除去される。

主な貢献

• 新規な合成フレームワーク： ワッサースタイン・バリセントによる形状／位置の補間と、最適輸送写像による強度／ラベルの転送を独自に組み合わせたOTLesMixを導入した。
• 強化された多様性： 本手法は、既存のミックスベースの手法が持つ限界を明示的に解決し、既存データの単なるコピーや単純な回転ではない、多様な形状と位置を持つ病変を生成する。
• 最先端の性能： 本手法は、標準的なミックスベースのデータ拡張技術と比較して、ダウンストリームのセグメンテーション・タスクの性能を向上させる上で優れた性能を示すことが実証された。

実験結果
著者らは、nnUNetバックボーンを用いて、3つの脳病変セグメンテーション・データセットでOTLesMixを評価した：

1. BraTS 2020： 膠芽腫および低グレード神経膠腫のセグメンテーション（2Dスライス）。
2. ATLAS v2.0： 脳卒中ニューロイメージング（3Dボリューム）。
3. ISLES 2022： 虚血性脳卒中病変のセグメンテーション（3Dボリューム）。

定量的知見：

• OTLesMixは、すべてのデータセットにおいて、最先端のミックスベース手法（MixUp、CutMix、CarveMix）およびベースライン（従来のデータ拡張のみを用いたnnUNet）を上回った。
• ベースラインに対する改善： 合成データなしで訓練されたモデルと比較して、Dice Similarity Coefficient (DSC) が 2.9から6.6ポイント向上した。
• 競合手法に対する改善： OTLesMixは、2番目に優れた手法に対して 0.8から1.8ポイントのDSC向上を達成した。
• 統計的有意性： 改善は他の生成手法と比較して統計的に有意であった（$p \le 0.05$）。最も大きなp値は、ATLASにおけるCarveMixとの比較で0.012であった。

定性的知見：
視覚的検査により、OTLesMixは他の手法と比較して、より多様な形状と位置を持つ病変を生成することが確認された。本手法は、マルチクラス・セグメンテーション（BraTS）や複数の連結成分を持つ病変（ATLAS、ISLES）をうまく扱い、補間された形状の一貫性を維持した。

意義と限界
論文は、OTLesMixが、データの乏しい医療画像におけるデータ拡張の堅牢なソリューションを提供し、単純なミックス手法と、完全な生成モデルを訓練することの複雑さとの間の「中間領域」を提供すると主張している。最適輸送を活用することで、合成サンプルが基礎となるデータの幾何学的および強度的な分布と整合していることを保証している。

しかし、著者らは以下の特定の限界も認めている：

• 計算コスト： 1つのサンプル生成に、2D画像では約1秒、3Dボリュームでは解像度に応じて1〜3分を要する。
• リアリズムの制約： 本手法はピクセル強度の線形補間に依存している。そのため、混合に使用される2つのソース画像間に大きなドメインギャップ（例：異なるスキャナーや撮像プロトコル）がある場合、不自然なインペイントが生じる可能性がある。
• 今後の方向性： 著者らは、合成された強度のリアリズムをさらに向上させるために、OTLesMixのマスク生成手順を、マスク条件付きのディープラーニング生成モデルと組み合わせることを今後の課題として挙げている。