Graph Forward Distribution Matching for Molecular Inverse Design

Yihan Zhu; Yuhan Liu; Brett Savoie; Tengfei Luo; Meng Jiang

cs.LGSep 25, 2026

Graph Forward Distribution Matching for Molecular Inverse Design

Authors: Yihan Zhu, Yuhan Liu, Brett Savoie, Tengfei Luo, Meng Jiang

Organizations: University of Notre Dame

Abstract

Achieving precise control over multiple properties without sacrificing chemical validity remains a central challenge in molecular inverse design. Existing reinforcement learning (RL) methods fine-tune graph diffusion models by treating reverse sampling as a sequential policy, using a single terminal reward to optimize hundreds of coupled decisions. They often suffer from instability, validity collapse, and limited property gains. We introduce GraphFDM (Graph Forward Distribution Matching), a new online RL paradigm for graph diffusion that performs optimization through the forward process. GraphFDM uses valid generations to define a reward-tilted target distribution jointly optimized over graph size and molecular structure for each property condition, incorporating reinforcement signals into supervised learning without storing reverse trajectories. We derive the unique optimal target, prove a condition-wise improvement guarantee, and show that the fixed graph-size prior of standard graph diffusion leaves an irreducible matching gap. In multi-conditional polymer and small-molecule generation, GraphFDM achieves the lowest MAE on every target property, with reductions of up to 53.0% relative to the strongest baselines and chemical validity above 0.99. It further generalizes to out-of-distribution property combinations.

Figures & tables

Figure 1: Overview of GraphFDM . (a) Reverse-policy RL (e.g., GRPO) assigns terminal rewards across coupled reverse transitions which can reinforce structural errors. (b) At each online round, GraphFDM constructs a reward-tilted teacher $\pi^{*}(G\mid c)$ over valid graphs $G=(N,S)$ . SDC reallocates probability across the number of active nodes $N$ while reward-weighted denoising refines molecular structures $S$ at each $N$ .

Model	Validity	Distribution Learning			Condition Control
Model	Valid. (raw) $\uparrow$	Div. $\uparrow$	Sim. $\uparrow$	Dis. $\downarrow$	Synth. $\downarrow$	O 2 Perm $\downarrow$	N 2 Perm $\downarrow$	CO 2 Perm $\downarrow$	Avg. MAE $\downarrow$
GraphGA	1.000 (N.A.)	0.883	0.927	9.188	1.331	1.984	2.290	1.949	1.888
LSTM-HC	0.991 (N.A.)	0.894	0.932	12.122	1.413	1.381	1.641	1.356	1.448
MARS	1.000 (N.A.)	0.838	0.928	7.562	1.166	1.576	1.833	1.607	1.546
DiGress	0.839 (0.117)	0.918	0.515	20.256	2.104	1.543	1.861	1.477	1.746
GDSS	1.000 (0.003)	0.857	0.004	37.627	1.164	1.260	1.477	1.285	1.297

Table 1: Multi-Conditional Generation of 10K Polymers: Results on four properties (synthetic score, gas permeability for O 2 , N 2 , CO 2 ). MAE is computed between the input conditions and evaluated properties of the generated polymers. Best results are in red .

Model	Validity	Distribution Learning			Condition Control
Model	Valid. (raw) $\uparrow$	Div. $\uparrow$	Sim. $\uparrow$	Dis. $\downarrow$	Synth. $\downarrow$	$\mu$ $\downarrow$	$\alpha$ $\downarrow$	$\varepsilon_{\mathrm{HOMO}}$ $\downarrow$	$\varepsilon_{\mathrm{LUMO}}$ $\downarrow$	$C_{v}$ $\downarrow$	Overall Rank $\downarrow$
GraphGA	1.000 (N.A.)	0.927	0.922	0.777	0.946	0.702	4.903	0.019	0.020	2.121	8
LSTM-HC	0.990 (N.A.)	0.916	0.957	1.400	1.016	1.478	8.489	0.022	0.052	4.300	16
MARS	1.000 (N.A.)	0.931	0.770	5.731	1.037	0.844	7.597	0.014	0.021	2.961	10
DiGress	0.915 (0.849)	0.918	0.941	4.696	1.612	0.912	3.897	0.012	0.018	1.900	7
GDSS	1.000 (0.737)	0.908	0.733	6.559	1.340	1.804	10.371	0.021	0.052	3.586	16

