[论文解读] 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 学习基于流的生成策略，按正向奖励的概率比例抽样候选，实现非迭代、多样化的高奖励输出（例如分子）。

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.

研究动机与目标

• 促使对高奖励候选的分布进行采样，而非仅选取一个最高候选。
• 引入生成过程的流网络视角，以处理导致同一最终对象的多条行动序列。
• 给出一个学习目标，强制流的一致性并得到与奖励成正比的策略。
• 在合成域和分子生成任务中展示可扩展性与多样性提升。
• 展示训练目标的离线/离策略收敛性性质。

提出的方法

• 将生成过程建模为一个流网络，含单一源点和对应候选对象的终端汇点。
• 将动作概率定义为 pi(a|s) = F(s,a)/F(s)，其中 F 表示边上的流量，且 F 满足流一致性方程。
• 在对数流上使用流匹配目标以稳定训练：最小化进入流和流出流的对数差的平方（式(12)）.
• 证明流目标的全局最小值产生一个在终端状态下与 R(x) 成正比的 pi(x) 的策略。
• 证明该方法是离策略的：训练可以使用来自广泛支持策略的轨迹，同时收敛到期望的流和分布（命题3）。
• 应用于合成的 DAG/MDP 域以及带有基于对接的奖励代理的大规模分子生成。

实验结果

研究问题

• RQ1 基于流网络的策略是否在存在多条行动序列到达同一最终状态时，仍能使样本的概率与给定奖励 R(x) 成正比？
• RQ2 流匹配目标是否能够恢复所有高奖励模式，并在多样性和样本效率方面优于传统的 MCMC/RL 基线？
• RQ3 该方法是否能扩展到如分子生成等大型组合空间，并对离线/离策略训练具有鲁棒性？
• RQ4 在合成与分子领域的实际应用中，生成的分布与归一化后的奖励分布相比如何？

主要发现

• GFlowNet 在 C 为非单射的情况下（存在到达同一状态的多条路径），仍以与 R(x) 成正比的概率对终端状态进行采样。
• 对数流匹配目标给出与期望分布对应的全局最小值，并使得离策略/离线训练成为可能。
• 在超网格域中，GFlowNet 收敛到 pi(x) ∝ R(x)，样本量少于 MCMC，并对模态分离具有鲁棒性。
• 在大规模分子生成中，GFlowNet 比 MARS、PPO 及 JT-VAE+BO 基线产生更高奖励且更具多样性的分子，速度更快。
• GFlowNet 发现并覆盖了许多模态（例如，R>8 的模态数量>1500），且顶端样本之间的平均两两相似度低于基线，表明更高的多样性。
• 该方法可扩展到具有可观动作空间的分子图，并在相同计算预算内产出更多独特且高分的分子。