Source-linked AI summary
Embedded Graph Flows for Categorical Graph Generation
Ethan Ma, Zihan Wang, Chris Siu Yeung Chow, Xinguo Feng, Qingqing Li, Rui Jiang, Naipeng Dong, Guangdong Bai
TL;DR
Categorical graph generators need compatible node and edge types without depending on node order, but fixed one-hot endpoints impose equidistant category geometry. EGF learns node and unordered-edge anchors, transports Gaussian noise through them with a permutation-equivariant graph transformer, and decodes categories at the endpoint. EGF performs best on all four QM9 metrics and has the lowest NSPDK-based MMD on ZINC250k, while its scope is limited by dense edge storage, omitted chemical attributes, and single-seed runs.
Problem
Categorical graph generation must coordinate node and edge types, variable graph sizes, and permutation symmetry, while fixed one-hot encodings impose equidistant category geometry.
Method
EGF learns trainable anchors for node and unordered-edge categories, transports Gaussian noise toward them with a permutation-equivariant graph transformer, and applies a terminal categorical readout.
Results
EGF has the best value on all four QM9 metrics and the lowest NSPDK-based MMD on ZINC250k.
Takeaways & Limitations
Learned endpoint geometry supports categorical graph transport while retaining masked, undirected, and permutation-equivariant generation.
Takeaways & Limitations
EGF uses dense N^2 edge storage, omits stereochemistry, isotope identity, and radical-electron counts, and reports runs from one training seed.
Abstract
from arXiv · showhide
Generating categorical graphs requires choosing node and edge types that form a coherent structure without depending on node order. Many graph generators encode categories as fixed one-hot vectors, which can impose an artificial geometry in which categories are equidistant. We propose Embedded Graph Flows (EGF), a generative model that learns continuous embeddings for node and unordered-edge categories and transports Gaussian noise towards these learnt endpoints using a permutation-equivariant graph transformer. A terminal readout maps the embeddings back to discrete graph categories. Across molecular benchmarks, EGF achieved competitive performance. On QM9, EGF gives the best result on all four reported metrics among the three methods, including a Fréchet ChemNet Distance (FCD) of 0.150, compared with 0.717 for the categorical-diffusion baseline DiGress and 0.812 for the bridge-based baseline GruM. When applied to larger molecules in ZINC250k, EGF retains the lowest maximum mean discrepancy (MMD) using the neighbourhood subgraph pairwise distance kernel (NSPDK), indicating close agreement with the local substructures of the reference molecules. Our code is available at https://github.com/Trusted-System-Lab/EGF.
1 Introduction
Categorical graph generation must coordinate node and edge types while supporting variable sizes and permutation symmetry. EGF addresses fixed one-hot geometry by learning category anchors and transporting Gaussian noise through an equivariant continuous graph representation.
- Graph generation must jointly choose compatible node and edge types because local connectivity and chemical validity are coupled.
- Fixed one-hot encodings make all distinct category labels equidistant, motivating learned endpoint geometry for categorical graphs.
- EGF learns trainable anchors for node and unordered-edge categories in R^d and uses them as clean transport endpoints and regression targets.
- A permutation-equivariant graph transformer predicts clean embeddings along Gaussian-to-data paths, followed by ODE integration and terminal categorical decoding.
- EGF masks padded coordinates, shares one trajectory per unordered node pair, and preserves equivariance under node relabelling.
2 Related Work
Prior graph flow methods transport categorical or continuously encoded graph states using fixed representations or stochastic bridges. EGF instead learns a continuous anchor space and directly transports graph states while preserving padding, undirected-edge, and node-relabelling symmetries.
- Prior methods use score-based diffusion, Ornstein–Uhlenbeck bridge mixtures, categorical diffusion, or continuous-time category transitions.
- EGF transports Gaussian vectors toward learned node and edge anchors with a deterministic ordinary differential equation and categorical endpoint readout.
- Unlike fixed one-hot endpoint models, EGF predicts clean vectors in a learned anchor space before deriving the continuous field and reading categories at t = 1.
- EGF adapts shared continuous-to-discrete transport and terminal readout from sequence modeling to graph node and edge variables.
- For graphs, EGF masks padding, shares states across reversed edge indices, and makes transport and readout equivariant to node relabelling.
3 Method
EGF represents node and unordered-edge categories as learned continuous anchors, transports Gaussian noise toward these endpoints with a permutation-equivariant graph transformer, and decodes categories at the terminal state. Its graph representation masks padding, mirrors undirected edges, and supports node-aligned diagonal states.
- Graph representation and generation: EGF embeds each active node or unordered pair as a continuous vector and predicts clean embedded graphs from noisy states before terminal categorical decoding.The learned node and edge anchors define transport endpoints and regression targets.
- Graph representation and generation: The padded graph state uses active node and unordered-pair coordinates, while the mask excludes padding and the edge class none remains a valid target.The node output alphabet also includes absent for terminal readout, although reported training targets exclude that class.
- Graph representation and generation: Learned node and edge embedding maps place category anchors in R^d, allowing training to determine their relative geometry rather than using fixed one-hot distances.The node and edge maps have separate jointly trained weights whose columns define the anchors.
- Flow transport: EGF connects Gaussian source states to clean embedded endpoints along straight conditional paths and estimates the generation field from the current state toward the posterior mean endpoint.The clean endpoint is unknown during generation, so the graph transformer approximates the posterior mean used by the marginal field.
- Graph representation and generation: Undirected edge states are stored symmetrically, so one unordered-pair vector and one categorical decision represent both directions.The same symmetry is used for edge noise and decoded edge labels.
- Two-branch training: Training alternates denoising and endpoint-decoding branches: denoising predicts clean embeddings, while decoding trains the categorical readout on clean embeddings mixed with Gaussian noise.The separate node and edge coordinate counts prevent quadratic edge growth from reweighting the two coordinate types.
4 Experiments
The experiments evaluate EGF and prior methods on QM9 and ZINC250k using standardized molecular graph metrics. EGF leads all four QM9 metrics and achieves the lowest NSPDK-based MMD on ZINC250k.
- Datasets: EGF is evaluated on QM9 and ZINC250k, with 10,000 molecules assigned to both test and validation sets.The retained training sets contain 109,012 QM9 molecules and 229,455 ZINC250k molecules.
- Datasets and Encodings: The experiments use aromatic and kekulised encodings with distinct node and edge category definitions.Aromatic graphs use five edge categories, while kekulised graphs use four; neither encoding represents stereochemistry.
- Evaluation Metrics: All reported metrics use the same implementations, test reference, generated-sample count, and strict validity rule across methods.Strict validity requires reconstruction without correction into one connected molecule accepted by chemical sanitisation software.
- Evaluation Metrics: Scaffold similarity compares relative scaffold frequencies after removing overall count scale.The comparison uses cosine similarity between corresponding scaffold count vectors.
- Results: QM9: EGF achieves the best value on all four metrics, including validity 99.87%, FCD 0.150, NSPDK-based MMD 0.0019, and scaffold similarity 0.9936.The corresponding DiGress and GruM values are 98.92% and 99.69% validity, 0.717 and 0.812 FCD, 0.0031 and 0.0036 MMD, and 0.7984 and 0.8086 scaffold similarity.
- Results: ZINC250k: EGF achieves the lowest NSPDK-based MMD at 0.0024, while GruM has better FCD, validity, and scaffold similarity.EGF records FCD 3.326, validity 96.56%, and scaffold similarity 0.5546; the comparisons are reported in the same passage.
5 Discussion
The discussion interprets the metrics as complementary and contrasts EGF’s strong local-structure agreement on QM9 and ZINC250k with weaker broader-distribution and validity results on larger graphs.
- Metric Interpretation: NSPDK-based MMD, scaffold similarity, and FCD measure complementary aspects of generated molecular graph distributions.NSPDK-based MMD compares local graph neighbourhoods, scaffold similarity compares Bemis–Murcko framework frequencies, and FCD compares learned ChemNet feature statistics.
- QM9: On QM9, EGF’s strong results suggest agreement with both local graph structure and broader molecular distributional properties.QM9 contains graphs with up to 9 nodes and 36 unordered pairs.
- ZINC250k: On ZINC250k, EGF maintains strong NSPDK-based MMD agreement, but its higher FCD and lower validity than GruM show that local substructure matching is not sufficient for broader agreement or validity.ZINC250k increases capacity to 38 nodes and 703 unordered pairs, adding more jointly constrained atom–bond assignments.
6 Limitations and Future Work
The paper identifies unresolved component effects, scaling constraints, omitted molecular attributes, and uncertainty from using a single training seed. It proposes richer context-dependent embeddings as future work.
- Open Questions: A component study is needed to separate the effects of learned endpoint geometry, the graph transformer, two-branch training, and terminal readout.The paper does not isolate these components experimentally.
- Embedding Expressivity: Category-only endpoint embeddings may be less expressive across the broader range of environments in larger molecules and may contribute to higher ZINC250k FCD.Future work could use context-dependent node and edge representations from a pretrained molecular graph encoder.
- Scalability: Dense edge states scale quadratically with graph capacity, so larger graphs require additional memory and pairwise computation.The largest graphs considered use N = 38, and the edge tensor contains N^2 entries.
- Representation Scope: The reported encodings omit stereochemistry, isotope identity, and radical-electron counts, which would require extensions to category sets and reconstruction rules.These molecular attributes are outside the reported representation scope.
- Reproducibility: EGF runs use one training seed, so replicate runs are needed to quantify uncertainty in modest ZINC250k MMD and QM9 validity margins.The paper identifies these margins as requiring uncertainty assessment.
- Future Applications: The paper proposes future extensions for model usage control and data authorisation in proprietary molecular-generation settings.These extensions would use EGF’s continuous graph states and separated embedding, transformer, and readout parameters.
7 Conclusions
EGF transports Gaussian states through learned node and edge embeddings and decodes them into categorical graphs while preserving graph symmetries. It performs best on all four QM9 metrics and has the lowest NSPDK-based MMD on ZINC250k.
- Method: EGF transports Gaussian states toward learned node and unordered-edge anchors using a permutation-equivariant transformer and terminal categorical readout.Padding is masked, and each unordered edge uses one shared state and terminal decision.
- Results: EGF achieves the best value on all four QM9 metrics among the three methods.
- Results: EGF achieves the lowest NSPDK-based MMD on ZINC250k among the three methods.