Table 2: Multi-Conditional Generation of 10K QM9 Molecules: Results on six properties (synthetic score and five QM9 properties). MAE is computed between the input conditions and evaluated properties of the generated molecules. Overall Rank aggregates property-wise MAE ranks across units and scales. Best results are in red .

Figure 2: Out-of-distribution multi-property alignment. Joint t-SNE visualization of normalized target vectors (gray stars) and per-target centroids of valid generated molecules (colored circles). GraphFDM more consistently follows the target modes than competing post-training methods.

Model	Validity	Distribution Learning			Condition Control
Model	Valid. (raw) $\uparrow$	Div. $\uparrow$	Sim. $\uparrow$	Dis. $\downarrow$	Synth. $\downarrow$	$\mu$ $\downarrow$	$\alpha$ $\downarrow$	$\varepsilon_{\mathrm{HOMO}}$ $\downarrow$	$\varepsilon_{\mathrm{LUMO}}$ $\downarrow$	$C_{v}$ $\downarrow$
VIDD	1.000 (0.535)	0.923	0.664	10.985	2.002	1.262	10.708	0.032	0.059	4.426
TreeDiff	0.944 (0.038)	0.934	0.687	5.488	0.996	1.229	9.879	0.021	0.031	3.372
GDPO	0.804 (0.717)	0.931	0.829	3.496	0.574	1.227	6.875	0.019	0.030	2.438
GraphGRPO	0.955 (0.705)	0.927	0.787	9.873	0.692	1.163	7.127	0.029	0.026	2.681
CoMole	0.875 (0.716)	0.926	0.744	3.207	0.642	1.210	6.688	0.019	0.030	2.792

Table 3: Out-of-distribution Multi-Conditional Generation of 10K QM9 Molecules: We evaluate post-training methods on 200 unseen targets with rare property combinations. Best results are in red .

Constraint	$\mathcal{C}_{N}:\mathcal{C}_{S}$	Validity	Distribution		Condition Control
Constraint	$\mathcal{C}_{N}:\mathcal{C}_{S}$	Valid. (raw) $\uparrow$	Div. $\uparrow$	Dis. $\downarrow$	Synth. $\downarrow$	Avg. Gas $\downarrow$	Avg. MAE $\downarrow$
Joint KL	Adaptive ( $\approx 2.5{:}1$ )	0.988 (0.979)	0.834	8.764	0.184	0.725	0.590
Separate 1:1	$1{:}1$	0.991 (0.986)	0.834	9.642	0.183	0.618	0.509
Separate 1:2	$1{:}2$	0.986 (0.965)	0.840	9.418	0.189	0.648	0.534

Table 4: KL budget allocation in SDC:

\mathcal{C}_{N}:\mathcal{C}_{S}

denotes the ratio of structural-marginal to conditional-structure KL change. Avg. Gas averages the three gas-permeability MAEs. Best results are in bold .

Appendix figures & tables10 assets

Supplementary material from the paper’s appendix.

Appendix

Benchmark	Unlabeled PT	Conditional Training	Train/Val/Test	Motif tokenizer
Small molecule	100,000	10,000	6,000/2,000/2,000	V300-R80
Polymer	12,792	553	337/107/109	V1000-R80

Appendix

Table 5: Overview of the datasets used by GraphFDM . PT uses molecular structures only and does not access target labels or conditional split identities.

