PaperlayerView exact revision on arXiv

Source-linked AI summary

Graph Forward Distribution Matching for Molecular Inverse Design

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

arXiv:2609.32056v1cs.LG

TL;DR

Molecular inverse design requires simultaneous control of target properties while preserving chemical validity, but reverse-policy reinforcement learning optimizes many coupled graph-diffusion decisions with a terminal reward. GRAPHFDM instead performs online optimization through forward distribution matching, jointly adjusting graph size and molecular structure. It achieves the strongest reported property control across molecular and polymer benchmarks, with high validity and generalization to unseen targets.

  • Problem

    Molecular inverse design needs precise simultaneous control of target properties while preserving chemical validity, whereas reverse-policy RL applies a single terminal reward across hundreds of coupled decisions.

    Source: §1 · p. 1

  • Method

    GRAPHFDM reweights valid generations into a reward-tilted target distribution and matches it through forward denoising, jointly optimizing graph size and structure.

    Source: §1 · p. 1, §3 · p. 3, §3.1 · p. 4

  • Results

    GRAPHFDM outperforms the strongest baseline on every property across molecular and polymer benchmarks, reduces MAE by up to 53.0%, maintains validity above 0.99, and generalizes to unseen property combinations.

    Source: §1 · p. 2, §6 · p. 10

  • Takeaways & Limitations

    Forward distribution matching is an effective alternative to reverse-policy optimization for controllable molecular inverse design across the reported molecular and polymer benchmarks.

    Source: §6 · p. 10

    GRAPHFDM operates within the structural support of its pretrained backbone, so support-expanding updates such as adaptive graph growth or vocabulary expansion remain future work.

    Source: §A.4 · p. 22

Abstract

from arXiv · show

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.

1 INTRODUCTION

GRAPHFDM shifts graph-diffusion optimization from reverse trajectories to forward distribution matching, addressing the challenge of controlling multiple molecular properties while preserving validity. It jointly adjusts graph size and molecular structure; across molecular and polymer benchmarks, it improves every target property while maintaining high validity.

1 INTRODUCTION

  • Existing reverse-policy RL uses a single terminal reward across hundreds of coupled decisions, and structural edits can shift multiple properties or invalidate a molecule.GDPO and GraphGRPO are reported to drive validity close to zero while providing limited property improvement.Source: §1 · p. 1
  • GRAPHFDM constructs a reward-tilted target by reweighting valid generations and matches it through denoising with the same sample weights.SDC reallocates probability across graph sizes while reward-weighted denoising refines structures at each size.Source: §1 · p. 1, Figure 1 · p. 2
  • A condition-independent fixed graph-size prior leaves an irreducible gap to the target distribution; SDC jointly matches graph-size and structure distributions under separate KL constraints.Molecular size influences property control and chemical validity.Source: §1 · p. 2
  • Across six-property molecular and four-property polymer benchmarks, GRAPHFDM outperforms the strongest baseline on every property, reduces MAE by up to 53.0%, and maintains chemical validity above 0.99.It also generalizes to unseen, out-of-distribution property combinations.Source: §1 · p. 2

2 PRELIMINARIES

The preliminaries represent a molecular graph by its active-unit count and structure, then describe diffusion denoising and reward-based post-training. Generation samples graph size from a fixed, condition-independent prior before denoising the structure.

2 PRELIMINARIES

  • A molecular graph is represented as G = (N, S), where N is the active-unit count and S contains node, edge, and attachment-position labels.Each active unit is an atom or motif, and N is represented by a fixed active-node mask during diffusion.Source: §2 · p. 2
  • The motif-aware discrete diffusion backbone corrupts the clean structure and uses a conditional denoiser to predict it from the noisy structure, mask, timestep, and property condition.The denoising loss is described as masked categorical cross-entropies for node, edge, and attachment-position labels.Source: §2 · pp. 2–3
  • Generation samples graph size from a fixed, condition-independent empirical prior, then reverse diffusion induces a clean-structure distribution conditioned on size and target properties.The active-node mask stays fixed throughout diffusion.Source: §2 · p. 3
  • Reverse-policy optimization treats denoising steps as a finite-horizon policy and propagates terminal molecular reward through updates to the reverse-transition probabilities.The same terminal feedback weights transitions along the sampled reverse trajectory.Source: §2 · p. 3

3 GRAPH FORWARD DISTRIBUTION MATCHING

GRAPHFDM turns valid terminal feedback into a weighted teacher and matches it through forward denoising. Structural Distribution Control (SDC) lets that target adjust graph size as well as molecular structure.

