MUDiff: Unified Diffusion for Complete Molecule Generation

• MUDiff applies discrete noise to graph structure and continuous noise to atom features and coordinates; MUformer couples 2D/3D attention while preserving 3D roto-translation equivariance.The denoiser predicts clean graph structure and estimates noise for atom features and coordinates.Source: §1 · p. 2
• MUDiff and MUformer can operate when either 2D graph structure or 3D geometric information is missing, supporting complete representations from incomplete molecular inputs.This is useful when 3D coordinates are unavailable because of experimental or computational limitations, including the cost of conformational analysis.Source: §1 · p. 2
• 7.9% more stable molecules and 2% greater molecular uniqueness than existing methods; MUDiff also remains competitive when trained with limited 3D structures against methods trained with complete 3D structures.Source: §1 · p. 2

2.1 Diffusion Models

• A diffusion model adds noise through a Markov process and trains a denoising network to estimate the clean data from a noisy point.For continuous data, the forward process is Gaussian; the signal factor α_t and noise scale σ_t control the retained signal and added noise, with α_t = √(1 − σ_t²) for variance preservation.Source: §2.1 · p. 2
• For discrete data, transition matrices corrupt one-hot states, while cumulative transitions give noisy marginals and Bayes’ rule gives the posterior over earlier states.This approach is used because Gaussian noise may damage graph sparsity and connectivity.Source: §2.1 · p. 3

2.2 The Basics of Molecules

• A molecule is represented as M = (H, E, X): atom features, a symmetric fully connected bond-type graph with a no-bond type, and 3D atom coordinates.Pairwise 3D distances are encoded with exponential normal radial basis functions; a cosine cutoff smoothly reaches zero at 5 Å.Source: §2.2 · p. 3

3.1 Diffusion Process

• MUDiff applies continuous noise to atom features and coordinates, while discrete noise changes edge types through uniform transitions and preserves undirected graph symmetry.Coordinate noise is constrained to the zero-center-of-gravity subspace, so perturbations preserve the molecule’s center of gravity and support translation-invariant denoising.Source: §3.1 · pp. 3–4

3.2 Denoising Process

• MUformer denoises noisy atom features, coordinates, and edges together, predicting atom-feature and coordinate noise alongside edge-type probabilities while maintaining equivariance.For atom features and coordinates, training minimizes the distance between true and predicted noise; edge types are trained with cross-entropy classification loss.Source: §3.2 · p. 4

4 MUformer: Molecule Unified Transformer

• MUformer maps a molecule’s atom features, edges, and coordinates to predicted atom features, edges, and coordinates using 6 encoding functions, 4 attention biases, and 2 attention channels.The paper denotes the input as M = (H, E, X) and the predicted molecule as ˆM = (ˆH, ˆE, ˆX).Source: §4 · p. 5
• With only 2D input, MUformer activates its invariant channel to predict atom and edge features; with only 3D input, it activates its equivariant channel to predict atom features and coordinates.When both 2D and 3D data are provided, both channels activate to predict the complete molecule, including its atom features, structure, and geometry.Source: §4 · p. 5

4.1 Encodings

• Of MUformer’s 6 encoding functions, 3 use message passing to incorporate atomic, positional, and structural information into representations suited to graph-structured inputs.Source: §4.1 · p. 5

• The method in uses each atom’s in-degree and out-degree from the 2D molecular graph to incorporate centrality into atom-wise encoding.Source: §1 · p. 5

2. Bond Encoding

• A permutation-invariant message-passing function embeds pairwise atom information into each edge, making the representation independent of atom order.The embedding combines projected features of both endpoint atoms and the edge: zeij = Wcomb1([Watom2hi + Wedge1eij + Watom2hj]) + bcomb1.Source: §2 · p. 5
• The final edge encoding averages the representations in both directions to ensure symmetry.zeij = (zeij + zeji)/2.Source: §2 · p. 5

3. Graph Encoding

• Graph-level features include cycle counts up to size 6, connected components, eigenvalue features, atom valencies, and molecular weight.The spectral features include the multiplicity of eigenvalue 0 and the first five nonzero eigenvalues.Source: §3 · p. 5

4. 2D Neighborhood Encoding

• 2D neighborhood encoding aggregates information from each atom’s immediate graph neighbors using message passing and atom–edge embedding interactions.The messages use the Hadamard product of atom and edge embeddings.Source: §3 · p. 6

5. 3D Neighborhood Encoding

• 3D neighborhood encoding uses interatomic Euclidean distances and exponential radial basis features to form messages from nearby atoms.Messages combine atom and edge embeddings with the Hadamard product.Source: §3 · p. 6
• A cosine cutoff includes only atoms with positive cutoff values, and setting it to 1 for 2D-connected atoms combines graph and spatial locality.The cutoff provides a smooth way to incorporate spatial locality into message passing.Source: §3 · p. 6

6. Combine Encoding

• Atomic embeddings concatenate atom, 2D-neighborhood, and 3D-neighborhood encodings.Bond and graph encodings are also used to calculate attention.Source: §6 · p. 6

4.2 Attention Biases

