[論文レビュー] Learning GFlowNets from partial episodes for improved convergence and stability

Kanika Madan, Jarrid Rector-Brooks+7 more|arXiv (Cornell University)|Sep 26, 2022

Reinforcement Learning in Robotics被引用数 5

ひとこと要約

本稿では、部分的アクション部分列（長さが異なる）から学習可能な、勾配バイアスと分散のバランスを取る新しい訓練目的であるSubtrajectory Balance（SubTB）を提案する。SubTBは収束を加速し、スパースまたは長時間スケールの環境における安定性を向上させ、ペプチドやタンパク質生成を含む合成的および実世界のタスクにおいて、既存の目的（TB や FM）を上回る報酬と多様性の指標を達成する。

ABSTRACT

Generative flow networks (GFlowNets) are a family of algorithms for training a sequential sampler of discrete objects under an unnormalized target density and have been successfully used for various probabilistic modeling tasks. Existing training objectives for GFlowNets are either local to states or transitions, or propagate a reward signal over an entire sampling trajectory. We argue that these alternatives represent opposite ends of a gradient bias-variance tradeoff and propose a way to exploit this tradeoff to mitigate its harmful effects. Inspired by the TD($λ$) algorithm in reinforcement learning, we introduce subtrajectory balance or SubTB($λ$), a GFlowNet training objective that can learn from partial action subsequences of varying lengths. We show that SubTB($λ$) accelerates sampler convergence in previously studied and new environments and enables training GFlowNets in environments with longer action sequences and sparser reward landscapes than what was possible before. We also perform a comparative analysis of stochastic gradient dynamics, shedding light on the bias-variance tradeoff in GFlowNet training and the advantages of subtrajectory balance.

研究の動機と目的

• 局所的目的（例：詳細釣り合い）は低分散だが高バイアスであり、軌道レベルの目的（例：軌道釣り合い）は低バイアスだが高分散であるGFlowNet訓練におけるバイアス-分散トレードオフを解消すること。
• 完全でないまたは部分的なエピソードから学習可能な柔軟な訓練目的を開発し、長時間のアクション列やスパース報酬を持つ環境でのより効率的で安定した訓練を可能にすること。
• 特に困難な確率的モデリング設定において、収束速度の向上とハイパーパrameter選択へのロバスト性を向上させること。
• 学習可能な $\lambda$ パrameterを介して、局所的整合性と完全軌道への責任帰属を統合するフレームワークを提供すること。

提案手法

• 強化学習におけるTD($\lambda$)にインspiredされた、部分軌道の長さを変化させた信用配分のバランスを取る訓練目的であるSubTB($\lambda$)を提案する。
• 部分軌道における時系列差分ターゲットの重み付き和を用い、その重みはエリギビリティトレースパラメータ $\lambda$ によって決定され、局所的（DBに類似）からグローバル（TBに類似）な学習への補間を可能にする。
• 軌道釣り合いの一般化としての目的を導出することで、未完了のエピソードから学習可能にし、将来の報酬に基づいて中間状態への信用配分を可能にする。
• 部分軌道セグメントを用いて勾配を計算する確率的勾配更新ルールを採用し、全軌道バックプロパゲーションに比べて分散を低減する。
• 動的に部分軌道をサンプリングおよび重み付けする微分可能メカニズムを導入し、データ効率性と収束性を向上させる。
• ポリシーをニューラルネットワークでパrameter化することで、未観測状態に対しても一般化可能なGFlowNetsの訓練に本目的を適用する。

実験結果

リサーチクエスチョン

• RQ1 部分エピソードから学習するGFlowNet訓練目的は、全軌道または局所的目的と比較して収束速度と安定性を向上させることができるか？
• RQ2 SubTB($\lambda$)におけるハイパーパラメータ $\lambda$ は、GFlowNet訓練におけるバイアス-分散トレードオフにどのように影響を与えるか？
• RQ3 SubTB($\lambda$)は、従来の手法が失敗する長時間アクション列やスパース報酬関数を持つ環境でも、効果的な訓練を可能にするか？
• RQ4 SubTB($\lambda$)は、軌道釣り合いと比較して、勾配分散をどの程度低減するか？ また、その影響は訓練ダイナミクスにどう現れるか？
• RQ5 SubTB($\lambda$)は、順序生成タスクにおけるサンプル品質（報酬）と多様性の両方を向上させることができるか？

主な発見

• 合成的および実世界の環境において、SubTB($\lambda$)はTB、FM、DBの各目的と比較して、より速い収束とハイパーパラメータ選択への低感受性を達成した。
• 抗菌ペプチド生成タスクでは、SubTB($\lambda$)が平均トップ100報酬0.96 ± 0.02を達成し、TB（0.90 ± 0.03）やFM/DB（0.78 ± 0.05）を顕著に上回った。
• 237長さの配列を持つ蛍光タンパク質生成タスクでは、SubTB($\lambda$)がトップ100報酬1.18 ± 0.10、多様性204.44 ± 0.45を達成し、TB（0.76 ± 0.19 および 204.31 ± 0.44）や他のベースラインを上回った。
• SubTB($\lambda$)で訓練されたモデルは、より多くのモードを発見し、訓練初期段階でターゲット分布とのスピアマン相関が高かったため、より速く効果的な探索が可能であった。
• SubTB($\lambda$)では、TBと比較して確率的勾配の分散が顕著に低く抑えられ、SubTBがバイアスを低く保ちつつ勾配分散を低減するという仮説が裏付けられた。
• SubTB($\lambda$)は、スパース報酬や長時間の軌道を持つ環境でも、従来のGFlowNet目的が困難または収束しなかった状況においても、効果的な訓練を可能にした。