Benchmark	Target	Unit	Full range	Train mean/std
Small molecule	SA	-	1.111-7.750	4.236 / 0.924
	$\mu$	D	0-9.071	2.665 / 1.433
	$\alpha$	$a_{0}^{3}$	13.21-127.46	75.274 / 8.226
	HOMO	Ha	$-0.3963$ - $-0.1538$	$-0.2404$ / 0.0219
	LUMO	Ha	$-0.1466$ -0.1171	0.0128 / 0.0469
	$C_{v}$	$\mathrm{cal\,mol^{-1}\,K^{-1}}$	6.469-46.306	31.663 / 4.118

Appendix

Table 6: Target ranges and train-split normalization statistics. Gas-permeability statistics are computed in the raw Barrer scale.

Figure 3: Target distributions for the small-molecule and polymer benchmarks, shown separately for the training, validation, and test splits. Polymer gas permeabilities are visualized in

\log_{10}

Barrer, consistent with the evaluation metric.

Benchmark	Real atoms	Hetero atoms	Rings	Mol. weight	Active motif units
Small molecule	8.80/9	2.44/7	1.74/8	122.72/152.04	4.84/9
Polymer	27.97/48	5.28/17	3.68/11	391.78/866.20	6.66/46

Appendix

Table 7: Structural statistics of the conditional benchmarks. Values are reported as mean/maximum.

Figure 4: Distributions of real-atom counts and active motif-unit counts in the two conditional benchmarks. Polymer wildcard atoms are excluded from the real-atom count. The active motif-unit count corresponds to the structural variable

N

used by GraphFDM .

Configuration	Molecule	Polymer
PT epochs	200	400
PT batch size	256	64
SFT epochs	200	1200
SFT batch size	256	128
Online rounds	10	10
Candidates per condition	32	32

Appendix

Table 8: Dataset-specific training configurations.

Model	Validity	Distribution Learning
	Valid. (raw) $\uparrow$	Div. $\uparrow$	Sim. $\uparrow$	Dis. $\downarrow$
DiGress	0.914 $\pm$ 0.004 (0.827 $\pm$ 0.020)	0.919 $\pm$ 0.002	0.943 $\pm$ 0.002	4.742 $\pm$ 0.142
CoMole	0.937 $\pm$ 0.037 (0.831 $\pm$ 0.067)	0.916 $\pm$ 0.001	0.930 $\pm$ 0.010	4.253 $\pm$ 0.019
GDPO	0.914 $\pm$ 0.003 (0.849 $\pm$ 0.010)	0.917 $\pm$ 0.001	0.950 $\pm$ 0.003	0.795 $\pm$ 0.087
GraphGRPO	0.973 $\pm$ 0.003 (0.836 $\pm$ 0.015)	0.910 $\pm$ 0.001	0.914 $\pm$ 0.016	2.953 $\pm$ 0.158
GraphFDM	0.994 $\pm$ 0.000 (0.959 $\pm$ 0.001)	0.916 $\pm$ 0.001	0.907 $\pm$ 0.007	0.612 $\pm$ 0.004

Appendix

Table 9: Multi-seed results of the top-5 models on QM9 multi-conditional generation. Results on six properties (synthetic score and five QM9 properties) are reported as mean

\pm

std. Overall Rank aggregates property-wise MAE ranks across different units and scales. Best results are in red .

Oracle	Metric	$O_{2}$	$N_{2}$	$CO_{2}$	$\mu$	$\alpha$	$\varepsilon_{\mathrm{HOMO}}$	$\varepsilon_{\mathrm{LUMO}}$	$C_{v}$
RF	MAE $\downarrow$	0.412	0.456	0.421	0.194	1.343	0.002	0.002	0.540
SVR	MAE $\downarrow$	0.480	0.539	0.488	0.862	4.476	0.010	0.019	2.056
GP	MAE $\downarrow$	0.872	0.989	0.874	1.153	6.304	0.015	0.039	3.148

Appendix

Table 10: Training Performance of Oracle Methods: we train the models on all polymers or small molecules in a task to simulate the Oracle. Results from the random forest model are in bold because it has the lowest training MAE.

