Source-linked AI summary
E(n) Equivariant Normalizing Flows
Victor Garcia Satorras, Emiel Hoogeboom, Fabian B. Fuchs, Ingmar Posner, Max Welling
TL;DR
Molecular generative models need to represent 3D positions and features while respecting Euclidean symmetries, a capability not provided by discriminative EGNNs alone. E-NFs integrate a stabilized EGNN into a continuous-time invertible flow and outperform existing normalizing-flow methods across particle and molecular benchmarks, while jointly modeling molecular features and positions in 3D.
Problem
Generative modeling of 3D molecular structures requires efficient equivariant approaches, while EGNNs cannot generate new molecular structures.
Method
E-NFs integrate a stability-adapted EGNN as the derivative of a continuous-time equivariant normalizing flow and jointly model positions with invariant molecular features.
Results
E-NFs considerably outperform existing normalizing flows in log-likelihood on DW4, LJ13, and QM9, and improve generated molecule and atom stability.
Takeaways & Limitations
E-NFs provide a general-purpose equivariant flow for jointly modeling molecular structure and features in 3D without additional domain knowledge.
Takeaways & Limitations
ODE-based training is computationally expensive, rare loss peaks can cause divergence, and the model does not yet support datasets with edge data.
Abstract
from arXiv · showhide
This paper introduces a generative model equivariant to Euclidean symmetries: E(n) Equivariant Normalizing Flows (E-NFs). To construct E-NFs, we take the discriminative E(n) graph neural networks and integrate them as a differential equation to obtain an invertible equivariant function: a continuous-time normalizing flow. We demonstrate that E-NFs considerably outperform baselines and existing methods from the literature on particle systems such as DW4 and LJ13, and on molecules from QM9 in terms of log-likelihood. To the best of our knowledge, this is the first flow that jointly generates molecule features and positions in 3D.
1 Introduction
Equivariant methods have improved discriminative learning on 3D data, but efficient and powerful rotation-equivariant generative approaches remain limited. E-NFs address this gap by integrating an EGNN into a stable continuous-time flow for molecular generation.
- EGNNs are computationally cheap and effective for equivariant molecular regression and classification, but cannot generate new molecular structures.
- E-NFs are generative models for E(n)-equivariant data, including molecules represented in 3D.
- E-NFs use a continuous-time flow whose derivative is modeled by an EGNN adapted for stability.
- The method significantly outperforms non-equivariant variants and previous equivariant generative methods, while generating realistic QM9 molecules.
2 Background
Normalizing flows learn invertible transformations between simple base distributions and complex data distributions, enabling likelihood computation. Continuous-time flows express this transformation through neural-network dynamics, while E(n) equivariance preserves Euclidean symmetry for molecular coordinates and features.
- Normalizing Flows: Normalizing flows use a learnable invertible transformation from a simple base distribution to a complex data distribution.The inverse transformation enables exact likelihood computation through the change-of-variables formula.
- Normalizing Flows: Continuous-time normalizing flows define invertible transformations with a differential equation whose first-order derivative is predicted by neural-network dynamics.
- Normalizing Flows: Continuous-time flow training can estimate the Jacobian trace with Hutchinson’s estimator and integrate the dynamics numerically.Regularizing the dynamics is often used to improve training speed and solution stability.
- Equivariance in Normalizing Flows: An equivariant transformation applied to an invariant base distribution produces an invariant data distribution, and bijectivity transfers equivariance to the inverse.
- Equivariance: E(n) transformations consist of rotations, reflections, and translations that preserve Euclidean distances.For graph data, positions transform equivariantly while node features remain invariant and can still affect positions.
- E(n) Equivariant Graph Neural Networks: An EGNN stacks equivariant graph convolutional layers operating on node features, coordinates, and edge information.The resulting EGNN defines the dynamics of the ODE flow; when adjacency is unavailable, inferred soft edges act as attention over neighbors.
3 Related Work
Prior work includes equivariant neural networks, continuous-time normalizing flows, and specialized molecule-generation methods. E-NFs extend this landscape with a general flow for sets containing both positions and invariant features, supporting joint molecular modeling.
- Equivariant neural networks have been effective across discriminative tasks and dynamical-system modeling.
- Normalizing flows provide exact likelihood computation and can support fast inference and sampling, while continuous-time variants use first-order temporal dynamics.
- Existing molecule-generation methods variously use autoregressive generation, discrete graph structures, or position-only modeling.
- Prior equivariant flows target positional data, graphs, SU(n) symmetries, or other specialized settings.
- E-NFs provide a general-purpose flow for node sets containing positional and invariant features while remaining equivariant to E(n) transformations.
- The method jointly models molecular features and 3D structure without additional domain knowledge.
4 Method: E(n) Equivariant Normalizing Flows
E-NFs define an invertible, E(n)-equivariant transformation by integrating an EGNN as continuous-time dynamics, while jointly modeling positions and node features. The method handles translation-invariant position representations, discrete feature values, and Jacobian computation within this flow.
- Flow construction: E-NFs use an invariant latent distribution and an invertible equivariant map to generate graph nodes with positions x and features h.The model factors the latent distribution into positional and feature parts and maps latent representations to molecular data.
- Flow construction: The flow is defined as an ODE whose dynamics φ is predicted by an EGNN and integrated from data space to latent space.The transformation uses x(0)=x, h(0)=h and produces x(1)=zx, h(1)=zh.
- Dynamics: The EGNN outputs feature derivatives directly, while coordinate derivatives use the difference between the EGNN output and current coordinates.This makes coordinate dynamics rotate equivariantly and remain unaffected by translations.
- Dynamics: To stabilize ODE integration, E-NFs normalize relative coordinate differences by their norm plus a constant C, set to 1 in practice.The normalization is introduced because the original EGNN coordinate update can become unstable in an ODE.
- Symmetry constraints: Translation invariance is enforced by restricting positions to the zero-center-of-gravity subspace and removing the mean dynamics output.Under this condition, the ambient-space Jacobian determinant equals the determinant in the constrained subspace.
- Base distribution: The base distribution uses independent positional and feature components, with the positional Gaussian normalized according to the (M −1)n-dimensional subspace.The feature component can use a standard Gaussian because node features are already invariant to E(n) transformations.
- Discrete features: Ordinal and categorical node features are lifted to continuous variables using variational dequantization and argmax flows before flow training.Sampling reverses this lifting with rounding for ordinal features and argmax for categorical features, yielding a lower bound on the discrete log-likelihood.
5 Experiments
Experiments evaluate E-NF on synthetic particle systems and QM9 variants, measuring likelihood, distributional similarity, stability, and qualitative molecular structure. E-NF consistently achieves strong likelihood performance, while QM9 molecule experiments also assess validity-related outcomes and sample quality.
- DW4 and LJ13: DW4 and LJ13 evaluate equivariant flow performance across multiple training-set sizes using MCMC-generated particle configurations.Both systems possess translation, rotation, and reflection symmetries; validation and test sets each contain 1,000 samples.
- DW4 and LJ13: E-NF outperforms GNF, GNF-att, GNF-att-aug, Simple Dynamics, and Kernel Dynamics in test negative log-likelihood across all data regimes.The comparison averages reported numbers over three runs.
- QM9 Positional: QM9 Positional contains 19-atom molecules represented only by 3D coordinates, enabling comparison with methods that model positional data alone.Evaluation reports negative log-likelihood and Jensen–Shannon divergence between generated and training-set inter-atomic-distance histograms.
- QM9 Positional: E-NF achieves the best negative log-likelihood on QM9 Positional and converges faster than the other methods during optimization.E-NF and GNF-att-aug achieve the best Jensen–Shannon divergence.
- QM9 Molecules: QM9 molecule generation models positions, atom features, and molecule size, with likelihood computed using the continuous-time change-of-variables formula.The dataset uses 100K, 18K, and 13K training, validation, and test molecules, respectively.
- QM9 Molecules: E-NF captures sharp inter-atomic-distance distributions and generates 85% valid atoms and 4.9% stable molecules, compared with 75% and 0.5% for the best baseline.Among 10,000 generated examples, 491 stable molecules were obtained; 99.80% were unique and 93.28% were novel.
- QM9 Molecules: Qualitative samples are mostly structurally accurate, but random samples can contain disconnected molecules or atoms with incorrect bond counts.Stable-selected samples additionally visualize multiple bonds inferred from inter-atomic distances.
6 Limitations and Conclusions
E-NFs are generative models equivariant to Euclidean symmetries, using stable EGNN-parametrized continuous-time flows. They achieve strong likelihood and generation results, while retaining several computational, stability, scope, and reflection-related limitations.
- Limitations: ODE-based training is computationally expensive because solving the ODE requires multiple sequential forward operations.
- Limitations: The model has rare QM9 loss peaks, may diverge in one very rare case, lacks edge-data modeling, and assigns mirrored chiral molecules equal likelihood.
- Societal Impact: Molecule-generation methods may support drug research and material discovery, including new medications and catalyst materials for cheaper green-energy production.
- Conclusions: E-NFs are continuous-time normalizing flows that use an EGNN with improved stability as their parametrization.
- Conclusions: E-NFs considerably outperform existing normalizing flows in log-likelihood on DW4, LJ13, and QM9, and in generated-molecule and atom stability.
A Experiment details
The experiments instantiate E-NF and graph-flow dynamics through EGNN-related operations, with variants differing in architecture and equivariance. The appendix specifies each operation and baseline construction.
- EGNN components: The EGNN module comprises edge, coordinate, node, and edge-inference operations implemented with multilayer perceptrons and activation functions.
- Dynamics: E-NF, GNF, GNF-att, and GNF-att-aug use dynamics defined by different graph-based flow constructions.
- Dynamics: E-NF dynamics are the Section 2.1 EGNN dynamics adapted for the method in Section 4.
- Baselines: GNF treats dataset coordinates as node features and uses eij = 1 because the datasets have no adjacency matrices.
- Baselines: GNF-att infers edges as eij = φinf(mij), interpreting the mechanism as attention.
- Baselines: GNF-att-aug is identical to GNF-att except that node positions are rotated during preprocessing for data augmentation.
A.1 DW4 and LJ13 experiments
The DW4 and LJ13 appendix experiments use established particle-energy systems and report averaged negative log-likelihood results with implementation details and standard deviations.
- Datasets: DW4 contains four two-dimensional particles governed by an energy function with pairwise coupling and multiple metastable states.
- Datasets: LJ13 contains thirteen particles in three dimensions and uses a Lennard-Jones potential modeling repulsive and attractive interactions.
- Implementation: Training uses Adam, weight decay 10^-12, batch size 100, and learning rate 10^-3 except 5 · 10^-4 for E-NF.
- Results: DW4 and LJ13 results are averaged over three runs, with standard deviations included in the appendix tables.
- Results: DW4 and LJ13 appendix results report negative log-likelihood comparisons on test partitions across different training-sample amounts.
A.2 QM9 Positional and QM9
The QM9 appendix reports training settings and computational costs for QM9 and QM9 Positional flows. Training times vary substantially across methods and increase per epoch as ODE solving becomes more demanding.
- Implementation details: QM9 and QM9 Positional use batch size 128 and weight decay 10^-12, with reduced learning rates for E-NF and Simple Dynamics.
- Training cost: QM9 flows train for 30 epochs in approximately two weeks on two NVIDIA 1080Ti GPUs, while QM9 Positional flows train for 160 epochs on one GPU.
- Training cost: Simple Dynamics trains in less than a day, Kernel Dynamics in around two days, and other methods can take up to seven days.
- Training cost: Training becomes slower per epoch as performance improves because the ODE solver requires additional steps.
- Evaluation: QM9 results include negative log-likelihood averaged over three passes, with variance from dequantization and the trace estimator.
A.3 Stability of Molecules Benchmark
The QM9 stability benchmark evaluates molecules by checking whether atom–bond counts match prescribed valencies, with molecular stability requiring every atom to pass.
- A.3 Stability of Molecules Benchmark: An atom is stable when its number of bonds matches its valence.The listed valencies are H: 1, C: 4, N: 3, O: 2, and F: 1.
- A.3 Stability of Molecules Benchmark: A molecule is stable only when all of its atoms are stable.
- A.3 Stability of Molecules Benchmark: Bond types are determined from atom-pair distances and the types of atoms involved before stability is assessed.
A.4 Further QM9 analysis | Validity, Uniqueness, Novelty
The further QM9 analysis reports validity-related metrics for generated molecules and compares atom-count and atom-type distributions between stable generated and training molecules.
- A.4 Further QM9 analysis | Validity, Uniqueness, Novelty: Validity is the ratio of generated molecules that are valid.
- A.4 Further QM9 analysis | Validity, Uniqueness, Novelty: Uniqueness is the number of unique valid generated molecules divided by the number of all valid generated molecules.
- A.4 Further QM9 analysis | Validity, Uniqueness, Novelty: Novelty is the number of valid generated molecules absent from the training set divided by the total number of valid generated molecules.
- A.4 Further QM9 analysis | Validity, Uniqueness, Novelty: Because all stable molecules are valid in this analysis, only Uniqueness and Novelty are reported for 10.000 generated examples.
- A.4 Further QM9 analysis | Validity, Uniqueness, Novelty: Figure 6 compares the number of atoms per molecule and atom-type distributions for stable generated molecules and ground-truth molecules.
B Lifting Discrete Features to Continuous Space
The paper lifts discrete node features into continuous variables and uses deterministic rounding or argmax mappings, with learned dequantization distributions providing probabilistic inverses.
- B Lifting Discrete Features to Continuous Space: Discrete node features are separated into ordinal and categorical components, which share representation but receive different treatment.
- B Lifting Discrete Features to Continuous Space: Ordinal continuous features map to discrete values by rounding, while categorical features map by argmax.
- B Lifting Discrete Features to Continuous Space: The discrete-to-continuous transformation is deterministic, placing all conditional probability mass on the corresponding rounded or argmax event.
- B Lifting Discrete Features to Continuous Space: The variational inverse distribution must have support only where the deterministic discrete mapping assigns probability one.
- B Lifting Discrete Features to Continuous Space: For ordinal features, Gaussian noise is transformed through a sigmoid and added within (0, 1), preserving the rounded discrete value.
- B Lifting Discrete Features to Continuous Space: For categorical features, Gaussian noise is transformed with a softplus-based construction so the target category remains the argmax.