Coupled Optimal Transport with Landmark Constraints

技術要約：ランドマーク制約付き結合最適輸送

問題提起
既存の最適輸送（OT）モデルは、主に2つの確率分布間の輸送コストまたは歪みを最小化する輸送写像または輸送計画を探索する。分布のマッチングには効果的であるが、コストのみを最小化することは、幾何学的に意味のある変換を特定できない場合が多い。例えば、画像や形状のレジストレーションにおいて、コスト駆動型の写像は、幾何学的には短い経路であっても、意味論的には誤った（例：異なる解剖学的部位を一致させてしまう）経路で質量を移動させることがある。逆に、疎なアノテーションによるランドマークのみに依存する方法は、グローバルな分布の遷移や、変化を支配する基礎となる変形場を捉えきれない可能性がある。本論文は、グローバルな分布の整合性と、疎なランドマークによって提供される局所的な幾何学的制約の両方を満たす、一貫した変形場の回復という課題に取り組む。

手法
著者らは、輸送計画の最適化と変形場の推定を統合する、新しい**結合最適輸送（Coupled Optimal Transport: COT）**フレームワークを提案する。

• モデル定式化: フレームワークは、輸送の変形ベースの視点を導入する。$u: \Omega \to \mathbb{R}^p$ を変形場とし、$T_u = \text{Id} + u$ を関連する輸送写像とする。アノテーションされたランドマークのペアの集合を $S = \{(x_i, y_i)\}_{i=1}^m$ とするとき、モデルは結合 $\Pi(\mu, \nu)$ と変形場のソボレフ空間 $U$ の集合上で汎関数 $F(\pi, u)$ を最小化する：
  $$\inf_{\pi \in \Pi(\mu, \nu), u \in U} F(\pi, u)$$
  この汎関数は、主に以下の3つの要素で構成される：

  1. 結合輸送コスト: 古典的な輸送コスト $c(x, y)$ と変形一貫性ペナルティ $h(T_u(x) - y)$ の加重結合。この項は、輸送計画 $\pi$ が変形場 $u$ と整合していることを保証する。
  2. ランドマークによる教師あり学習: $\sum w_i \rho(A_i^r(u) - \Delta_i^{lm})$ という項であり、観測オペレータ $A_i^r$ を介した、観測されたランドマークの変位と変形場の局所平均との間の不一致を罰する。
  3. 正則化: 空間的な一貫性と滑らかさを促進するための弾性エネルギー項 $E_{el}(u)$（線形弾性エネルギー）。これにより、ドメイン全体への疎なランドマーク情報の伝播を容易にする。

• 理論的特性:

  • 良定常性: 著者らは、（例：$\mu$ の絶対連続性、弾性エネルギーの強コエルシビティなどの）緩やかな条件下で、最小化子 $(\pi^\star, u^\star)$ の存在を確立している。
  • 極限挙動: バランシング・パラメータ $\alpha \to 1$（変形支配的）および $\alpha \to 0$（コスト支配的）の場合のモデルを分析する。$\alpha \to 0$ のとき、モデルはデカップルされた問題へと回帰する。そこでは、輸送計画はコストを最小化し、変形はランドマーク損失を独立して最小化し、結合項はデカップルされた最小化子間の選択基準として機能する。
  • 命題 2.1: 二乗ユークリッドコストの下で、結合された問題は、緩和された変形写像によるソース測度のプッシュフォワードとターゲット測との間のワッサースタイン距離、および正則化項を含む汎関数を最小化することと等価であることが示される。

• 数値アルゴリズム:

  • 離散化: 変形場に対して、連続テンソル積 $Q_1$ 要素を用いた有限要素法（FEM）を採用する。測度は、重み付き原子測度として離散化される。
  • 交互反復法: 解は以下の交互スキームを用いて計算される：
    1. 輸送更新: $u$ を固定し、周辺制約を扱うためのエントロピー的またはブレグマン型の正則化（KLダイバージェンス）を用いたSinkhorn型の反復を用いて、輸送計画を更新する。
    2. 変形更新: $\pi$ を固定し、$h$ と $\rho$ の二次的な選択（เนื่องจาก一次の最適条件から導かれるため）に基づき、変形係数を更新する。
  • 収束性: 理論的分析により、反復列は有界であり、目的関数値は単調に減少し、 successive な反復間の差は消失することが証明されている。エントロピー正則化版については、反復列が制約付き臨界点に収束することが示されている。

主な貢献

1. 新しいフレームワーク: 輸送計画の最適化と変形場の推定を統合する、新しい結合OTモデルの提案。これは、ランドマークベースのレジストレーションと輸送ベースの分布マッチングの架け橋となる。
2. 理論的基礎: 一般的な変分設定におけるモデルの定義可能性の確立、および限定的分析を通じた古典的OTとの関係の特性付け。
3. アルゴリズムの開発: 系統的な収束分析（KL近接およびエントロピー正則化版を含む）を伴う、有限要素ベースの数値アルゴリズムの開発。
4. 疎な教師あり学習: 疎なランドマークによる教師あり学習とグローバルな分布マッチングを組み合わせることで、高密度で幾何学的に意味のある変形場を回復する能力の提示。

実験結果
提案手法の有効性は、2種類の実験によって検証された。

1. 合成実験: 正解の変形場を持つ魚型の分布を使用。
   • 指標: 変形誤差（$L_2$）、保持されたランドマーク誤差、および密度誤差（$L_1$）。
   • 知見: COT法（KL近接およびエントロピー版の両方）は、「OT (barycentric)」（ランドマーク制約を欠く）や「Landmark-only」（グローバルな分布制約を欠く）を一貫して上回った。特に、ランドマークが疎な領域において顕著であった。手法は、ランドマーク数が増えるにつれて、より低い誤差で正解の変形を回復することに成功した。目的関数と反復列の収束も観察され、理論的予測と一致した。
2. 実形状マッチング: 手動で指定されたランドマークを持つ手書き数字画像（MNIST）に適用。
   • 知見: ベースラインと比較して、COTは空間的に一貫しており、かつ規定されたランドマークの対応関係と一致する変形場を生成した。「OT (barycentric)」はグローバルな分布を捉えるものの不規則な場を生成し、「Landmark-only」はランドマーク付近でのみ滑らかな場を生成したが、COTは分布とランドマークの両方の情報をソースのサポート全体に伝播させることに成功した。

意義と主張
本論文は、提案されたフレームワークが、ランドマークベースのレジストレーションと輸送ベースの分布マッチングの間の原理的な接続を提供すると主張している。輸送計画と変形場を相互一貫性制約を通じて結合することにより、本手法は、疎な幾何学的教師あり学習から輸送写像を回復することを可能にする。著者らは、このアプローチが、輸送コストのみを最小化すること（意味論的に誤った経路を生む可能性がある）や、疎なランドマークのみに依存すること（グローバルな変形を捉えきれない可能性がある）の限界を克服すると断言している。本研究は、モデルの理論的正当性を確立し、形状マッチングタスクにおける実用的な有用性を示しており、幾何学的に意味のある分布変換を必要とするアプリケーションに対する堅牢なソリューションを提供している。

著者らは、フレームワークは一般的であるものの、現在の数値実装は低次元の空間ドメイン向けに設計されていることを謙虚に述べている。そして、ニューラルネットワークによるパラメータ化などを用いた、高次元および大規模問題のための手法の開発を将来の研究方向として挙げている。