Rank	Random Forest	Support Vector Regression	Gaussian Process
1	GraphFDM	GraphFDM	GraphFDM
2	CoMole	CoMole	CoMole
3	GraphDiT	GDSS	GDSS
4	CSGD	GDPO	GDPO
5	GDSS	GraphDiT	GraphDiT
6	DeFoG	LSTM	LSTM

Appendix

Table 11: Oracles for Generation Evaluation: Models are ranked from 1 to 9 across the polymer benchmark by Avg.MAE under each evaluator. Underlined entries preserve the same relative ordering across all evaluators.

Polymer
Method	Online unit	Evaluations/unit	Seconds/unit	ms/evaluation	Total hours
GDPO	reverse trajectory	500	1.312	2.625	12.28
GraphGRPO	group trajectory	500	1.940	3.884	19.30
VIDD	score trajectory	2,400	9.291	3.871	3.55
CoMole	reverse trajectory	500	0.291	0.582	30.37
GraphFDM	clean endpoint	500	0.556	1.112	16.65

Appendix

Table 12: Online post-training and inference-time search cost. We report wall-clock cost on one H100, including sampling, oracle evaluation, optimization, and validation. Evaluations denote the dominant model calls for each method’s generative dynamics. CoMole and GraphFDM use the same SFT initialization and backbone.

Explore similar work

Controllable Molecular Generative Foundation Models
May 14, 2026Yihan Zhu, Yuhan Liu, Weijiang Li +2Molecular Optimization Molecular Graph Generation
University of Notre Dame
Fine-Tuning Diffusion Models for Molecular Generation via Reinforcement Learning and Fast Sampling
May 31, 2026Guang Lin, Shikui Tu, Lei XuMolecular Optimization Diffusion Model Fine-Tuning
Department of Computer Science and Engineering, Shanghai Jiao Tong University, Shanghai, China · Guangdong Laboratory of Artificial Intelligence and Digital Economy (SZ), Shenzhen, Guangdong, China
MolHIT: Advancing Molecular-Graph Generation with Hierarchical Discrete Diffusion Models
Feb 19, 2026Hojung Jung, Rodrigo Hormazabal, Jaehyeong Jo +5Molecular Graph Generation Diffusion Models
KAIST AI · LG AI Research · Seoul National University

May 14, 2026J/K move · Enter open · S save

Controllable Molecular Generative Foundation Models

Yihan Zhu, Yuhan Liu, Weijiang Li, Tengfei Luo+1

University of Notre Dame

Molecular Optimization Molecular Graph Generation Graph Foundation Models RL for Diffusion Models Molecular Generation Conditional Generative Modeling

Open paper AskPDF arXiv

Despite the success of foundation models in language and vision, molecular graph generation still lacks a unified framework for heterogeneous design tasks with reliable controllability. While reinforcement learning (RL) offers a natural post-training mechanism for task-specific optimization, applying it to graph generative models is hindered by the vast atom-wise action spaces and chemically invalid intermediate states. We propose Controllable Molecular Generative Foundation Models (CoMole), built with a unified motif-aware graph diffusion pipeline. By learning a motif-aware graph space, CoMole transfers pretrained structural priors into controllable generation, where RL optimizes conditional reverse policies over chemically meaningful decisions. We theoretically characterize the bottleneck of atom-level RL and justify motif-aware policy optimization. Across three heterogeneous benchmarks spanning materials and drug discovery, CoMole ranks first in controllability on all nine targets, reduces MAE by up to 48.2% relative to the strongest baselines, and maintains validity above 0.94 without rule-based correction or post-hoc filtering. We further show that CoMole transfers controllability to unseen properties by optimizing only task embeddings with the generator frozen, achieving performance competitive with strong task-specific baselines.

Abstract

Explore similar work

Controllable Molecular Generative Foundation Models

Fine-Tuning Diffusion Models for Molecular Generation via Reinforcement Learning and Fast Sampling

MolHIT: Advancing Molecular-Graph Generation with Hierarchical Discrete Diffusion Models