3 GRAPH FORWARD DISTRIBUTION MATCHING

  • GRAPHFDM re-corrupts clean graphs from a valid-conditioned reference and uses one density ratio to transfer teacher preferences to denoising and structural-controller objectives.Here, “forward” means constructing supervision by reapplying the fixed forward-corruption process.Source: §3 · p. 3, §3.1 · p. 4

3.1 FORWARD-DENOISING MATCHING

  • Under simultaneous realizability, population minimizers recover teacher-induced denoising posterior marginals at every noise level and the teacher's structural marginal.The denoiser and structural controller share parameters across target conditions.Source: §3.1 · p. 4

3.2 BEYOND FIXED-MASK DENOISING: STRUCTURAL-MARGINAL GAP

  • With fixed-mask denoising, the structural marginal stays fixed, so even an oracle conditional denoiser leaves mismatch equal to the KL divergence between teacher and fixed marginals.The best conditional distribution is the teacher conditional; the remaining minimum divergence is between the teacher's structural marginal and the fixed prior.Source: §3.2 · p. 4

3.3 CLOSING THE GAP: STRUCTURAL DISTRIBUTION CONTROL

  • Under finite supports and active constraints with positive duals, SDC's unique teacher tilts structures within each N and reallocates mass across N.Its two budgets separately constrain cross-N reallocation and conditional refinement within each N.Source: §3.3 · p. 5
  • Each online iteration samples graphs, retains valid candidates, evaluates terminal rewards, estimates SDC weights, and updates the denoiser and structural controller.The updated generator supplies candidates for the next iteration.Source: §3.3 · p. 6

4 EXPERIMENTS

Across polymer and QM9 benchmarks, GRAPHFDM leads in condition control while maintaining high validity. The experiments also test generalization to rare unseen targets and KL-budget allocation in structural distribution control.

4.1 EXPERIMENTAL SETUP

  • The study compares four-property polymers and six-property QM9 using fixed 6:2:2 splits, 10,000 generated examples, and a broad set of baselines.Evaluation covers validity, diversity, similarity, distribution distance, and conditional accuracy.Source: §4.1 · p. 6

4.2 RQ1: CONDITIONAL MOLECULAR GENERATION

  • Raw validity improves from 0.909 to 0.986 on Gas and 0.921 to 0.960 on QM9; GRAPHFDM also retains high fragment similarity and diversity without severe distributional drift.On polymers, fragment similarity is 0.923; on QM9, distribution distance is 0.615 and diversity is 0.917.Source: §4.2 · pp. 7–8
  • GRAPHFDM reduces polymer Avg. MAE to 0.509 versus CoMole (0.743), a 31.5% reduction, and achieves the lowest MAE on all four polymer and six QM9 properties.The reported comparison also places GRAPHFDM first in Overall Rank on QM9.Source: Table 1 · p. 7, Table 2 · p. 7, §4.2 · p. 8

4.3 RQ2: ABLATION STUDIES AND MODEL ANALYSIS

  • On 200 unseen QM9 targets, GRAPHFDM achieves the lowest MAE on every property while maintaining high raw validity; the targets represent eight rare joint signatures.Each signature contributes 25 targets, and SAScore is included as a sixth condition.Source: §4.3 · p. 8, Table 3 · p. 9
  • With the same total KL radius, 1:1 allocation lowers Avg. MAE from 0.590 to 0.509 (13.7%) and is selected for the main experiments.The 1:1 allocation also reduces Avg. Gas MAE from 0.725 to 0.618.Source: §4.3 · p. 9, Table 4 · p. 9

5 RELATED WORK

Prior work spans several approaches to molecular inverse design and reward-guided post-training. GRAPHFDM is distinguished by matching reward-tilted graph-size and molecular distributions for categorical graphs.

5 RELATED WORK

  • Molecular inverse design has used evolutionary search, autoregressive generation, latent-space optimization, and discrete graph diffusion.GraphDiT incorporates multiple property conditions into denoising, GrIDDD supports node insertion and deletion, and DeFoG uses discrete flow matching.Source: §5 · p. 9
  • GDPO and CoMole optimize reverse policies, while GraphGRPO applies group-relative policy optimization through analytic graph-flow transitions.CoMole operates in a motif-aware graph space.Source: §5 · p. 9
  • Diffusion alignment work covers discrete sequences, continuous visual generators, and reverse-trajectory matching; GRAPHFDM targets categorical graphs.It jointly matches the reward-tilted target's graph-size marginal and conditional molecular distributions under separate KL constraints.Source: §5 · p. 9

6 CONCLUSION

GRAPHFDM post-trains molecular graph diffusion by matching a reward-tilted target distribution through forward denoising. It combines theoretical guarantees with strong multi-property control, high validity, and generalization to unseen targets.

