Graph-GRPO: Training Graph Flow Models with Reinforcement Learning

TL;DR

Graph-GRPO introduces a reinforcement learning framework for training graph flow models with verifiable rewards, improving graph generation quality and efficiency, especially in molecular optimization tasks.

Contribution

It derives an analytical transition probability for GFMs and proposes a localized exploration strategy, enabling effective RL training and self-improvement.

Findings

Achieves 95.0% validity on synthetic datasets.

Attains state-of-the-art results in molecular optimization.

Uses only 50 denoising steps for high-quality generation.

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.