DiGress: Discrete Denoising diffusion for graph generation

In plain words

A graph is a set of dots with lines between some pairs, such as a molecule with atoms and bonds. Generating a new graph is hard because the same graph can be written in many orders, and most of the entries that describe it are zeros.

This paper adapts the "diffusion" recipe behind image generators to graphs. In training, the true graph is gradually scrambled: some atom types change, some lines appear or disappear. A neural network learns to undo one small scramble at a time. To make a new graph, the process starts from a random scramble and the network is applied again and again until a clean graph is left.

The authors scramble with discrete choices, so each atom and each bond is always one of a fixed set of types. The earlier graph-diffusion methods added continuous noise, which made the noisy graphs dense. The network itself is a transformer that predicts a label for every pair of nodes. The paper reports that the result matches or beats earlier one-shot graph generators, and on a large molecule dataset it reaches the level of a model that builds molecules one piece at a time.

TL;DR

• DiGress is a discrete denoising diffusion model for graphs with categorical node and edge attributes, where "no edge" is one of the edge classes §3.
• Each node and each edge is corrupted independently by a Markov chain of category changes; the reverse model is a graph transformer that predicts clean classes with cross-entropy §3.1 §3.2 eq. 3.
• Three additions carry most of the reported gains: marginal noise that keeps training-set class frequencies §4.1, cycle and spectral features of the noisy graph §4.2, and a discrete regressor guidance rule for graph properties §5.
• Author-reported: best one-shot results on SBM and planar benchmarks; QM9 validity 99.0% tab. 2; MOSES and GuacaMol at the scale of about one million molecules tab. 3 tab. 4.
• Use it when you need a one-shot (non-autoregressive) graph generator with labelled nodes and edges; expect size-extrapolation limits on MOSES tab. 7.

Problem

Continuous diffusion on graphs (Gaussian noise on adjacency and features) destroys sparsity: noisy graphs become dense, and structural properties such as connectivity or cycle counts stop being defined on them §1. Diffusion for discrete data exists for text, images and audio, but graphs have varying sizes, permutation symmetry, and no tractable universal approximator, so those methods do not transfer directly §2. Non-autoregressive graph generators trailed autoregressive and motif-based models, which can use much more domain knowledge §6. The paper targets one-shot generation of graphs with categorical node and edge attributes, and reports results on datasets of more than a million molecules §1 §7.4.

Method

Representation. A graph with n nodes has node attributes X in R^(n x a) and edge attributes E in R^(n x n x b), each row a one-hot class. The absence of an edge is one edge class §3. Undirected graphs are handled by noising only the upper triangle and symmetrising §3.1.

Forward (noising) process. Noise acts per node and per edge as a Markov chain with transition matrices Q^t; the cumulative matrix is Q̄^t, so noisy states can be drawn in one step from the clean graph §2 §3.1 eq. 2. The paper uses a transition matrix that moves each class toward the empirical class frequencies m of the training set, instead of toward a uniform distribution §4.1 eq. 6:

Q^t_X = \alpha^t I + \beta^t \, \mathbf{1}_a m_X^\top, \qquad Q^t_E = \alpha^t I + \beta^t \, \mathbf{1}_b m_E^\top \qquad [\text{eq. 6}]

With alpha-bar_t = cos^2( 0.5 pi (t/T + s) / (1 + s) ) for a small s, the cumulative kernel has the same form with alpha-bar_t and 1 - alpha-bar_t §4.1. The limit distribution is then the product of the marginals, which the paper shows is the l2 projection of the data distribution onto product distributions §4.1, Theorem 4.1.

Reverse posterior. The closed-form posterior of one step is proportional to the one-step transition times the noisy-to-clean kernel §2 eq. 1. The network outputs a distribution p̂ over clean classes for each node and edge. The reverse step marginalises the clean class under that prediction, and the product over nodes and edges gives the joint reverse distribution §3.1 eq. 4 eq. 5. Nodes and edges are then sampled independently alg. 2.

Training loss. Cross-entropy on nodes plus lambda times cross-entropy on edges, summed over all node and edge positions §3.2 eq. 3. The paper says lambda sets the relative weight of nodes and edges but gives no value §3.2.

