E(n) Equivariant Normalizing Flows

TL;DR

This work introduces E(n) Equivariant Normalizing Flows (E-NFs), a continuous-time normalizing flow whose dynamics are parameterized by an E(n) equivariant graph neural network (EGNN) to generate 3D molecular structures with both coordinates and invariant features. By centering coordinates to enforce translation invariance, lifting discrete features via variational dequantization, and using a subspace-based base distribution, E-NFs achieve exact likelihoods while maintaining Euclidean symmetry. Empirically, E-NFs outperform non-equivariant variants and prior equivariant flows on DW4, LJ13, and QM9 datasets in log-likelihood and molecule stability, and demonstrate joint generation of atom types, charges, and 3D positions. The approach holds promise for efficient, symmetry-aware molecular generation, with caveats around computational cost and enantioselectivity limitations, guiding future improvements in efficiency and model expressivity.

Abstract

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.

Figures (6)

• Figure 1: Overview of our method in the sampling direction. An equivariant invertible function $g_\theta$ has learned to map samples from a Gaussian distribution to molecules in 3D, described by ${\mathbf{x}}, {\mathbf{h}}$.
• Figure 2: Overview of the training procedure: The discrete $\mathbf{h}$ is lifted to continuous $\boldsymbol{h}$. Then the variables ${\mathbf{x}}, \boldsymbol{h}$ are transformed by an ODE to ${\mathbf{z}}_x, {\mathbf{x}}_h$. To get a lower bound on $\log p_V({\mathbf{x}}, {\mathbf{h}})$ we sum the variational term $-\log q(\boldsymbol{h} | \mathbf{h})$, the volume term from the ODE $\int_0^1 \mathrm{Tr}\, J_\phi(\mathbf{z}(t)) \mathrm{d}t$, the log-likelihood of the latent representation on a Gaussian $\log p_Z({\mathbf{z}}_x, {\mathbf{x}}_h)$, and the log-likelihood of the molecule size $\log p_\mathrm{M}(M)$. To train the model, the sum of these terms is maximized.
• Figure 3: Normalized histogram of relative distances between atoms for QM9 Positional and E-NF generated samples.
• Figure 4: The table on the left presents the Negative Log Likelihood (NLL) $-\log p_V({\mathbf{x}})$ for the QM9 Positional dataset on the test data. The figure on the right shows the training curves for all methods.
• Figure 5: Sampled molecules by our E-NF. The top row contains random samples, the bottom row also contains samples but selected to be stable. Edges are drawn depending on inter-atomic distance.