6 CONCLUSION

  • Across molecular and polymer benchmarks, GRAPHFDM achieves the strongest multi-property control with high chemical validity and generalizes to unseen targets from rare property combinations.The method reward-reweights valid terminal graphs and uses the same sample weights in the denoising objective.Source: §6 · p. 10
  • GRAPHFDM jointly optimizes the graph-size marginal and conditional molecular distributions at each size.Its target comes from reward reweighting valid terminal graphs, with the same sample weights used in denoising.Source: §6 · p. 10
  • The analysis derives a unique optimal teacher, proves condition-wise policy improvement, and establishes an irreducible matching gap under a fixed graph-size prior.Source: §6 · p. 10

A.1 ADDITIONAL DETAILS ON DATASET

The evaluation uses small-molecule and polymer benchmarks, with dataset roles summarized in Table 5. Pretraining uses molecular structures without target labels or conditional split identities.

A.1 ADDITIONAL DETAILS ON DATASET

  • The evaluation covers a six-property small-molecule benchmark and a four-property polymer benchmark.Source: §A.1 · p. 12
  • All compared methods use the same frozen splits and train-derived normalization statistics, and random subset construction and splitting use seed 42.Source: §A.1 · p. 12
  • Table 5 summarizes the datasets used for unconditional pretraining and conditional post-training.Pretraining uses molecular structures only and does not access target labels or conditional split identities.Source: Table 5 · p. 12, §A.1 · p. 12

A.1.1 SMALL-MOLECULE BENCHMARK

The small-molecule benchmark is built from a cleaned collection of neutral structures, with a sampled corpus for pretraining and a fixed conditional subset. It also defines rare compositional-OOD targets from combinations of property ranges.

A.1.1 SMALL-MOLECULE BENCHMARK

  • Starting from 133,885 rQM9 v1 entries, filtering problematic, unsanitizable, charged, and duplicate structures yields 132,382 unique neutral molecules.The removed counts are 303, 88, 742, and 370, respectively.Source: §A.1.1 · p. 12
  • Pretraining samples 100,000 structures from the cleaned collection, while the 10,000-molecule conditional benchmark is a fixed subset of that corpus.Property labels and conditional split identities are not used during pretraining.Source: §A.1.1 · p. 12
  • The OOD construction selects eight rare eligible joint property signatures, each occurring only two to five times in the 6,000-molecule conditional training split.It samples 25 targets per signature; selected values remain within their corresponding training-set marginal ranges, and SA is added as the sixth target property.Source: §A.1.1 · p. 12

A.1.2 POLYMER BENCHMARK

The polymer benchmark uses 553 repeat-unit records with jointly observed gas-permeability and synthetic-accessibility targets. Gas-property conditioning and evaluation use different scales.

A.1.2 POLYMER BENCHMARK

  • The benchmark contains 553 polymer repeat-unit records with jointly observed O2, N2, and CO2 permeability and an SA score.Wildcard atoms in repeat-unit SMILES mark polymerization sites.Source: §A.1.2 · p. 12
  • Gas-property errors are evaluated in log10 Barrer space, while model conditioning uses train-split z-score normalization.Permeabilities are stored in Barrer.Source: §A.1.2 · p. 12
  • The polymer backbone is initialized from CoMole’s 12,792-polymer unconditional pretraining corpus.Source: §A.1.2 · p. 12

A.1.3 TOKENIZER

The tokenizer uses domain-specific node-pair motif encodings, with a compact V300-R80 choice for rQM9 and CoMole’s V1000-R80 tokenizer reused for polymers. The supplied benchmark visuals distinguish target distributions and reporting scales.

A.1.3 TOKENIZER

  • The domain-specific motif tokenizers use CoMole’s node-pair encoding procedure, with V denoting motif vocabulary size and R the ring-token budget.Source: §A.1.3 · p. 12
  • For polymers, the tokenizer directly reuses CoMole’s V1000-R80 tokenizer learned from its pretraining corpus.Source: §A.1.3 · p. 12
  • Figure 3 separates target distributions by training, validation, and test split and plots polymer gas permeability in log10 Barrer; Table 6 reports gas-permeability statistics in raw Barrer.Source: Figure 3 · p. 13, Table 6 · p. 13
  • V300-R80 is selected for rQM9 as the smallest tokenizer at a compression plateau: raising V to 1000 reduces mean motif count by only 1.6% and leaves its 95th percentile unchanged, while adding 700 prediction classes.The ring budget of 80 consistently shortens motif sequences; tokenizer selection uses representation-level statistics, not conditional labels or downstream validation performance.Source: §A.1.3 · pp. 12–13

