Scalable Bi-causal Optimal Transport via KL Relaxation and Policy Gradients

技術的サマリー：KL 緩和と方策勾配によるスケーラブルな双因果最適輸送

問題定義
双因果最適輸送（OT）は、非予知の情報制約下での確率過程の比較と結合（カップリング）のための枠組みを提供する。古典的な OT と異なり、双因果 OT は、各過程の次の状態の条件付き法則が、現在の時点までの結合履歴のみに依存するように、許容されるカップリングを制限する。これにより、いかなる過程も他方の過程の未来を予知しないことが保証される。この構造は、モデル不確実性下でのヘッジング（ロバストファイナンス）や時系列のダウンスケーリング（逐次的不確実性定量化）などの応用において極めて重要である。

しかしながら、双因果 OT の実用的な応用は、計算上の非実行性によって阻害されている。エントロピー正則化と Sinkhorn 反復、またはフィットされた値反復などの既存の手法は、連続分布、高次元の状態空間、および長い時間地平において困難に直面する。主要な難点は、経路空間全体にわたる硬い周辺制約を強制することにある。これは、直接最適化を組み合わせ論的に困難かつ計算コストの高いものとしている。

手法
著者らは、硬い周辺制約を微分可能なカルバック・ライブラー（KL）ペナルティに置き換えるスケーラブルな確率的最適化枠組みを提案し、方策勾配法の利用を可能にする。

1. KL ペナルティ緩和
   著者らは、任意の結合遷移カーネルを含むように許容カップリングの集合を拡大することで、元の双因果 OT 問題を緩和する。指定された周辺遷移（$T_n, T'_n$）および初期周辺（$\mu_0, \mu'_0$）からの逸脱を制御するため、KL 発散ペナルティを導入する。緩和された目的関数 $W_{rel}$ は、期待累積コストとパラメータ $\beta$ で重み付けされたペナルティ項の和を最小化する：
   $$W_{rel}(\mu, \mu') = \inf_{\pi \in \Pi_{rel}} \mathbb{E}_\pi \left[ \sum_{n=0}^N c_n(Y_n, Y'_n) + \beta \cdot \text{KL ペナルティ} \right]$$
   この定式化は問題の再帰的構造を保持し、動的計画法を可能にする。

2. 理論的基盤

   • 動的計画法: 本論文は、元の制約付き問題と緩和された定式化の両方に対して動的計画法の原理を確立する。これにより、経路空間全体にわたるグローバルな最適化が、初期カップリングと遷移カーネルにわたる一連の 1 段階最適化へと帰着される。
   • $\Gamma$-収束: 著者らは、ペナルティパラメータ $\beta \to \infty$ において、緩和された目的関数とその最適カップリングが、元の双因果 OT 問題とその最適解に収束することを証明する。これにより、この緩和が単なるヒューリスティックではなく、原理的な代理であることを保証する。

3. 方策勾配アルゴリズム

   • パラメータ化: カップリングは、正確な対数密度評価を可能にする表現力のある条件付きモデル、具体的には正規化フロー（ニューラルスプライン層）を用いてパラメータ化される。
   • 勾配表現: 著者らは、緩和された目的関数に対する明示的な方策勾配表現を導出する。勾配は 2 つの成分に分解される：
     • 値成分（$\nabla J_{val}$）: 現在のパラメータ化されたカップリング下での期待値であり、残りの累積コストで重み付けされる。
     • KL 成分（$\nabla J_{KL}$）: 指定された周辺ダイナミクス（基準測度）下での期待値であり、ペナルティ項を捉える。
   • 推定: アルゴリズムは不偏ミニバッチ勾配推定量を採用する。分散を低減するため、補助的な軌道にフィットさせた状態依存の制御変数を利用する。
   • 収束: 本論文は、Polyak–Łojasiewicz 条件下で平均後悔が $O(1/K)$ のオーダーであることを示す非漸近的収束保証を提供する。ここで $K$ は反復回数である。

主な貢献

• 緩和理論: KL 緩和問題に対する動的計画法の原理の確立と、元の制約付き問題への $\Gamma$-収束の証明。
• スケーラブルなアルゴリズム: 不偏推定量と分散低減を備えた方策勾配アルゴリズムの開発。状態空間や行動空間の離散化を必要とせず、一般の周辺（連続分布および高次元の場合を含む）を処理可能。
• 収束保証: 提案された確率的最適化手法に対する非漸近的後悔 bound（$O(1/K)$）の導出。
• 実証的検証: 合成設定および実世界応用において、周辺法則と時間的依存性の両方を捉える能力の示唆。

結果
本論文は、合成実験と 2 つの具体的な応用を通じて枠組みを検証する：

1. 合成実験:

   • 本手法は、単峰および双峰の両方の設定において、周辺分布（ガウス分布および非ガウス/ベータ分布）と時間的依存性（隣接ステップ間の相関）を正確に回復する。
   • 性能指標（ワッサーシュタイン距離、KS 距離、スライスド・ワッサーシュタイン距離）は、生成された分布と真の分布の間の低い不一致を示す。
   • 本手法は、より高次元（$d=3$）およびより長い地平（$N=10$）へも効果的にスケーリングする。

2. ロバスト・ヘッジング（金融応用）:

   • この枠組みは、非予知制約下での金融派生商品のロバスト・サブヘッジ価格の計算に応用される。
   • 学習されたカップリングは、閉形式の解と比較して相対誤差 1% 未満の近似的な最適価格帯を回復する。
   • 本手法はマルティンゲールダイナミクスを保持し、計算速度において既存の手法（例えば、フィットされた値反復法よりも高速）を上回る。

3. 時系列統計的ダウンスケーリング:

   • 本手法は、双因果 OT による低解像度データのバイアス除去を介して、低忠実度の観測から高忠実度の時系列を生成するために使用される。
   • 「OT 調整済み」の条件付き拡散モデルは、無条件生成および粗いデータへの直接条件付けを大幅に上回り、より低い分布的不一致とより正確な時間的相関を達成する。

意義と主張
本論文は、非予知情報制約が不可欠な設定において、理論的枠組みと実用的な実装の間のギャップを埋める、双因果 OT に対するスケーラブルな計算アプローチを提供すると主張する。

• 計算の実現可能性: 制約付き経路空間の問題を微分可能な確率的最適化問題に変換することで、本手法は連続分布および長い地平における双因果 OT を実行可能にし、次元の呪いと、以前の手法に関連するメモリコストを克服する。
• 汎用性: この枠組みは、離散的または経験的サンプルに限定されず、一般的な周辺に適用可能である。
• 実用性: 結果は、双因果 OT が、意思決定の質にとって情報フローの尊重が重要なロバスト・ヘッジングや統計的ダウンスケーリングなどのタスクにおいて、オペレーションズリサーチおよび定量ファイナンスにおける大規模な最適化および不確実性定量化パイプラインのための実用的な計算プリミティブとなり得ることを示唆している。

著者らは、将来の作業については控えめなトーンを維持しており、現在の分析はマルコフ設定に焦点を当てているが、この枠組みは非マルコフ設定や経路依存コストへ拡張可能であると指摘している。ただし、これにはさらなる理論的発展が必要である。