Denoiser. A graph transformer that carries node, edge and graph-level features §3.2 app. B.1. Node features are updated by self-attention, then conditioned on the edge and global features through FiLM layers; the edge features are updated from the unnormalised attention scores; graph-level features are pooled from nodes and edges. Layers have residual connections and layer normalisation. The timestep, normalised to [0, 1], is one global feature §3.2. Cost is Theta(n^2) per layer §3.2. The architecture is permutation equivariant and the loss is permutation invariant, so training does not need random node orderings or graph matching §3.3 eq. 3.

Structural and spectral features. Computed from the noisy graph at every step and appended to the network input §4.2 app. B.2. Node features count 3-, 4- and 5-cycles through each node; graph features count 3- to 6-cycles, given by closed-form expressions in powers of the adjacency matrix app. B.2. Spectral features use the Laplacian: the number of connected components and the first five non-zero eigenvalues go to the graph level; a non-largest- component indicator and the first two eigenvectors of the largest component go to the node level app. B.2. For molecules, valency and molecular weight are added app. B.2. The paper notes the eigendecomposition is O(n^3) and was not a bottleneck for graphs of up to 200 nodes app. B.2.

Sampling. Draw n from the training size distribution; start from the limit distribution; apply the reverse step for each timestep from T down to 1; return the final graph alg. 2 app. C.

Conditional generation (guidance). A regressor g_η is trained to predict target properties from noisy graphs. At each reverse step, the reverse distribution is tilted toward the target with a first-order term in the gradient of the squared error to the target, with a guidance scale §5 alg. 3. The scale value is not stated in the paper.

Substructure conditioning. A fixed subgraph is kept by overwriting the masked nodes and edges with the subgraph at every reverse step, in the style of inpainting app. E.

Likelihood. The model's likelihood is bounded by an ELBO with prior, diffusion and reconstruction terms, all KL divergences between categorical distributions app. C. The paper's likelihood results rely on this bound app. C tab. 2.

# Training step  [alg. 1]
# nodes: (n, a) one-hot classes; edges: (n, n, b) one-hot, no edge is a class
# Output: scalar loss

def training_step(nodes, edges):
    # Uniform timestep from 1..T  [code: src/diffusion_model_discrete.py#L413]
    step = sample_uniform_integer(1, T)
    # Sample the noisy graph from the closed-form kernel  [§3.1] [eq. 6]
    noisy_nodes, noisy_edges = corrupt(nodes, edges, step)
    # Structural and spectral features of the noisy graph  [§4.2] [app. B.2]
    extras = graph_features(noisy_nodes, noisy_edges, step)
    # Predict per-node and per-edge class probabilities  [§3.2]
    pred_nodes, pred_edges = denoiser(noisy_nodes, noisy_edges, extras, step)
    # Cross-entropy on nodes plus weighted cross-entropy on edges  [eq. 3]
    loss = cross_entropy(pred_nodes, nodes)
    loss += edge_weight * cross_entropy(pred_edges, edges)
    return loss

# Noisy class distribution of one node or one edge  [§4.1] [eq. 6]
# onehot: (k,) current class; marginal: (k,) training frequencies

def noisy_class_distribution(onehot, step, marginal):
    # Keep the current class with probability alpha_bar at this step  [§4.1]
    keep = alpha_bar[step]
    # Otherwise redraw from the training frequencies  [§4.1] [eq. 6]
    return keep * onehot + (1 - keep) * marginal    # (k,)

# Sampling  [alg. 2]
# Output: one graph with node and edge classes

def sample_graph():
    # Draw the number of nodes from the training distribution  [app. C]
    n = sample_size_from_training_data()
    # Start from the limit distribution, independently per node and edge  [§4.1]
    nodes, edges = sample_from_marginals(n)
    for step in range(T, 0, -1):
        # Features of the current noisy graph  [§4.2] [app. B.2]
        extras = graph_features(nodes, edges, step)
        # Predicted distribution over clean classes  [§3.2]
        pred_nodes, pred_edges = denoiser(nodes, edges, extras, step)
        # Marginalise the clean class to get the posterior  [eq. 1] [eq. 5]
        post_nodes = posterior(nodes, pred_nodes, step)
        post_edges = posterior(edges, pred_edges, step)
        # Sample the previous state independently per node and edge  [eq. 4]
        nodes, edges = sample_categorical(post_nodes, post_edges)
    return nodes, edges