• Attention biases encode 2D shortest-path distances and path-edge information, 3D Euclidean distances with an exponential radial basis function [9], and projected edge and graph features.These biases are integrated into attention to represent spatial relationships and molecular structure.Source: §4.2 · p. 7

4.3 Transformer Channels

• MUformer’s invariant channel predicts atom and edge features, while its equivariant channel uses 2D and 3D structure to predict atom features and coordinates.The channels’ atom representations are combined for subsequent transformer layers or final predictions; the output block uses atom, edge, and velocity embeddings to make final predictions.Source: §4.3 · p. 7

5.1 Graph Models and Equivariant Models

• Graph models capture atom–bond relationships, while equivariant graph models support geometric tasks by learning representations robust to transformations such as translation, rotation, and reflection.Equivariance preserves a function’s structure under group transformations, expressed as f(gx) = gf(x).Source: §5.1 · p. 8

5.2 Diffusion Models

• Diffusion-based molecular generation includes methods for low-energy conformations [3], joint denoising of atom features and molecular structures [6], and equivariant generation of atom features and coordinates [4].These works illustrate the integration of diffusion models with graph transformers or equivariant graph models.Source: §5.2 · p. 8

5.3 Joint Generative Models

• Joint generators include an autoregressive flow that sequentially generates atom types, bond types, and 3D coordinates [32], a diffusion model that jointly generates consistent atoms and bonds [33], and a method generating 2D and 3D molecular graphs for structure-based drug design.The flow predicts relative coordinates in a local spherical coordinate system; the diffusion model gradually adds guided noise to atom positions and types and bond types.Source: §5.3 · p. 8

6.1 Molecule Generation with Limited 3D Data

• With 30K molecules containing both 2D and 3D structures and 70K containing only 2D structures, MUDiff achieves competitive stable-molecule generation results.The two-channel design supports this limited-3D training framework; the authors suggest the results indicate that sufficient 2D structures can help infer 3D geometry.Source: §6.1 · p. 8, Table 1 · p. 8

6.2 Molecule Generation

• MUDiff significantly outperforms baselines on negative log-likelihood and molecule stability, matches SOTA atom stability, and produces more unique molecules while matching SOTA validity.Table 2 reports stability metrics over 10,000 generated samples across three runs; Table 3 reports validity and uniqueness over 10K molecules across three runs. Novelty is not reported.Source: Table 2 · p. 9, §6.2 · p. 9, Table 3 · p. 9

6.3 Conditional Generation

• Conditional MUDiff is trained on QM9 to generate molecules conditioned separately on α, ϵhomo, ϵlumo, ∆ϵ, µ, and Cv.The experimental setting follows [4].Source: §6.3 · p. 9

6.4 Ablation Study

• The paper conducts an ablation study in App D.4 to investigate how different MUDiff components affect performance.Experimental settings are described in App D.4.1.Source: §6.4 · p. 9

6.5 Property Prediction

• MUformer is compared with seven baselines on QM9 property prediction, estimating 12 chemical properties per molecule and reporting mean absolute error in App D.3.The baselines are SchNet [9], EGNN [2], PhysNet [17], DimeNet, Cormorant [36], PaiNN [42], and ET [21].Source: §6.5 · p. 9

7 Conclusion

• MUDiff generates complete molecule representations and produces more stable and valid molecules; jointly generating 2D and 3D structures improves overall performance.The conclusion also identifies scalability as an issue for future work.Source: §7 · p. 9

A.1 MUformer Encodings

• MUformer uses six encoding functions, three based on message passing, to combine atomic, positional, and structural information for graph-structured inputs.Source: §A.1 · p. 13

1. Atom Encoding

• Atom encoding adds 2D-graph in-degree and out-degree centrality to atom features, reflecting each atom’s relative importance.The resulting atom embedding includes degree centrality information.Source: §1 · p. 13

2. Bond Encoding

• Bond encoding combines endpoint atom features with edge types and uses permutation-invariant construction and symmetrization to produce order-independent edge representations.Source: §2 · p. 13

3. Graph Encoding

• Graph-level encoding combines cycle counts and connected components, eigenvalue features, atom valencies, and molecular weight, addressing GNN limits on cycle detection and following [6] [20].Source: §3 · p. 13
• Neighborhood encodings use 2D message passing over adjacent atoms and 3D distance-based messages with a cosine cutoff that excludes atoms when fcos(dij) ≤ 0.The 3D messages use radial basis features; the cutoff focuses computation on spatially nearby atoms [2].Source: §3 · pp. 13–14

6. Combine Encoding

• The final atomic embedding concatenates centrality, 2D-neighborhood, and 3D-neighborhood embeddings into a unified representation.Source: §6 · p. 14

A.2 Attention Biases

• MUformer uses four attention biases: 2D shortest-path and path-edge features, 3D Euclidean distances with radial basis functions, and projected edge and graph embeddings.The 3D radial basis approach follows [9].Source: §A.2 · pp. 14–15

A.3 Transformer Channels

