Reinforcement learning fine-tuning has proven effective for steering generative diffusion models toward desired properties in image and molecular domains. Graph diffusion models have similarly been applied to combinatorial structure generation, including neural architecture search (NAS). However, neural architectures are directed acyclic graphs (DAGs) where edge direction encodes functional semantics such as data flow-information that existing graph diffusion methods, designed for undirected structures, discard. We propose Directed Graph Policy Optimization (DGPO), which extends reinforcement learning fine-tuning of discrete graph diffusion models to DAGs via topological node ordering and positional encoding. Validated on NAS-Bench-101 and NAS-Bench-201, DGPO matches the benchmark optimum on all three NAS-Bench-201 tasks (91.61%, 73.49%, 46.77%). The central finding is that the model learns transferable structural priors: pretrained on only 7% of the search space, it generates near-oracle architectures after fine-tuning, within 0.32 percentage points of the full-data model and extrapolating 7.3 percentage points beyond its training ceiling. Bidirectional control experiments confirm genuine reward-driven steering, with inverse optimization reaching near random-chance accuracy (9.5%). These results demonstrate that reinforcement learning-steered discrete diffusion, once extended to handle directionality, provides a controllable generative framework for directed combinatorial structures.
Figures & tables
Fig. 1: Training dynamics of DGPO on (a) NB101 and (b) NB201 (CIFAR-10): validation accuracy (solid, left axis) and mean reward (dashed, right axis) over RL-FT epochs, with ±1σ bands across 3 seeds. Horizontal lines: random search and pretrained-only baselines. Both metrics converge reliably, confirming that RL fine-tuning steers the generation distribution toward higher-quality architectures.
Fig. 2: Distribution of generated architectures over RL-FT epochs on (a) NB101 and (b) NB201 (CIFAR-10). Each strip shows 300 sampled architectures (dots) at a given epoch, with mean accuracy (diamond) and top-5 architectures (stars). As training progresses (bottom to top), the mean and overall sample density shift markedly toward higher accuracy, demonstrating that RL fine-tuning reshapes the generative distribution rather than merely selecting isolated high performers. Single seed (42); n=300 samples per epoch.
Method
NB101
NB201 Val (%)
Val
C-10
C-100
IN-16
Optimum
95.06
91.61
73.49
46.77
Random Search
94.31
91.12
72.08
45.97
Reg. Evol.
94.47
–
–
–
Bayesian Opt.
94.57
91.54
73.26
46.43
BANANAS
94.73
91.56
73.49*
46.65
TABLE I: Comparison with SOTA NAS methods. Max validation accuracy (%) on NB101 and NB201. Prior results from [ 6 ] (mean over 10 runs); DGPO: mean ± std over 3 seeds (NB101) / deterministic API lookup (NB201). * = benchmark optimum. DGPO uses ∼ 2k queries; other methods use 150–192.
Method
NB101
NB201 (C-10)
Full Pretrain
94.50 ± 0.02
91.70 ± 0.02
Filtered Pretrain
94.18 ± 0.33
91.71 ± 0.02
TABLE II: RL-FT on full vs filtered pretraining: max validation accuracy (mean ± std, 3 seeds). Filtered models are pretrained exclusively on architectures below threshold T (NB101: T=0.87 , 7% of data; NB201: T=0.85 , 29%). Filtered pretraining loses only 0.32pp on NB101 and is within noise on NB201.
Fig. 3: OOD architecture discovery on (a, c) NB101 ( T=0.87 ) and (b, d) NB201 ( T=0.85 ). Top: distribution comparison - full pretrain (reference), filtered pretrain (epoch 0), and RL-FT (final epoch). Shaded region marks OOD architectures above T . Bottom: threshold crossing rate over RL-FT epochs. After pretraining on only sub-threshold architectures, RL-FT recovers above-threshold generation, demonstrating transferable structural priors. Single seed (42) for distributions; crossing rates are 3-seed aggregates.
Fig. 4: Bidirectional steering on (a) NB101 and (b) NB201 (CIFAR-10): forward (maximize, ↑ ) and inverse (minimize, ↓ ) DGPO trajectories over RL-FT epochs, with ±1σ bands (3 seeds). Dashed/dotted lines: expected max/min of random batches ( N=15 , bootstrap K=10,000 ). The inverse trajectory converges to near random-chance accuracy ( ∼ 9.5%), supporting reward-driven distribution steering.
Discrete graph generation has emerged as a powerful paradigm for modeling graph-structured data, yet state of the art models often rely on Graph Transformers or higher order architectures. We revisit this design assumption by introducing GenGNN, a modular message passing backbone for graph generation. GenGNN enables powerful generation by persisting edge fields through latent refinement of coupled node edge graph states, all without requiring global attention. Diffusion models integrating GenGNN achieve over 90 percent validity on standard benchmark datasets, performing within margins of Graph Transformer backbones and achieving up to 2x or even 5x faster inference. Systematic ablations isolate how GenGNN is resilient to oversmoothing during generative denoising, indicating each GenGNN component is necessary for downstream generation quality. Finally, representation-space analysis suggests GenGNN learns functionally similar representations to more theoretically-expressive architectures; even at deeper layers. As such, GenGNN uplifts local message-passing to challenge prevailing assumptions that performant discrete graph generation requires global attention or higher-order representations. Source Code Available Here
Jay Revolinsky, Harry Shomer, Jiliang Tang
Michigan State University East Lansing, Michigan, USA · University of Texas - Arlington Arlington, Texas, USA
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.
Baoheng Zhu, Deyu Bo, Delvin Ce Zhang +1
Beijing University of Posts and Telecommunications · National University of Singapore · University of Sheffield +1
Reinforcement learning can align visual generators with reward signals beyond their pretraining objectives, but existing score-function methods typically introduce stochasticity by perturbing the reverse denoising trajectory, leading to costly rollout-based optimization. We instead ask whether stochastic exploration can be introduced inside the generator while preserving its original sampler. Routing Policy Optimization (RPO) answers this question by turning attention routing into a stochastic exploration action. At a few selected denoising levels, each query block keeps its own block and samples further key blocks in proportion to the generator's native attention mass, or keeps dense attention. The sampled route perturbs the velocity of that transition, and a Gaussian likelihood around the dense velocity scores the perturbation with a score-function objective. RPO further uses a controlled two-rollout estimator: two rollouts share the same prompt, initial latent, routed denoising levels, and dense anchor layers, differing only in their sampled attention routes. Their reward difference provides an advantage signal for the routed transitions. Because the likelihood differentiates the dense velocity itself, the learned updates remain effective after routing is removed. At inference, RPO restores dense attention and retains the original deterministic sampler. Experiments on text-to-image and text-to-video generation show that RPO improves generation quality while achieving nearly 4x higher training efficiency than Flow-GRPO-style methods. Please see our project page at https://velox-world.com/rpo for visual results.