# Guided reverse step  [alg. 3]
# regressor predicts target properties from a noisy graph; target y

def guided_step(nodes, edges, step, target, scale):
    # Unconditional denoiser prediction  [§3.2]
    pred_nodes, pred_edges = denoiser(nodes, edges, extras, step)
    # Gradient of the squared error to the target w.r.t. the noisy graph  [§5]
    grad_nodes, grad_edges = regressor_gradient(nodes, edges, target)
    # Tilt the reverse distribution by exp(-scale * <grad, G^(t-1)>)  [alg. 3]
    tilt_nodes, tilt_edges = linear_tilt(grad_nodes, grad_edges, scale)
    return sample_guided(pred_nodes, pred_edges, tilt_nodes, tilt_edges)

Key equations (main text):

q(z^{t-1} \mid z^t, x) \propto z^t (Q^t)^\top \odot x \bar Q^{t-1} \qquad [\text{eq. 1}]

\ell(\hat p^G, G) = \sum_i \mathrm{CE}(x_i, \hat p^X_i) + \lambda \sum_{i,j} \mathrm{CE}(e_{ij}, \hat p^E_{ij}) \qquad [\text{eq. 3}]

p_\theta(G^{t-1} \mid G^t) = \prod_i p_\theta(x_i^{t-1} \mid G^t) \prod_{i,j} p_\theta(e_{ij}^{t-1} \mid G^t) \qquad [\text{eq. 4}]

p_\theta(x_i^{t-1} \mid G^t) = \sum_{x} p_\theta(x_i^{t-1} \mid x_i = x, G^t)\,\hat p^X_i(x) \qquad [\text{eq. 5}]

Figures: fig. 1 is a schematic of the forward and reverse chains and of the denoiser. fig. 2 compares reverse chains from uniform and marginal noise; with uniform noise, a sparse graph takes many steps to appear §4.1.

Hyperparameters

The paper's main text and appendix state few training values. The table gives what the paper states, then the values in the official configuration files at the code commit. I did not check that the configuration values reproduce the paper's runs. Code paths are relative to the repository root.

Results

Author-reported throughout. Setups: SBM and planar are two sets of 200 graphs each (SBM up to 200 nodes, planar 64 nodes), scored by degree, clustering and orbit distributions and by the valid-unique-novel rate §7.1 app. F.1. QM9 is scored on 10k generated molecules, validity by RDKit sanitisation §7.2. MOSES and GuacaMol have more than a million molecules each §7.4.

tab. 1 Unconditional generation on SBM and planar graphs (lower is better for Deg, Clus, Orb; higher is better for V.U.N.)

tab. 5 Community-20 (small SBM set, 200 graphs) app. F.1. Lower is better. Ratio is the MMD-based score.

tab. 2 Molecule generation on QM9 (training time: hours to reach 99% validity). NLL is an ELBO bound app. C. Validity is without the validity correction used by GDSS §7.2.

QM9 novelty is not reported in the main table, because QM9 enumerates small molecules exhaustively. The paper gives 33.4% novelty for DiGress on QM9 with implicit hydrogens, and 40.0% for ConGress app. F.2.

tab. 6 Ablation on QM9 with explicit hydrogens app. F.2. Test NLL on this setting is 129.7 app. F.3.

Conditional generation on QM9 §7.3. Mean absolute error between target and the property of a generated molecule. Properties are computed with Psi4 on RDKit conformers. 100 test molecules set the targets; 10 molecules are generated per target.

tab. 3 Molecule generation on MOSES §7.4. FCD, SNN and scaffold use the separate-scaffold test set. Higher is better except FCD.

tab. 4 Molecule generation on GuacaMol §7.4. Scores: higher is better for all metrics in the paper. Mapping about 20% of GuacaMol training molecules to a graph and back fails, so those molecules cannot be scored app. F.3.

Test NLL on GuacaMol is 308.1 and on MOSES 205.2 (separate-scaffold set) app. F.3.

tab. 7 Size extrapolation app. F.3. Proportion of valid and unique molecules when sampling graphs larger than the training maximum. The paper notes these molecules are novel because they are larger than any training graph.

