Fixed-Point Neural Optimal Transport without Implicit Differentiation

技術的サマリー：陰微分を用いない固定点ニューラル最適輸送

問題定義
最適輸送（OT）は、所定の輸送コストを最小化する確率分布間の写像を求めるものである。近年の深層学習アプローチによりスケーラブルな OT ソルバが可能になったが、これらは頻繁に敵対的 min–max 最適化（例：WGAN）や、輸送写像と双対ポテンシャルを含むマルチネットワークアーキテクチャに依存している。これらの手法は、特に高次元設定において、訓練の不安定性、ハイパーパラメータへの感度、収束の困難さに悩まされることが多い。連続流に基づく代替アプローチは常微分方程式（ODE）を解く必要があり、大きな計算オーバーヘッドを伴う一方、Hamilton–Jacobi–Bellman（HJB）定式化は高次元におけるスケーラビリティに苦慮している。さらに、多くの既存手法は、双対実行可能性を強制するために補助ネットワークや交互最適化手順を必要とする。

手法
著者らは、敵対的訓練と陰微分を排除する、二次最適輸送のための単一ネットワークフレームワークを提案する。この手法の核心は、Kantorovich 双対問題から導出された陰的ニューラル定式化にある。

1. 単一ポテンシャルパラメータ化: 双対ポテンシャル $\phi$ は、単一のニューラルネットワークで表現される凸関数 $g_\theta$ の負としてパラメータ化される（$\phi_\theta(x) = -g_\theta(x)$）。これは、最適輸送写像が凸ポテンシャルの勾配であるとする Brenier の定理を利用している。
2. 近傍固定点再定式化: このパラメータ化の下で、双対目的関数の主要な構成要素である $c$-変換は、強凸な近傍最適化問題に帰着する：
   $$(-g_\theta)^c(z) = \inf_{y \in \mathbb{R}^d} \left\{ \frac{1}{2}\|y - z\|^2 + g_\theta(y) \right\}$$
   この問題の最小化点 $y^*_\theta(z)$ は、後方輸送写像（$T_{\nu \to \mu}(z) = y^*_\theta(z)$）を定義し、前方写像は $T_{\mu \to \nu}(x) = x + \nabla g_\theta(x)$ で与えられる。
3. 固定点反復による訓練: 訓練目的関数は、負の双対損失の最小化を含む。重要なのは、内部最小化（近傍演算子）が、展開最適化や陰微分ではなく、標準的な固定点反復（例：勾配降下ステップ）を用いて解かれる点である。
4. 陰微分なしの勾配計算: 重要な理論的洞察として、損失のネットワークパラメータ $\theta$ に対する勾配は、固定点反復を通じた逆伝播を必要としない。近傍演算子の一次最適性条件により、固定点の $\theta$ に対する微分（$dy^*/d\theta$）を含む項が完全に打ち消し合う。その結果、勾配は入力サンプルおよび近似固定点解において評価された $g_\theta$ の明示的な微分のみに依存する。これにより、標準的な逆伝播を用いた効率的な訓練が可能となる。
5. 収束性: 著者らは、近似誤差が有界であれば、内部固定点問題が指定された許容誤差内で近似して解かれた場合でも、確率的勾配降下法（SGD）が定常点に収束することを確立している。
6. クラス条件付き拡張: このフレームワークは、ポテンシャル $g_\theta(x, k)$ をクラスラベル $k$ に条件付けることで、クラス条件付き OT に自然に拡張可能であり、単一のネットワークで複数の輸送タスクを同時に処理できる。

主な貢献

• 単一ネットワーク再定式化: 著者らは、前方および後方の輸送写像の両方を回復するために単一のニューラルネットワークのみを必要とする二次 OT の定式化を導出した。これにより、敵対的 min–max 設定や補助共役ネットワークの必要性が排除された。
• 陰微分フリー訓練: 目的関数が、近傍演算子の最適性条件によって誘発される項の完全な打ち消しに依存し、陰微分なしで微分可能であることを実証した。
• 収束保証: 論文は、内部ループが機械精度まで解かれないという現実的な状況に対処し、不正確な固定点評価下での SGD の定常点への収束に関する理論的証明を提供している。
• 統合フレームワーク: この手法は、前方および後方写像を同時に回復し、アーキテクチャの変更なしにクラス条件付き設定に自然に拡張される。

実験結果
提案手法は、高次元ガウスベンチマーク、実世界の物理データセット（UCI）、および画像変換タスク（MNIST から Fashion-MNIST へ）で評価された。

• 高次元ガウス: $d=64$ までの次元を持つガウスベンチマークにおいて、この手法は NOT、MM-v1、MM:R、NCF、WGAN-QC などのベースラインと比較して、著しく低い説明不能分散率（UVP）を達成した。次元が増加するにつれて安定した精度を維持したのに対し、ベースラインは劣化を示した。
• 計算効率: この手法は、連続流ベースの手法（FFJORD、RNODE）や他のニューラル OT 手法よりも少ないパラメータ数と短い訓練時間を要し、有利な計算およびメモリ効率を示した。
• 実世界データ: 多様な UCI データセット（Power、Gas、HEPMASS、MINIBOONE）および BSDS300 画像パッチにおいて、この手法は著しく小さいモデルサイズで、競争力のある、あるいは優れた分布一致性能（最大平均不一致で測定）を達成した。
• 画像変換: クラス条件付き画像変換タスクにおいて、この手法はベースラインの中で最低の Fréchet Inception Distance（FID）と最高の分類精度を達成し、クラス構造と意味内容の優れた保持を示した。
• アブレーション研究: 実験により、この手法は固定点ソルバの停止許容誤差に対して頑健であることが示された。内部ループにおける中程度の不正確さは、精度と速度の間の有利なトレードオフをもたらす。さらに、重み投影による厳密な凸性の強制は固定点反復を安定させたが、全体の輸送精度を低下させた。これは、（硬い投影なしの）非凸パラメータ化が表現力と安定性のより良いバランスを提供することを示唆している。

意義と主張
本論文は、このフレームワークが既存のニューラル OT 手法に対する安定性、効率性、スケーラビリティに優れた代替手段であると主張している。敵対的訓練と、全体領域における ODE ソルバまたは HJB 方程式求解の計算オーバーヘッドを排除することで、この手法は最適性条件を強制する「因果性フリー」のメカニズムを提供する。著者らは、このアプローチが広範に適用可能であり、高次元およびクラス条件付き問題の処理が可能で、双方向に正確な輸送写像を回復できると述べている。この研究は、OT の最適化風景を単純化する一歩として位置づけられ、複雑な min–max 定式化から、理論的根拠を持ち実用的に効率的な単一目的最小化問題へと移行することを意図している。