A.1.4 TARGET STATISTICS AND NORMALIZATION

This appendix supports GraphFDM’s theory and implementation by detailing experimental procedures and examining robustness across training seeds and learned property evaluators. It also compares post-training cost with CoMole and discusses limitations involving structural support, reward evaluation, and objective scope.

A.1.5 STRUCTURAL STATISTICS

  • Active motif-unit counts differ substantially from real-atom counts and form the structural variable sampled before denoising; target-condition normalization statistics are computed from the corresponding training split.The appendix reports online post-training wall-clock costs on one H100, including generation, oracle evaluation, optimization, and validation, but excluding base-model training and final evaluation.Source: §A.1.5 · p. 13, §A.1.4 · p. 13, Figure 4 · p. 14

A.2.1 VALIDITY CONDITIONING AND STRUCTURAL SUPPORT

  • Validity conditioning reweights the structural count distribution by count-dependent validity, and the original and conditioned marginals coincide only when that validity probability is constant over the reference support.The corollary proof establishes a strictly positive lower bound when the teacher differs from the reference distribution.Source: §A.2.1 · p. 14, §A.2.6 · p. 17

A.2.2 PROOF OF FORWARD DISTRIBUTION MATCHING

  • The appendix derives change-of-measure identities for denoising and structural log loss, and averaging over target conditions preserves both identities.The benchmark setup uses a normalized absolute-error reward and runs GraphFDM for 10 online rounds with K = 32 candidates per training condition.Source: §A.2.2 · p. 15, §A.3.1 · pp. 19–20

A.2.3 PROPER-SCORING INTERPRETATION

  • Under simultaneous realizability, every population minimizer of the denoising and structural objectives recovers the corresponding teacher-induced posterior marginals.The structural objective recovers µN,k(· | c) for each target condition with positive empirical mass.Source: §A.2.3 · p. 15

A.2.5 PROOF OF THEOREM 3.3

  • Theorem 3.3’s proof obtains a unique condition-wise maximizing joint distribution by applying Gibbs variational identities to structural and molecular factors.The proof connects this Lagrangian maximizer to the constrained problem using Slater’s condition and KKT conditions, assuming positive optimal dual variables and positive KL budgets.Source: §A.2.5 · p. 16

A.2.7 FINITE-SAMPLE STRUCTURED PROJECTION

  • The finite-sample projection shares temperatures across conditions to enforce the condition-averaged KL budgets, and its estimates converge almost surely under the stated support and integrability assumptions.A condition with no valid candidates contributes no update in that iteration.Source: §A.2.7 · p. 19

A.3.2 ROBUSTNESS ACROSS SEEDS

  • GRAPHFDM achieves the best overall rank with low variance across three training seeds.The evaluation fixes the data split, property oracle, and test targets across runs; post-training methods reuse the same reference checkpoint.Source: §A.3.2 · p. 20, Table 9 · p. 20

A.3.3 ROBUSTNESS ACROSS ORACLES

  • GRAPHFDM ranks first under all three evaluators, while GRAPHFDM ≻ CoMole ≻ GraphDiT ≻ LSTM ≻ MELD remains the consistent partial ordering.The same generated molecules are rescored with SVR and GP oracles; RF is used as the reward oracle and primary evaluator because it has the lowest fitting MAE.Source: §A.3.3 · p. 21, Table 11 · p. 21

A.3.4 COMPUTATIONAL COST

  • GRAPHFDM requires 16.65 vs. 30.37 hours on polymers and 49.18 vs. 50.30 hours on molecules relative to CoMole under matched SFT initialization and backbone.Endpoint-based updates reuse clean samples through forward re-noising without retaining reverse trajectories. Oracle evaluation takes 1.12 ms per polymer and 0.79 ms per QM9 molecule, approximately 12 and 152 seconds per GraphFDM round, respectively, a small fraction of total online cost; VIDD and GDPO have lower costs under their native configurations.Source: §A.3.4 · p. 21, Table 12 · p. 22

A.4 LIMITATIONS AND FUTURE WORK

  • GRAPHFDM operates within its pretrained backbone’s structural support, reallocating probability across represented graph sizes rather than expanding that support.Adaptive graph growth, vocabulary expansion, broader chemical spaces, and other structured domains are proposed directions for future work.Source: §A.4 · p. 22
  • Current benchmarks rely on learned property oracles and a normalized multi-property reward, leaving integration of costly high-fidelity feedback and broader objective types for future work.The authors suggest sample-efficient physics-based simulation or experimental feedback, as well as task-specific preferences, constraints, or Pareto-aware objectives.Source: §A.4 · p. 22

Related papers

Loading 2609.32056v1…