Paper vs code

Official code: https://github.com/cvignac/DiGress, MIT license, commit bdfa10cdedc9e274c755ac17dc42d633ea8f5909. The repository is linked from the paper's footnote §7 and is by the authors.

• Edge weight lambda. Paper: lambda is in R+ with no value given §3.2 eq. 3. Code: lambda = 5 on edges and 0 on the graph-level term in every experiment config discrete.yaml qm9_with_h.yaml train_metrics.py.
• Loss normalisation. Paper: eq. 3 uses sums over nodes and edges eq. 3. Code: the cross-entropy metric returns total loss divided by sample count, a mean per row abstract_metrics.py. If the training loss takes that value, each term is a mean, not a sum, so the effective weight of lambda differs from eq. 3. I did not measure the effect.
• Optimizer. Paper: not stated §7. Code: AdamW with amsgrad always on diffusion_model_discrete.py. The config key train.optimizer is not read by this code path; the GuacaMol config sets adam guacamol.yaml, which has no effect.
• Residual on the noisy input. Paper: the network "predicts the clean graph" §3.1. Code: the output adds the noisy input's node and edge features (the first out_dim channels) to the predicted logits, and the edge output is masked on the diagonal and symmetrised transformer_model.py transformer_model.py.
• Cycle features. Paper: formulas only app. B.2. Code: node and graph cycle counts are divided by 10 and clipped at 1 extra_features.py. Large counts (dense graphs) are therefore clipped.
• Eigenvalue features. Paper: first five non-zero eigenvalues app. B.2. Code: eigenvalues are divided by the node count before selection extra_features.py extra_features.py.
• Six-cycle formula. Paper's printed y_6 formula has a Frobenius-norm term, 9 ||A (A^2 ⊙ A^2)||_F, and a 3 ||A^3||_F term app. B.2. Code computes sum(A ⊙ (A^2)^2) for the first and the sum of all entries of A^3 for the second extra_features.py. The code is what produced the released samples. I did not verify whether the paper's notation was intended differently.
• Posterior at sampling. Paper: p_θ(x^(t-1) | x_i = x, G^t) = q(...) when q(x^t | x) > 0, else 0 eq. 5. Code: sums over clean classes and normalises per node and per edge, with a fallback of 1e-5 for rows that sum to zero diffusion_model_discrete.py.
• Noise schedule. Paper: cosine schedule with "small s" §4.1. Code: s = 0.008 diffusion_utils.py. Per-step noise is clamped at 0.9999 before alpha-bar is built noise_schedule.py.
• Guidance and substructure conditioning. Paper: guidance §5 alg. 3 and scaffold extension app. E. Code: the reverse sampler contains no regressor term and no masking step diffusion_model_discrete.py. The only guidance-related code is a dataset transform that reads a regressor attribute, which the repository never sets qm9_dataset.py. The paper's guided results cannot be reproduced from this repository alone.
• QM9 validation size. Paper: 100k / 20k / 13k §7.2. Code: 100k train, 10% test, the remainder validation qm9_dataset.py.
• GuacaMol sample count. Paper: not stated app. F.3. Code: 3 final samples; the 18000 setting is commented out guacamol.yaml. Reproducing a GuacaMol score from this config would generate 3 molecules.

Pitfalls

