技術的概要：プロキシを超えて：オフライン GFlowNet 学習のための軌道蒸留ガイダンス

問題定義

生成フローネットワーク（GFlowNets）は、報酬に比例した構成オブジェクトのサンプリングを行う強力なモデルである。しかし、分子発見、生物学的配列設計、推薦システムなど、多くの実用的応用において、実験コストの高さや人間の評価の必要性から、報酬を問い合わせるための環境との能動的な相互作用は不可能である。その結果、GFlowNets は静的なオフラインデータセットを用いて学習しなければならない。

現在のオフライン学習パラダイムは、通常、データセットから報酬信号を近似するプロキシ報酬モデルを学習し、GFlowNet が方策最適化のためにそのモデルに報酬を問い合わせる方式に依存している。このアプローチは以下の重大な課題に直面している：

1. データ不足とコスト: 信頼性の高いプロキシモデルの構築には、大規模で多様なデータセットと専門的なドメイン知識が必要となる。
2. 誤差の伝播: 分布外（OOD）のサンプルや状態に対する報酬問い合わせが、プロキシによって不正確に推定された場合、その誤差は勾配を通じて伝播し、学習された方策を損なう可能性がある。
3. プロキシフリー手法の限界: 既存のプロキシフリー手法（例：RO-GFlowNets、COFlowNet）は、方策をデータセットに整合させるために粗い制約を課すことが多い。これらの制約は汎化を制限し、効果的な探索を阻害し、データ品質への感度を高める可能性がある。

ここで扱われる中心的な問いは以下の通りである：終端報酬が信頼性の高いプロキシモデルの構築に不十分な場合、オフライン軌道を活用して GFlowNets を効果的に学習するにはどうすればよいか？

手法：軌道蒸留 GFlowNet（TD-GFN）

著者は、TD-GFNを提案する。これは、プロキシ報酬モデルに依存することなく、環境の有向非巡回グラフ（DAG）から豊富な構造的情報を抽出して方策学習を導く、新規のプロキシフリーフレームワークである。この手法は以下の 3 つの主要フェーズから構成される：

1. 逆強化学習（IRL）によるエッジ報酬の抽出

終端報酬を予測する代わりに、TD-GFN は IRL を利用して、オフライン軌道から遷移レベルのエッジ報酬（$R_E$）を蒸留する。

