graph-grpo：使用强化学习训练图流模型

作者：Baoheng Zhu（北京邮电大学）、Deyu Bo（新加坡国立大学）、Delvin Ce Zhang（谢菲尔德大学）、Xiao Wang（北京航空航天大学，通讯作者） 发表信息：Proceedings of the 43rd International Conference on Machine Learning, Seoul, South Korea. PMLR 306, 2026.

1. 研究背景与目标

图生成（graph generation）是机器学习领域的一项基础任务，在药物发现、材料设计等应用中具有广泛价值。近年来，基于离散流匹配（discrete flow matching）的图生成方法——即图流模型（graph flow model，GFM）——由于性能优越且采样灵活，逐渐受到关注。这类方法通过解耦训练目标与采样过程，实现了更有效的生成建模和更灵活的推理。然而，如何将图流模型与复杂的人类偏好或任务特定目标有效对齐，仍然是一个关键挑战。以药物发现为例，理想的小分子需要同时具备高结合亲和力与低毒性，但生成此类分子在庞大的生成空间中计算代价极高。

近期，在线强化学习（online reinforcement learning，RL）在通过最大化特定奖励函数来对齐生成模型与人类先验方面展现出显著潜力。已有研究探索了RL与图生成对抗网络、图扩散模型的结合，但将RL扩展至图流模型面临两个根本性困难。第一，现代RL算法依赖策略梯度（policy gradient），要求策略模型对每个动作的转移概率可微分，然而现有图流模型通过蒙特卡洛采样（Monte Carlo sampling）估计动作概率，切断了梯度流。第二，RL探索需要充分的反馈信号来定位任务特定区域，但图流模型通常采用从头生成（de novo generation）方式，容易产生稀疏奖励信号，即大多数生成的图无效，导致探索效率低下。

针对上述问题，本文提出graph-grpo，一个在线强化学习框架，采用组相对策略优化（group relative policy optimization，GRPO）将图流模型与任务特定目标对齐。该方法的两个核心贡献是：推导出图流模型转移概率的解析表达式，从而替代蒙特卡洛采样，实现完全可微分的rollout；以及提出一种精炼策略（refinement strategy），通过随机扰动图中的特定节点和边并进行再生成，实现局部化探索和生成质量的自改进。

1. 研究方法与工作流程

2.1 解析速率矩阵的推导

图流模型的去噪过程被建模为连续时间马尔可夫链（continuous-time Markov chain，CTMC）。在每个去噪步骤中，需要估计速率矩阵（rate matrix）r_t，将图从状态gt变换到g{t+dt}。现有方法（如defog）通过蒙特卡洛采样获得伪图作为真实数据的替代，进而计算条件速率矩阵。但这种方式存在两个问题：转移概率不可微分，导致无法进行RL优化；训练与推理不一致，新旧模型可能采样到不同的伪图。

为解决这些问题，本文在命题3.1中给出了速率矩阵的解析表达式：

r^θ_t(zt, z{t+dt}) = pθ(z{t+dt})·v_1 + (1 − p_θ(z_t) − pθ(z{t+dt}))·v_2

其中v_1和v_2是两个可在生成前预先计算的统计量。该解析表达式直接由去噪器（denoiser）的预测概率计算得到，完全可微分，并且保证了训练与推理过程的一致性。这一创新使得图流模型能够通过策略梯度进行端到端的RL训练。

2.2 graph-grpo的训练框架

graph-grpo的训练流程包括两个主要阶段：rollout收集和RL训练。在rollout收集阶段，给定一个从先验分布p_0采样的噪声图g_0，策略模型π_θ并行生成k条独立的去噪轨迹{τ^{(k)}}，每条轨迹记录采样的图状态序列以及每一步的转移概率。到达终止状态t=1后，使用任务特定的奖励函数r(g)对最终图进行评估，获得奖励r^{(k)}。RL训练阶段的目标是最大化期望奖励，采用GRPO的目标函数，包含策略优化损失和KL散度惩罚项。策略优化使用组相对优势（group-relative advantage）和重要性采样比率（importance sampling ratio），KL散度用于防止RL优化后的模型偏离基础模型过远。

