Flow Network based Generative Models for Non-Iterative Diverse Candidate Generation
Emmanuel Bengio, Moksh Jain, Maksym Korablyov, Doina Precup, Yoshua Bengio
NeurIPS 2021
cs.LG
2021-06-08
把物体生成看成DAG上的流量网络,按终端奖励成比例采样;分子对接里高分骨架超过1500个,远多于MARS和PPO。
标准强化学习把期望回报推到最大,策略概率会塌到少数几条最高分轨迹上。分子设计这类任务要的是另一件事:给定一个正的奖励函数,按奖励大小成比例地抽出一批候选,而且这批候选要盖住多个峰,不能围着一个局部最优打转。
MCMC能从非归一化密度里采样,但峰与峰之间隔着低奖励的荒地时,链很难跳过去。把生成过程当成树来做的方法更糟:同一分子可以有很多条组装路径,树式价值估计会按路径条数给分子加权,大分子被指数级优待。药物发现里一轮实验可以测上万个候选,需要的是线性成本的多样化采样,不是再找一个单点最优。
GFlowNet把状态看成有向无环图上的节点,边是确定性动作。源节点是空物体,汇节点是完成的候选,汇上的流出量钉死为该候选的正奖励 R(x)。对每个中间节点,流入必须等于流出。策略按边上的流量比例选动作。只要流量守恒成立,终端状态的采样概率就等于 R(x)/Z,Z 是所有终端奖励之和,跟有多少条路径无关。
训练目标仿 Temporal Difference:在轨迹经过的每个节点上,让预测的流入对数与流出对数对齐。流量跨层级差几个数量级,所以网络输出 log-flow,再用 log-sum-exp 算进出总量,避免根节点流量把叶子冲掉。这个目标不依赖采样策略,只要探索策略能覆盖到状态空间,就可以离策略训练。实现里用 95% 当前策略加 5% 均匀动作做探索。
分子实验里,动作是往连接树上接一个片段,环境保证化学有效性。奖励来自预训练的 MPNN 代理,预测分子对可溶性环氧化物水解酶的对接分数,缩放到正区间。
超立方网格上奖励只在 2^n 个角落附近有峰。把谷底奖励 R0 从 0.1 降到 0.001,MCMC 访问所有峰所需样本指数上升,GFlowNet 的 L1 分布误差几乎不受影响,也比高熵正则的 PPO 更快扫完所有峰。
分子任务状态空间约 10^16,每步动作 100 到 2000 个。训练最多看 10^6 个分子。代理训练集里分数超过 8 的只有 233 个;GFlowNet 训练过程中找到 2339 个不重复的分数>8 分子,其中仅 39 个出现在训练集。按 Bemis-Murcko 骨架计,奖励>8 的模式超过 1500 个,MARS 不到 100。Top-1000 候选的平均成对 Tanimoto 相似度:GFlowNet 0.44,PPO 0.62,MARS 0.59,数字越低越多样。多轮主动学习里,用代理拟合再查询真对接,GFlowNet 的 top-k 对接分数持续高于 MARS 和随机采集。
| 方法 | 指标 | 结果 |
| GFlowNet | 分数>8 的骨架数 | >1500 |
| MARS | 分数>8 的骨架数 | <100 |
| GFlowNet | Top-1000 Tanimoto | 0.44 |
| PPO | Top-1000 Tanimoto | 0.62 |
这是把能量函数摊成一次前向生成的一条路,不必每抽一个样本都跑一轮 MCMC。对要一次性提交大批候选、且代理本身不可靠的黑盒优化,覆盖多个峰比抠一个最高点更有用。分子结果说明,流量匹配在组合 DAG 上跑得通,不只是玩具网格。
它解决的是「按奖励比例采样」,不是「找到全局最优」。如果任务只关心单点最优,PPO 一类方法更直接。
作者自己写了:和 TD 方法一样,bootstrap 会带来优化困难,限制最终拟合质量。生成的是峰附近的区域,不是精确的局部极大,后续还需要局部搜索来精修。分子奖励只是对接代理,没有药代、毒性、可合成性。GFlowNet 匹配的是代理诱导的分布,代理错了,采样也会跟着偏。超网格实验规模很小(4 维边长 8),外推到更大组合空间时,流量估计的误差如何累积,论文没有给出保证。