Graph-GRPO: Training Graph Flow Models with Reinforcement Learning
Abstract
Graph generation is a fundamental task with broad applications, such as drug discovery. Recently, discrete flow matching-based graph generation, a.k.a., 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. Code is available in https://github.com/Zhubaoheng/Graph-GRPO.
Lay Summary
How can we get AI models to generate graph structures, such as molecules or networks, that meet a specific goal we have in mind? This question is becoming increasingly important. For example, we might want to generate a drug molecule that binds to a particular protein. But the strongest graph generators available today, known as graph flow models, are very hard to steer through feedback-based training (reinforcement learning), and have so far resisted being reliably aligned to concrete goals. We found that the obstacle lies in a random sampling step inside graph flow models. This randomness severs the path along which feedback signals would normally flow back into the model, leaving it with nothing to learn from. We derived a new mathematical formula that replaces this random step with a precise, computable expression, allowing the feedback signal to travel all the way back. We also propose a refinement strategy: instead of generating each candidate from scratch, we keep the ones that already look promising, perturb a small part of them, and regenerate, much like revising a draft rather than rewriting it every time. We call this method Graph-GRPO. It achieves new state-of-the-art results across graph and molecular design benchmarks, offering a principled and general recipe for aligning graph generators.