2.3 精炼策略

标准采样流程中，图流模型通过从噪声图g_0逐步去噪至干净图g_1来生成新图，属于从头生成。但现实应用往往只需要具有特定性质的图，这些图仅存在于整个生成空间的小区域内。因此，从头生成会产生大量无效或低质量图，无法有效探索高潜力区域。为解决该问题，graph-grpo提出精炼策略：维护一个优先级池B，记录奖励分数最高的m个图。每次迭代中，这些候选图经历“重新加噪”和“再生成”两个步骤。重新加噪阶段，根据条件概率路径将候选图g_1回退至中间噪声状态t_ϵ，其中t_ϵ控制扰动幅度——t_ϵ越大意味着噪声幅度越小。实际中从一个候选图采样多个噪声版本来探索其潜力。再生成阶段，对每个噪声图重新执行去噪过程生成新的候选图，评估奖励后与池中已有图比较，保留分数最高的m个图。通过反复迭代，graph-grpo逐步聚焦于生成空间中的高潜力区域，提升生成质量。此外，研究还引入了动态先验（dynamic prior）机制，从高奖励缓冲区中非参数化地估计节点类型、边类型和图规模的分布，并在采样阶段使用这些自适应的先验分布替代原始的静态先验。

2.4 实验设置

实验框架构建于defog基础模型之上，采用配备了相对随机游走概率（relative random walk probabilities，RRWP）结构编码的图Transformer作为去噪网络。针对合成任务（planar/tree）使用10层模型且不使用RRWP，针对分子任务（蛋白质对接/PMO）使用12层模型且RRWP步数为20。优化器采用AdamW，学习率初始为2×10^{-5}，奖励平台期衰减至最低1×10^{-5}，有效批大小为200。GRPO的组大小k=60，裁剪参数ε_low=0.2、ε_high=0.28，优势裁剪至[−5, 5]，KL惩罚系数β=0.005。精炼策略中，首先使用300次oracle调用进行从头生成以初始化候选池，然后对池中每个分子生成150个变体直至消耗2,000次oracle调用，之后将探索强度提高至每个分子500个变体直至10,000次oracle调用。动态先验的更新动量α=0.05，维护容量为1,000的缓冲区，当奖励改进超过0.001时触发更新。精炼的噪声水平固定为t_ϵ=0.8，候选池大小m=5。

1. 主要实验结果

3.1 通用图生成

在planar和tree两个合成数据集上评估通用图生成能力。两个数据集均包含64个节点的图，训练集128个图，验证集32个，测试集40个。评估指标为有效-唯一-新颖（valid-unique-novel，V.U.N.）分数和与训练集的结构比率。奖励函数设计为硬有效性约束与软分布匹配度量的组合。结果显示，在planar数据集上，graph-grpo的比率为1.5，低于基础模型defog的3.2，同时V.U.N.分数保持95.0%不变，说明在不降低有效性的前提下实现了更好的结构对齐。在tree数据集上改进更为显著：V.U.N.从defog的73.5%跃升至97.5%，同时比率从2.5降至2.2。值得注意的是，graph-grpo仅用50步去噪步骤，就超越了使用1,000步的图扩散模型（如digress、gbd、disco）以及图扩散策略优化方法gdpo。

3.2 蛋白质对接