• MUformer activates its invariant channel for 2D inputs, its equivariant channel for 3D inputs, or both for complete molecule predictions, including atom features, molecular structure, and coordinates.The invariant channel predicts discrete 2D structure and invariant atom and graph features; the equivariant channel predicts atom features and coordinates.Source: §A.3 · pp. 15–17
• MUDiff uses O(nd + 3n + n^2b) memory, more than EDM’s O(nd + 3n) and DiGress’s O(nd + n^2b), while representing both 2D and 3D molecular information.Source: §A.3 · p. 17

B Importance of 2D and 3D Attention Biases

• MUformer’s 2D shortest-path and edge-type biases and 3D Euclidean-distance radial-basis bias encode topological and geometric relationships to focus attention on relevant atom connections.Source: B Importance of 2D and 3D Attention… · pp. 17–18

C Analysis of Interdependence of Generation of 2D and 3D Structures

• The paper states that generating 3D structures can refine 2D structures by providing information about conformational space, stereochemistry and chirality, and additional constraints on molecular geometry.For example, when the 3D structure indicates that two atoms are connected by a ring, the algorithm can adjust the 2D structure to ensure the ring is planar.Source: C Analysis of Interdependence of… · pp. 18–19

D.1 Molecule Generation on DRUG

• The DRUG evaluation compares MUDiff with GDM, EDM, MDM, and GeoLDM on atom stability, molecule stability, and validity across three runs.Table 4 provides the evaluation scope; the supplied passages do not give its values.Source: §D.1 · p. 19, Table 4 · p. 19

D.2 Conditional Generation

• Conditional MUDiff samples more closely resemble QM9 subset B than the baselines, indicating better generation of molecules with desired properties and structural similarities.The comparison uses a property classifier; Table 5 reports its mean absolute error for subset B, MUDiff samples, and three baselines.Source: §D.2 · p. 19, Table 5 · p. 19

D.2.1 Conditional Generation

• Conditional generation samples a property condition and generates molecules from it, covering polarizability, orbital energies, energy difference, dipole moment, and heat capacity at 298.15K.The six conditioned properties are α, ϵhomo, ϵlumo, ∆ϵ, µ, and Cv.Source: §D.2.1 · p. 20

D.3 Property Prediction on QM9

• MUformer is evaluated for property prediction on all QM9 targets using mean absolute error and standard deviation averaged over three runs; baseline results come from prior literature.The reported setup uses learning rate 1e-3, 1e-4, 5e-4, and 1e-5, weight decay 5e-5, 128 hidden dimensions, and six layers; the supplied passages do not provide outcome values.Source: §D.3 · p. 20, Table 6 · p. 20

D.4 Ablation Study; D.4.1 Ablation Study

• The full MUformer configuration produces the most stable and valid molecules among the ablation models.Ablations use smaller models and 1,000 generated molecules per variant; the variants add combinations of structure, neighborhood, attention-bias, and graph features.Source: §D.4 · p. 20, Table 7 · p. 21, §D.4.1 · p. 21

E Limit Distribution for Edges

• As t approaches T, the edge-noising distribution converges to a uniform distribution over edge categories, independent of the clean edges.The transition follows [6] [16] and uses a parameter αt that smoothly decreases from 1 to 0.Source: E Limit Distribution for Edges · p. 22

F Likelihood and Lower Bound

• The variational log-likelihood lower bound comprises prior, reconstruction, and diffusion losses, with the diffusion loss covering atom features, edges, and coordinates.The reconstruction term separates likelihoods for atom features, edge features, and coordinates.Source: F Likelihood and Lower Bound · pp. 22–23

G Posterior Distribution of Edge Features p(˜Et−1|˜Et)

• The reverse posterior factorizes across atom features, edge features, and coordinates; during sampling, the unknown clean edges are replaced by the network estimate.The edge posterior is written as q(Et−1|Et, Ê).Source: G Posterior Distribution of Edge… · p. 23

H Zeroth Likelihood Estimation

• Zeroth-likelihood estimation handles edge features directly and uses different calculations for integer and categorical atom features.Categorical atom types use one-hot encoding and a normalized categorical distribution, following [4].Source: H Zeroth Likelihood Estimation · p. 24

I Node2Graph & Edge2Graph Functions

• Node2Graph and Edge2Graph aggregate feature values using mean, max, and min, concatenate them, and apply a linear transformation to obtain graph-level features.Node2Graph operates on node features, while Edge2Graph operates on edge features.Source: I Node2Graph & Edge2Graph Functions · p. 24

J Molecule Generation

• RDKit-based validity metrics can misclassify charged QM9 molecules, and novelty is difficult to interpret because QM9 enumerates molecules under specified constraints [2] [4] [5] [38].The stated RDKit procedure adds hydrogens to heavy atoms to match valency; novelty may indicate failure to capture the dataset distribution [38].Source: J Molecule Generation · p. 25

K.1 Scalability Issue

• MUDiff generates dense n × n × 4 bond tensors, which capture bond details but incur computational costs that can limit scalability.Sparse adjacency tensors and multi-resolution molecular representations are described as alternatives for reducing computation while retaining structural information.Source: §K.1 · pp. 25–26