[論文レビュー] Trajectory balance: Improved credit assignment in GFlowNets

Nikolay Malkin, Moksh Jain+3 more|arXiv (Cornell University)|Jan 31, 2022

Topic Modeling被引用数 12

ひとこと要約

本稿では、局所的なフロー制約ではなく、完全な軌道全体を最適化することによって、長時間にわたるアクション列における信用配分を改善する新しい学習目的「トラジェクトリーバランス」を提案する。理論的に、トラジェクトリーバランスのグローバル最小化が、ターゲット分布からの正確なサンプリングを保証することを示し、実験的にも、長時間スケールおよび大規模アクション空間設定において、収束速度の向上、サンプルの多様性の向上、ロバストネスの向上を確認した。

ABSTRACT

Generative flow networks (GFlowNets) are a method for learning a stochastic policy for generating compositional objects, such as graphs or strings, from a given unnormalized density by sequences of actions, where many possible action sequences may lead to the same object. We find previously proposed learning objectives for GFlowNets, flow matching and detailed balance, which are analogous to temporal difference learning, to be prone to inefficient credit propagation across long action sequences. We thus propose a new learning objective for GFlowNets, trajectory balance, as a more efficient alternative to previously used objectives. We prove that any global minimizer of the trajectory balance objective can define a policy that samples exactly from the target distribution. In experiments on four distinct domains, we empirically demonstrate the benefits of the trajectory balance objective for GFlowNet convergence, diversity of generated samples, and robustness to long action sequences and large action spaces.

研究の動機と目的

• フローマッチングや詳細バランスのような従来のGFlowNet学習目的では、長時間アクション列において時間的信用配分が遅くなる問題を解決する。
• 構成的オブジェクト生成における複雑で長い軌道全体にわたる、より効率的かつ正確な信用配分を可能にする新しい学習目的を開発する。
• トラジェクトリーバランスがグローバルに最小化されると、ターゲットの非正規化密度 R(x) からの正確なサンプリングが保証されることを証明する。
• 4つの異なるドメインにおいて、収束速度、サンプル多様性、長時間スケールおよび大規模アクション空間におけるロバストネスという観点から、トラジェクトリーバランスの優位性を実証的に検証する。
• GFlowNetsにおける詳細バランス目的の最初の実証的評価を提供し、トラジェクトリーバランスおよびフローマッチングと比較可能な分析を可能にする。

提案手法

• 初期状態から終端状態までの完全なアクション列（軌道）を対象として、局所的エッジフローまたはペアワイズ状態遷移ではなく、グローバルな軌道を最適化する新しい学習目的「トラジェクトリーバランス」を提案する。
• トラジェクトリーバランス目的を、サンプリングされた各軌道の寄与度の合計として定義し、各軌道の寄与度は、経路に沿ったフローとターゲット報酬 R(x) の乖離に基づく。
• エッジフロー F(u→v) をニューラルネットワークでパラメータライズし、ポリシーを F(u→v)/∑v′F(u→v′) として定義することで、確率的アクション選択を保証する。
• 勾配降下法を用いてトラジェクトリーバランス目的を最適化し、微分可能なモデルではパスワイズ微分と再パラメータライゼーションを用いて勾配を計算する。
• トラジェクトリーバランス目的の任意のグローバル最小化解が、ターゲット分布 p(x) ∝ R(x) からの正確なサンプリングを実現するポリシーを生成することを証明する。
• 複数のベンチマーク環境において、標準的なGFlowNet学習パイプラインを用いて、トラジェクトリーバランスをフローマッチングおよび詳細バランス目的と比較実装する。

実験結果

リサーチクエスチョン

• RQ1 局所的フローベースの目的（フローマッチングや詳細バランス）と比較して、軌道ベースの目的はGFlowNetsにおける信用配分効率を向上させるか？
• RQ2 トラジェクトリーバランスは、長時間スケール生成タスクにおいて収束速度の向上とサンプル多様性の向上をもたらすか？
• RQ3 構成的オブジェクト生成において、大規模アクション空間および高次元状態空間に対しても、トラジェクトリーバランスはロバストか？
• RQ4 サンプル品質および学習安定性という観点から、実証的にトラジェクトリーバランスはフローマッチングおよび詳細バランスを上回るか？
• RQ5 トラジェクトリーバランスは、従来の方法では到達できなかった長時間スケールのシーケンスでもGFlowNetsの学習を可能にするか？

主な発見

• トラジェクトリーバランスは、長時間シーケンスおよび大規模アクション空間設定を含む、すべての評価ドメインでフローマッチングおよび詳細バランスよりも収束速度が速い。
• エントロピーおよびカバレッジなどの指標で測定したサンプル多様性が、ベースライン目的と比較して顕著に向上している。
• トラジェクトリーバランスにより、100ステップに達するアクションシーケンスのGFlowNetsの学習が成功裏に達成可能となった。これは、従来の目的では収束しなかった分野である。
• 実証的結果から、生成サンプルとターゲット分布 R(x) 間の乖離が小さくなることが示され、望ましい報酬関数との整合性が向上している。
• 詳細バランス目的の最初の実証的評価により、理論的整合性が確認されたが、トラジェクトリーバランスと比較して収束が遅く、サンプル分散が高かった。
• トラジェクトリーバランス目的のグローバル最小化が、理論的にも正確にターゲット分布 p(x) ∝ R(x) からのサンプリングを保証することを示した。