• リバランス: オフラインデータが専門家行動を代表していないという問題に対処するため、著者はリバランス戦略を導入する。軌道は終端報酬に比例する確率（$P(\tau) \propto R(s_T)$）でサンプリングされ、疑似専門家データセットが作成される。
• 敵対的学習: 修正された GAIL（Generative Adversarial Imitation Learning）フレームワークを用いて、識別子 $D_\phi$ がリバランスされたデータセット内の遷移と、方策 $\pi_\psi$ によって生成された遷移を区別する。
• 報酬導出: 識別子から不偏なエッジレベルの報酬関数が抽出される：
  $$R_E(s, s') = \log D_\phi(s, s') - \log(1 - D_\phi(s, s'))$$
  理論的に、この報酬は対数条件付き流入フロースコア（$\log P_B(s|s')$）を近似し、効果的な方策を学習するための特定の遷移の重要性を定量化する。重要なのは、これらの報酬が終端予測ではなく、密な中間ガイダンスを提供する点である。

2. 報酬誘導型 DAG プリニング

堅牢性を確保し、潜在的に不正確な IRL 推定からの誤差伝播を防ぐため、TD-GFN はエッジ報酬を直接方策勾配の形成に使用しない。代わりに、それらを用いて環境 DAG をプリニング（剪定）する。

• 閾値処理: 報酬分布の平均から $K$ 倍の標準偏差を引いた統計的閾値より低い報酬を持つエッジは削除される。
• 理論的保証: 著者は、このプリニング戦略が「高ボトルネック」モード（エッジスコアの高い経路を介して到達可能な終端状態）をすべて保持しつつ、低有用性の遷移を除去することを証明している。これにより、高価値の軌道を優先しつつ、方策を IRL 近似誤差から保護する、コンパクトで表現力のある行動空間が作成される。

3. 優先付き backward サンプリング

学習は、優先付き backward サンプリングメカニズムを用いて、プリニングされた部分グラフ上で行われる。

• サンプリング戦略: 終端状態は、データセットとプリニングされたグラフの交差部分から、真の報酬に比例する確率でサンプリングされる。
• backward 軌道構築: サンプリングされた各終端から、学習されたエッジ報酬によって定義される学習済み backward 方策 $P_B$ を用いて、軌道が逆方向に根に向かって構築される：
  $$P_B(s_t | s_{t+1}) = \frac{\exp\{R_E(s_t, s_{t+1})\}}{\sum \exp\{R_E(s, s_{t+1})\}}$$
• 勾配更新: GFlowNet 方策は、これらのサンプリングされた軌道とそれらの真の終端報酬を用いて最適化される。勾配更新に記録された終端報酬のみを依存させることで、この手法はプロキシベースの手法に内在する誤差伝播の問題を回避する。

主要な貢献

1. 新規プロキシフリーパラダイム: TD-GFN はプロキシモデルと OOD 報酬問い合わせへの依存を排除する。IRL を用いてリバランスされたデータセットから遷移レベルのエッジ報酬を蒸留し、方策学習と汎化のための構造化ガイダンスを提供する。
2. 堅牢なガイダンスメカニズム: このフレームワークは、DAG プリニングと優先付き backward サンプリングという 2 つの間接的ガイダンスメカニズムを導入する。この設計により、勾配更新が真の報酬にのみ依存し、誤差伝播を防ぎつつ、効率的な探索のための密な構造的ガイダンスを活用することを保証する。
3. 最先端のパフォーマンス: 広範な実験により、TD-GFN はプロキシベース手法（Proxy-GFN）、オフライン RL アルゴリズム（CQL、IQL）、その他のプロキシフリー GFlowNet 変種（COFlowNet、Dataset-GFN）を含む多様なベースラインと比較して、優れた収束速度とサンプル品質を達成することが示された。

実験結果

著者は TD-GFN を 3 つのベンチマークで評価した：

• Hypergrid（合成）: 84 次元グリッドでスパースかつ多モーダルな報酬が設定された環境において、TD-GFN は 5,000 回未満の状態訪問で 16 個のすべての報酬モードを特定し（ベースラインの 6 倍の改善）、COFlowNet を含む最良の競合他社よりも 2 倍速く最小の経験的 L1 誤差に収束した。また、ノイズの多いデータ（Mixed データセット）やデータ不足（データの 10%）下でも強力な堅牢性を示した。
• Biosequence Design（抗菌ペプチド）: TD-GFN は、上位 100 配列の中で最高報酬かつ最大の多様性を持つ配列を生成し、プロキシベース手法および他のオフライン手法を凌駕した。
• Molecule Design: TD-GFN は、より多くの高報酬分子と多様なモードを発見し、収束に必要な軌道数を大幅に削減した（Proxy-GFN の 43.7 × 10^4 に対し 2.75 × 10^4）。また、真の報酬を直接使用する Oracle-GFN と同等の性能で高報酬分子を特定した。

意義と主張

本論文は、TD-GFN がオフライン GFlowNet 学習のためのより堅牢で安定し、効率的なパラダイムを確立すると主張している。オフライン軌道に内在する構造的情報を IRL を通じて活用することで、この手法は限られたオフラインデータセットと生成モデルの複雑な探索要件の間のギャップを埋める。

著者は、TD-GFN が以下の点で優れていることを強調している：

• 可能性の解放: 能動的な報酬問い合わせが困難な複雑な領域において、GFlowNets の展開を可能にする。
• 誤差伝播の防止: ガイダンスメカニズム（エッジ報酬）と勾配更新メカニズム（終端報酬）を分離することで、プロキシモデルの落とし穴を回避する。
• 探索の強化: 密な構造的ガイダンスにより、モデルは訓練データで疎に表現されているか未観測であっても、高報酬領域を効果的に探索できる。

この研究は、プロキシベースの手法でしばしば見落とされてきた軌道レベルの情報が、効果的なオフライン GFlowNet 学習の礎石であり、創薬、生物学的設計、組み合わせ最適化における実世界応用のためのスケーラブルな解決策を提供することを示唆している。