Your GFlowNet Secretly Learns an Optimal Transport Plan

技術要約：あなたのGFlowNetは、密かに最適輸送プランを学習している

問題提起
生成フローネットワーク（GFlowNet）は、有向グラフにおける確率的な軌跡を通じて、構造化された離散オブジェクトをサンプリングするように設計されている。もともとは非巡回グラフ向けに開発されたが、近年の研究では非巡回環境への拡張が行われている。このような設定では、サンプリング効率のために期待軌跡長を最小化することが極めて重要であり、「最小フロー（minimum-flow）」GFlowNetの定式化へとつながる。同時に、最適輸送（OT）は、輸送コストを最小化することで確率分布を比較するための原理的な枠組みを提供しており、通常は線形計画法（カントロヴィッチ定式化）を通じて解かれる。しかし、大規模なグラフ上でOTを解くことは、厳密なソルバーを用いると計算量的に困難な場合が多い。

本論文は、この二つの分野の間のギャップを埋めるべく、最小フロー定式化を用いた非巡回GFlowNetが、本質的に最適輸送問題を解いているのではないかという問いを調査する。

手法
著者らは、非巡回GFlowNetとグラフ上の離散最適輸送との間の理論的な等価性を確立している。核心となる手法は以下の通りである：

1. 理論的再定式化： 著者らは、非巡回GFlowNetにおける最小フロー最適化問題から出発する。これは、詳細バランス（detailed balance）と報酬一致（reward matching）の制約の下で、総フロー（期待軌跡長に相当）を最小化するものである。初期エッジフロー分布（ソース状態 $s_0$ から出るフロー）を特定のソース分布 $L(u)$ に一致させるように固定することで、GFlowNetの目的関数を線形計画法（LP）問題へと変換する。
2. 等価性の証明： 彼らは、この制約付きLPが、カントロヴィッチの最適輸送問題と数学的に等価であることを証明した。この定式化において：
   • ソース分布は、一連の初期状態 $U$ 上で定義される $L(u)$ である。
   • ターゲット分布は、終端状態 $X$ 上で定義される報酬分布 $R(x)$ である。
   • 輸送コスト $d(u, x)$ は、グラフにおける $u$ と $x$ の間の最短経路距離として定義される。
   • 最適なGFlowNetポリシーは、総輸送コストを最小化する最適な結合（coupling）$\Pi$ を暗黙的に学習する。
3. 学習アルゴリズム： 大規模なグラフにおいて厳密なLPが実行不可能であるため、著者らはGFlowNetのニューラルパラメータ化を採用する。彼らは、フロー正則化を伴う軌跡バランス（Trajectory Balance; TB）目的関数を利用する。主要な修正点は、ターゲット分布に対する標準的な報酬一致に加え、固定された初期分布 $L(u)$ を強制するための「leward（左報酬）」項を損失関数に含めたことである。

主な貢献

1. 理論的等価性： 本論文は、最小フローGFlowNetにおいて初期フロー分布を固定すると、その目的関数がグラフ誘起の最短経路コストを持つカントロヴィッチOT問題に正確に還元されることを証明した。その結果、最適なGFlowNetの順方向ポリシーは、最適な輸送プランをエンコードし、そこからの軌跡サンプリングは最適な結合を復元する。
2. ベックマン問題への接続： 著者らは、この定式化がグラフベースの離散的なベックマン問題と等価であることを示し、GFlowNetを古典的なネットワークフロー理論へと結びつけた。
3. スケーラブルなOTソルバー： 本研究は、GFlowNetの学習フレームワークが、厳密なOTソルバーでは手に負えないグラフベースのOT問題の近似解を得るために適用できることを示している。ニューラルネットワークと確率的な軌跡サンプリングを用いることで、この手法は組合せ空間へとスケールすることができる。
4. 最短経路理論の回収： 補遺として、この構成は非巡回GFlowNetに対して以前に確立された最短経路理論を回収することを示し、最適なフローが最短経路のサブグラフ上にのみサポートされることを示している。

実験結果
著者らは、以下の2つの環境で理論と手法を検証した：

• ハイパーグリッド環境： 小規模なハイパーグリッド（$H=10, 15, 20$）において、学習されたGFlowNetポリシーは、古典的なLPソルバー（POT）から得られる厳密解と一致する輸送コスト（$OT^\star$）を達成した。この手法は、サンプリングにバイアスを与えることなく、異なるソース分布（ボール形状およびムーン形状）にわたって最適な結合を正常に復元した。
• 置換環境： 対称群 $S_n$ のケイリーグラフにおいて、手法は $n=4$ および $n=8$ で（厳密な解が実行可能な範囲で）準最適な解を復元した。厳密なOTが困難な $n=20$ においても、手法は固定点確率において低い誤差を伴う妥当な近似を生成した。
• アブレーション解析： 正則化係数 $\lambda$ に関する実験により、軌跡の長さ（サンプリング効率）とサンプリングの最適性（輸送プランの正確性）の間のトレードオフが確認された。$\lambda$ の値が低いほど、より正確な輸送プランが得られるが、軌跡は長くなる。

意義と主張
本論文は、フローを最小化するように定式化されたGFlowNetが、暗黙的に最適輸送問題を解いているという「秘密の」つながりを明らかにしていると主張している。本研究の意義は以下の点にある：

• 統一： GFlowNet学習と最適輸送を統合する新しい理論的視点を提供し、GFlowNetが輸送マップを学習する手法として見なせることを示唆している。
• 実用的応用： ニューラルネットワークとサンプリングによってスケールするGFlowNet学習アルゴリズムを用いることで、厳密なソルバーでは計算量的に不可能な大規模グラフのOT問題への適用を可能にする。
• フレームワークの拡張： サンプリングや生成タスクを超えて、計算論的最適輸送の領域へとGFlowNetフレームワークの汎用性を拡張し、質量が「どこへ」輸送されるかだけでなく、実現可能な局所的な移動を通じて「どのように」輸送されるかを決定する、近最適な経路コストを持つ確率的ポリシーを提供する。

著者らは、将来的な影響については謙虚な姿勢を維持しており、本手法は計算効率が高いものの、最適性と効率性のバランスをとるために正則化係数 $\lambda$ の慎重なチューニングが必要であると述べている。彼らは、今後の研究として、非正規化分布間の輸送マップを学習するためにこれらの手法を適用することを検討している。