蛋白质对接任务使用zinc250k数据集，评估生成与特定蛋白质对接的分子的能力，涉及五个靶蛋白：parp1、fa7、5ht1b、braf和jak2。指标包括命中率（hit ratio）和前5%对接分数（docking score，DS）。奖励函数由QED、合成可及性（SA）、新颖性和对接分数的加权组合构成。结果显示，graph-grpo在五个靶蛋白上均取得了最优或次优的对接分数。在采样效率方面表现尤为突出：以parp1为例，graph-grpo生成分子的命中率达60.763%，是最佳基线gdpo（9.814%）的6倍以上。在jak2靶点上，命中率为52.897%，同样大幅领先于第二名的gdpo（13.405%）。该结果表明graph-grpo能够高效探索化学空间中的高潜力区域，并且较高的命中率也有利于对接分数，因为有更多有效候选分子参与评估。

3.3 目标性质优化

使用实用分子优化（practical molecular optimization，PMO）基准测试，包含23个多样化任务，在10,000次oracle调用预算下进行严格评估。设置了两种实验条件：预筛选（prescreening）和冷启动（cold-start）。在冷启动设置中，graph-grpo的总AUC-top10为18.987，显著优于最佳基线invirtuogen（16.676）和genmol（16.213）。在困难任务上优势尤为明显，如thiothixene rediscovery（graph-grpo为0.842，最佳基线为0.692）和troglitazone rediscovery（0.711对0.867，但整体而言graph-grpo表现更加均衡）。在预筛选设置中，graph-grpo进一步将AUC-top10提升至19.270，创造了新的最优结果。值得强调的是，graph-grpo的基础模型仅在zinc250k上预训练，而部分基线方法（如invirtuogen和genmol）使用了包含超过十亿分子的更大数据集进行预训练。

3.4 消融实验

对graph-grpo的四个核心组件进行了消融研究：RL训练、精炼、预筛选和动态先验。基础模型defog未使用任何组件时AUC-top10仅为11.079。单独添加精炼策略将其提升至15.251。仅使用RL训练（不含精炼和动态先验）可达到16.901，加入动态先验后提升至17.450。RL训练与精炼结合达到17.950，三者结合（RL+精炼+动态先验）达到18.987。四个组件全部启用时达到最优的19.270。结果表明RL训练是最主要的贡献因素，精炼策略在有和无RL的情况下均能带来显著增益。

3.5 额外分析

对精炼策略中重新加噪时间t_ϵ的分析表明，较大的t_ϵ值（如0.7或0.9）获得最佳结果，因为较小的噪声扰动允许更受控的探索。过小的t_ϵ会破坏分子结构导致性能退化，但仍优于从头生成（t_ϵ=0.0）。对不含解析转移概率的reinforce变体的对比实验显示，reinforce因训练目标与实际采样过程不一致而表现不稳定，奖励从1.03降至0.312，而graph-grpo将奖励提升至1.537，梯度裁剪率仅0.7%。

1. 结论与研究价值

graph-grpo提出了一个用于优化图流模型的在线强化学习框架。通过推导解析转移概率，该方法实现了完全可微分的rollout，使图流模型能够与现代RL框架无缝集成。精炼策略通过受控扰动和再生成实现局部化探索，有效避免了无效或低质量图的生成。实验结果表明，graph-grpo在通用图生成、蛋白质对接和目标性质优化三个任务上均实现了最先进的性能，超越了基于图和基于片段的RL方法以及经典遗传算法。该研究为将图流模型与复杂下游目标对齐提供了一个原理性的路径。未来工作的一个有前景的方向是将graph-grpo应用于更广泛的下游应用，如材料生成。

本研究的核心亮点包括：第一，首次推导出图流模型速率矩阵的解析表达式，解决了蒙特卡洛采样带来的不可微分问题，使端到端RL训练成为可能；第二，提出迭代精炼策略，通过局部探索实现生成质量的自改进，在复杂优化场景中相比从头生成展现出明显的性能优势；第三，在仅使用50步去噪步骤的条件下，超越了使用1,000步的图扩散模型，展现出极高的采样效率；第四，在冷启动设置下（无任何预筛选），graph-grpo即可匹敌甚至超越依赖昂贵预筛选的基线方法，说明其能够独立地探索高潜力化学区域。