• Unstated training values. Learning rate, epochs, batch size, depth, widths, T, lambda and the scale of s are absent from the paper's text. An agent implementing from the paper alone will have to choose them. The code values are in the table; use them as a starting point and report any change.
• Edge weight lambda. The most consequential unstated number is lambda. Code uses 5 on edges and 0 on the graph term. Since eq. 3 defines the loss with a free lambda, check the effective scale after normalisation, because the code's per-row means and the paper's sums differ abstract_metrics.py.
• Residual to the noisy input. The output adds the noisy input features to the network logits. An implementation that predicts logits from scratch will train to a different function and is not the released model transformer_model.py.
• Posterior numerics. The posterior divides by a sum over clean classes. The code guards rows that sum to zero with 1e-5 and renormalises diffusion_model_discrete.py. The fallback is absent from eq. 5, so an implementation that follows the paper literally can divide by zero on such rows.
• Clamping in the schedule. Betas are clamped at 0.9999 and alpha-bar is built from the clamped alphas noise_schedule.py. This changes the schedule only where a per-step beta is extreme, but the paper does not mention it.
• Feature clipping hides large counts. Cycle counts are scaled by 1/10 and clipped at 1 extra_features.py. Any count above 10 saturates, so the feature loses information on graphs with many cycles. Check the count distribution on your data before relying on it.
• Padding in pooling. The graph-level pooling takes mean, min, max and standard deviation over all rows layers.py layers.py. Padded rows enter as zeros through the mask. For graphs smaller than the batch maximum, min and max may include those zeros. I did not check whether this changes results on the released datasets; the paper does not discuss batching.
• Cost is quadratic and eigendecomposition is cubic. Every layer costs Theta(n^2) §3.2. Spectral features need an eigendecomposition per graph per step, O(n^3) app. B.2 extra_features.py. The paper says this was not a bottleneck at 200 nodes; at larger sizes it will be. Cycle features are recomputed every step app. B.2.
• Scale of results. Only the QM9 table reports variability (five runs). Other tables report single numbers §7.1 §7.4. Do not treat differences of a few points on MOSES or GuacaMol as significant.
• Validity conventions. QM9 validity is RDKit sanitisation without the relaxed charge rule used in GDSS app. F.2. MOSES and GuacaMol are converted to SMILES with GDSS's code that allows partial charges app. F.3. Scores therefore depend on the converter, and about 20% of GuacaMol training molecules cannot be converted back app. F.3.
• NLL is a bound. The NLL column is an ELBO-based bound, not an exact likelihood app. C tab. 2.
• Novelty and MMD conventions. QM9 novelty is left out of the main table app. F.2. SBM and planar ratios use a denominator from the SPECTRE table, and the paper notes that what it calls MMD is actually squared MMD app. F.1.
• Size extrapolation is uneven. Sampling larger than the training maximum fails on MOSES (2.2% valid-unique at +10) and works on GuacaMol (85.6%) tab. 7. The paper's explanation, dataset diversity, is a conjecture app. F.3. The network is not designed to be size-invariant; it uses sum aggregation in each layer app. F.3.
• Conditional guidance cannot be reproduced from the repository. The regressor is not in the repository, the guidance step is not in the sampler, and the guidance scale is not stated §5 alg. 3 diffusion_model_discrete.py. Any guided result needs the authors' code, or a new implementation whose numbers are then the implementer's own.
• Conditional MAE depends on conformers. Properties are computed with RDKit conformers and Psi4 on 100 targets, 10 samples each §7.3. Property errors mix generation error with conformer and Psi4 error; the paper does not separate them.
• Hydrogens. Main QM9 results use implicit hydrogens; explicit-H results are an ablation §7.2 tab. 6. Results with and without H are not comparable.
• Equivariance is not exchangeability of the data. The paper proves the model is permutation equivariant and the loss invariant §3.3. The exchangeability of the generated distribution relies on an external theorem the paper cites from GeoDiff app. D. Those are the claims the paper makes; it does not show exact likelihood invariance beyond them.

Lineage

• Builds on: diffusion models (Sohl-Dickstein et al. 2015; Ho et al. 2020) §2; discrete state-space diffusion, Austin et al. 2021 §2 (not carded; no arXiv id in the paper's bibliography); graph transformer, Dwivedi and Bresson 2020 §3.2 (not carded); classifier guidance, Dhariwal and Nichol 2021 §5 (not carded).
• Continuous graph diffusion it replaces: Niu et al. 2020 and GDSS (Jo et al. 2022, arXiv:2202.02514 per the bibliography) §1 §6. The authors' Gaussian model ConGress is a control with the same denoiser §7.1 app. A.
• Compared against: SPECTRE (Martinkus et al. 2022, arXiv:2204.01613 per the bibliography) §7.1 tab. 1; GraphRNN, GRAN, GG-GAN, Set2GraphVAE (Vignac et al. 2021, "Top-N", arXiv:2110.02096 per the bibliography), GraphNVP, JT-VAE, NAGVAE and GraphINVENT §7 tab. 3. None of these cards exist in this wiki yet, so no concept or result is taken from them here.