Graph-GRPO: Training Graph Flow Models with Reinforcement Learning

TL;DR

This paper proposes Graph-GRPO, an online reinforcement learning (RL) framework for training GFMs under verifiable rewards that achieves state-of-the-art performance on the molecular optimization tasks, outperforming graph-based and fragment-based RL methods as well as classic genetic algorithms.

Abstract

Graph generation is a fundamental task with broad applications, such as drug discovery. Recently, discrete flow matching-based graph generation, \aka, graph flow model (GFM), has emerged due to its superior performance and flexible sampling. However, effectively aligning GFMs with complex human preferences or task-specific objectives remains a significant challenge. In this paper, we propose Graph-GRPO, an online reinforcement learning (RL) framework for training GFMs under verifiable rewards. Our method makes two key contributions: (1) We derive an analytical expression for the transition probability of GFMs, replacing the Monte Carlo sampling and enabling fully differentiable rollouts for RL training; (2) We propose a refinement strategy that randomly perturbs specific nodes and edges in a graph, and regenerates them, allowing for localized exploration and self-improvement of generation quality. Extensive experiments on both synthetic and real datasets demonstrate the effectiveness of Graph-GRPO. With only 50 denoising steps, our method achieves 95.0\% and 97.5\% Valid-Unique-Novelty scores on the planar and tree datasets, respectively. Moreover, Graph-GRPO achieves state-of-the-art performance on the molecular optimization tasks, outperforming graph-based and fragment-based RL methods as well as classic genetic algorithms.

Figures (9)

• Figure 1: Reward curves on two molecular optimization tasks. We use DeFoG DeFoG as the base model. As the number of oracle calls increases, the score of RL-optimized models gradually rises, while the base model remains almost unchanged. In the tasks with highly selective reward, e.g., Valsartan SMARTS, refining promising candidates is more effective than de novo generation.
• Figure 2: The overall framework of Graph-GRPO. (1) Rollout: given a noisy graph, the policy model samples $K$ denoising trajectories and caches the graph state $z_{t+\Delta t}$ along with its transition probability $p_{t+\Delta t}^{\text{old}}$. In the meantime, we normalize the reward scores to calculate the advantage of each trajectory. (2) RL training: we select a graph $G_t$ from the rollouts and use the new policy model to estimate the transition probabilities $p_{t+\Delta t}$. Graph-GRPO maximizes the rewards by optimizing the advantage-weighted ratio in Eq. \ref{['eq: grpo']}.
• Figure 3: Refinement in Graph-GRPO. We first use GFMs to denoise a noise graph from $t=0$ to $t=1$. Subsequently, we re-noise the generated graph to time step $t_{\epsilon}$ and denoise it again using GFMs. This strategy explicitly increases the denoising steps of GFMs and improves the generation quality of Graph-GRPO.
• Figure 4: Visualization of the sampling trajectory. From left to right, the graphs $G_t$ move from a prior distribution ($t=0$) to the data distribution ($t=1$). Visualizations of the predicted clean state $\hat{z}_1$ at the corresponding time steps are provided in Figure \ref{['fig:vis_pred_z1']}.
• Figure 5: Visualization of training and evaluation curves in Graph-GRPO. We select five representative tasks from the PMO benchmark. (a) Training reward curves plotted against training steps. (b) Average top-10 scores plotted against the number of oracle calls. For brevity, we only visualize the first 5,000 oracle calls.

Theorems & Definitions (2)

• Proposition 3.1
• proof