Embedded Graph Flows for Categorical Graph Generation
Authors: Ethan Ma, Zihan Wang, Chris Siu Yeung Chow, Xinguo Feng, Qingqing Li, Rui Jiang, Naipeng Dong, Guangdong Bai
Organizations: School of Electrical Engineering and Computer Science, The University of Queensland, Australia · Institute for Molecular Bioscience, The University of Queensland, Australia · Qunar.com, China · Chinese Academy of Sciences, China · Department of Computer Science, City University of Hong Kong, Hong Kong SAR, China
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.
Figures & tables
Figure 1: Transport spaces and supervision for DiGress [ 21 ] , GruM [ 10 ] , and EGF, illustrated with C, N, and O. DiGress moves between simplex vertices with stepwise cross-entropy. GruM follows an Ornstein-Uhlenbeck bridge mixture over Gaussian-initialised, category-coded arrays and predicts G with weighted mean-squared error before quantisation. EGF moves a Gaussian source through Rd , using mean-squared error along the path and cross-entropy at t=1 . In its panel, W abbreviates Wr followed by argmax .
Figure 2: EGF trains a shared graph transformer in denoising and decoding modes. The maps WA,WB embed node and edge categories. The denoising branch predicts clean embeddings along the Gaussian path; the decoding branch feeds noise-perturbed clean embeddings to the same transformer and maps its output to categories with Wr . Unordered-edge states are mirrored, padding is masked, and D occupies the diagonal input.
Figure 3: EGF sampling uses S denoising calls from a Gaussian source followed by one terminal decoding call. Molecular reconstruction and sanitisation follow category selection; the intermediate panels visualise predicted clean embeddings.
QM9 (at most 9 nodes)
ZINC250k (at most 38 nodes)
Method
Valid (%) ↑
FCD ↓
NSPDK MMD ↓
Scaffold ↑
Valid (%) ↑
FCD ↓
NSPDK MMD ↓
Scaffold ↑
DiGress [ 21 ]
98.92
0.717
0.0031
0.7984
82.79
2.803
0.0032
0.4398
GruM [ 10 ]
99.69
0.812
0.0036
0.8086
98.65
2.276
0.0028
0.5802
EGF (ours)
99.87
0.150
0.0019
0.9936
96.56
3.326
0.0024
0.5546
Table 1: QM9 and ZINC250k results using the evaluation settings in Section 4.1.3 ; training settings are method-specific.