[論文レビュー] Flow Network based Generative Models for Non-Iterative Diverse Candidate Generation

Emmanuel Bengio, Moksh Jain+3 more|arXiv (Cornell University)|Jun 8, 2021

Machine Learning in Materials Science被引用数 49

ひとこと要約

GFlowNet は、正の報酬に比例する確率で候補をサンプルするフロー-based生成ポリシーを学習し、非反復的で多様な高報酬出力（例: 分子）を可能にする。

ABSTRACT

This paper is about the problem of learning a stochastic policy for generating an object (like a molecular graph) from a sequence of actions, such that the probability of generating an object is proportional to a given positive reward for that object. Whereas standard return maximization tends to converge to a single return-maximizing sequence, there are cases where we would like to sample a diverse set of high-return solutions. These arise, for example, in black-box function optimization when few rounds are possible, each with large batches of queries, where the batches should be diverse, e.g., in the design of new molecules. One can also see this as a problem of approximately converting an energy function to a generative distribution. While MCMC methods can achieve that, they are expensive and generally only perform local exploration. Instead, training a generative policy amortizes the cost of search during training and yields to fast generation. Using insights from Temporal Difference learning, we propose GFlowNet, based on a view of the generative process as a flow network, making it possible to handle the tricky case where different trajectories can yield the same final state, e.g., there are many ways to sequentially add atoms to generate some molecular graph. We cast the set of trajectories as a flow and convert the flow consistency equations into a learning objective, akin to the casting of the Bellman equations into Temporal Difference methods. We prove that any global minimum of the proposed objectives yields a policy which samples from the desired distribution, and demonstrate the improved performance and diversity of GFlowNet on a simple domain where there are many modes to the reward function, and on a molecule synthesis task.

研究の動機と目的

• 単一のトップ候補ではなく、高報酬の候補の分布をサンプリングする動機づけ。
• 生成過程をフロー・ネットワークの視点で捉え、同じ最終オブジェクトに至る多数の行動列を扱えるようにする。
• フローの整合性を課す学習目的を提供し、報酬に比例したポリシーを生み出す。
• 合成領域と分子生成タスクにおけるスケーラビリティと多様性の利得を示す。
• トレーニング目的の off-policy / offline 収束特性を示す。

提案手法

• 生成過程を、候補オブジェクトに対応する単一のソースと末端のシンクを持つフロー・ネットワークとしてモデル化する。
• アクション確率を pi(a|s) = F(s,a)/F(s) と定義する。ここで F は辺を通るフローを表し、F はフロー整合性方程式を満たす。
• トレーニングを安定化させるために、対数フロー上のフローメッチング目的を使用する：入ってくるフローと出ていくフローの対数の差の二乗を最小化する（式 12）。
• フロー目的関数の全体的最小値が、端点状態に対して pi(x) が R(x) に比例するポリシーを生み出すことを証明する。
• この方法が off-policy であることを示す：トレーニングは広いサポートを持つポリシーからの軌跡を使用しつつ、望ましいフローと分布へ収束する（命題 3）。
• 合成 DAG/MDP ドメインと、ドッキングベースの報酬プロキシを用いた大規模な分子生成に適用する。

実験結果

リサーチクエスチョン

• RQ1 複数の行動列が同じ最終状態に至る場合でも、流れネットワークベースのポリシーは、サンプルを与えられた報酬 R(x) に比例する確率で生成できるか？
• RQ2 フローメッチング目的はすべての高報酬モードを回復し、多様性とサンプル効率の点で従来の MCMC/RL ベースラインを上回るか？
• RQ3 分子生成のような大規模組合せ空間に対して方法はスケーラブルで、オフライン/オフポリシー学習に対して頑健か？
• RQ4 合成領域と分子領域の実践で、生成分布は正規化された報酬分布とどのように比較されるか？

主な発見

• GFlowNet は、C が非注入（同じ状態へ多くの経路がある場合）でも、端点状態を R(x) に比例する確率でサンプルする。
• 対数フローメッチング目的は、望ましい分布に対応する全体最小を生み出し、off-policy/offline トレーニングを可能にする。
• ハイパーグリッド領域では、GFlowNet は MCMC より少ないサンプルで pi(x) ∝ R(x) に収束し、モード分離に対して頑健である。
• 大規模分子生成では、GFlowNet は MARS、PPO、JT-VAE+BO のベースラインより高報酬でより多様な分子をより速く生成する。
• GFlowNet は多くのモードを発見・カバーする（例：R>8 のモードが 1500 超）し、トップサンプル間の平均ペアワイズ類似度をベースラインより低く達成し、より高い多様性を示す。
• このアプローチは sizable な行動空間を持つ分子グラフにもスケールし、同じ計算予算内でより多くのユニークな高得点分子を生み出す。