Equivariant Diffusion Model (EDM)

The core architectural constraint ensures that rotating or translating the input noise results in an identically transformed output molecule. This is achieved by restricting all learnable operations—such as message passing and coordinate updates—to be roto-translationally equivariant. This physical symmetry is not learned from data but hard-coded into the network via tensor products of irreducible representations or by operating on scalar-vector pairs, guaranteeing that generated geometries are physically consistent regardless of orientation.

EDMs learn to reverse a Markov chain that progressively adds Gaussian noise to the 3D coordinates of atoms. The forward process destroys the molecular structure, and the neural network is trained to predict the noise vector at each step. Key aspects include:

• Linear vs. Cosine Schedules: Controlling the rate of information destruction.
• Score Matching: The network learns the gradient of the log probability density (the score function) of the data distribution.
• Equivariant Noise Prediction: The predicted noise must transform equivariantly with the input, a condition satisfied by architectures like EGNNs or Tensor Field Networks.

While the denoising process operates on equivariant coordinates, the learned probability distribution over molecular geometries must be invariant to rotation and translation. This is a natural consequence of using an equivariant transition kernel with an invariant base distribution (e.g., a Gaussian centered at the origin). The model learns a likelihood that is independent of the molecule's absolute pose, focusing purely on its internal, physically meaningful geometry.

A specialized variant, Torsional Diffusion, applies the diffusion process exclusively to the torsional angles of rotatable bonds while keeping bond lengths and angles fixed at their equilibrium values. This dramatically reduces the degrees of freedom and prevents the generation of physically implausible, distorted local geometries. It operates on a hypertorus manifold, requiring a wrapped normal distribution for the noising process, and is highly effective for generating realistic conformers of drug-like molecules.

The denoising function is typically parameterized by an E(n) Equivariant Graph Neural Network (EGNN). Unlike higher-order tensor methods, EGNNs achieve equivariance efficiently by maintaining two feature tracks:

• Scalar Features: Invariant node embeddings updated via standard message passing.
• Vector Features: Equivariant coordinate updates computed as a weighted sum of relative position vectors, with weights derived from the invariant messages. This design avoids expensive Clebsch-Gordan tensor products, enabling faster training and inference.

The training objective minimizes the difference between the predicted and true noise added to the coordinates. Formally, this is denoising score matching with an equivariant constraint. The loss is computed on the vector field of coordinate displacements, ensuring the network learns a physically consistent force field. This connects EDMs directly to the concept of learning neural network potentials (NNPs), where the score is proportional to the atomic forces that drive molecules toward low-energy conformations.

The fundamental symmetry principle ensuring that the model's output transforms consistently with any input rotation and translation in 3D Euclidean space. In an EDM, this guarantees that rotating a molecule's noisy coordinates results in an identically rotated denoised structure, making the generative process physically consistent and independent of the molecule's orientation in space.

The generative framework that EDMs extend. DDPMs learn to reverse a Markov chain that gradually adds Gaussian noise to data. Key aspects include:

• Forward process: Systematically destroys structure in data
• Reverse process: A neural network learns to iteratively remove noise
• Score matching: The model learns the gradient of the log probability density (the score function) at various noise levels

A computationally efficient architecture often used as the denoising backbone in EDMs. EGNNs achieve E(n) equivariance without expensive higher-order tensor products by operating directly on scalar features and updating vector features (like coordinates) using messages that are invariant to rotation. This provides a lightweight yet powerful mechanism for processing 3D molecular geometries.

An alternative formulation of diffusion models that directly estimates the Stein score (the gradient of the log-density) of the data distribution at varying noise scales. The EDM leverages this concept by training a network to predict the noise added to atomic coordinates, which is mathematically equivalent to learning the score. Sampling proceeds via Langevin dynamics or stochastic differential equation solvers.

A specialized diffusion approach that operates on the torsional angle space of molecules rather than full Cartesian coordinates. By only modifying rotatable bonds while keeping bond lengths and angles fixed, it generates physically realistic 3D conformers with significantly fewer degrees of freedom. This contrasts with EDMs that operate on all atomic coordinates but must implicitly learn chemical validity constraints.

A machine-learned surrogate model that predicts the potential energy and atomic forces of a molecular system directly from atomic coordinates. EDMs are closely related to NNPs, as both must learn physically meaningful energy landscapes. In fact, the score function learned by an EDM can be interpreted as a force field, and many EDM architectures borrow equivariant message-passing designs from state-of-the-art NNPs like NequIP and MACE.