PaperlayerView exact revision on arXiv

Source-linked AI summary

SE(3) Equivariant Augmented Coupling Flows

Laurence I. Midgley, Vincent Stimper, Javier Antorán, Emile Mathieu, Bernhard Schölkopf, José Miguel Hernández-Lobato

arXiv:2308.10364v6cs.LGphysics.comp-ph

TL;DR

Standard coupling flows do not provide the Cartesian-coordinate symmetries needed for molecular distributions while retaining fast sampling and density evaluation. This paper builds an augmented equivariant coupling flow that transforms invariant representations of atom positions. It reports competitive performance, substantially faster sampling than equivariant CNFs and diffusion models, and energy-based training on particle systems.

  • Problem

    Molecular Boltzmann modeling needs SE(3) and permutation symmetry alongside fast sampling and density evaluation, but these properties are not jointly available in existing approaches.

    Source: §1 · p. 2

  • Method

    The flow updates Cartesian atom positions by transforming an invariant representation and mapping the result back to equivariant coordinates.

    Source: §1 · p. 2, §3.1 · p. 5

  • Results

    The flow achieves performance similar to CNFs while sampling and evaluating density more than an order of magnitude faster, and supports energy-based training on DW4 and LJ13.

    Source: §1 · p. 2, §6 · p. 10

  • Takeaways & Limitations

    The flow can learn full molecular Boltzmann distributions in Cartesian coordinates and use only target unnormalized densities on particle systems where internal coordinates are inadequate.

    Source: §1 · p. 2, §6 · p. 10

    Despite being faster than CNFs, the flow is more computationally expensive and less numerically stable than internal-coordinate flows.

    Source: §5 · p. 10

Abstract

from arXiv · show

Coupling normalizing flows allow for fast sampling and density evaluation, making them the tool of choice for probabilistic modeling of physical systems. However, the standard coupling architecture precludes endowing flows that operate on the Cartesian coordinates of atoms with the SE(3) and permutation invariances of physical systems. This work proposes a coupling flow that preserves SE(3) and permutation equivariance by performing coordinate splits along additional augmented dimensions. At each layer, the flow maps atoms' positions into learned SE(3) invariant bases, where we apply standard flow transformations, such as monotonic rational-quadratic splines, before returning to the original basis. Crucially, our flow preserves fast sampling and density evaluation, and may be used to produce unbiased estimates of expectations with respect to the target distribution via importance sampling. When trained on the DW4, LJ13, and QM9-positional datasets, our flow is competitive with equivariant continuous normalizing flows and diffusion models, while allowing sampling more than an order of magnitude faster. Moreover, to the best of our knowledge, we are the first to learn the full Boltzmann distribution of alanine dipeptide by only modeling the Cartesian positions of its atoms. Lastly, we demonstrate that our flow can be trained to approximately sample from the Boltzmann distribution of the DW4 and LJ13 particle systems using only their energy functions.

1 Introduction

Boltzmann generators can sample molecular equilibrium distributions independently, but existing approaches lack the combination of Cartesian-coordinate operation, SE(3) equivariance, and fast sampling. The paper proposes an equivariant coupling flow and reports competitive molecular-generation performance, faster sampling, and energy-based training on particle systems.

1 Introduction

  • Boltzmann generators provide independent samples whose expectations can be debiased by importance weighting, while fast sampling and density evaluation support energy-based training.Energy-based training uses the Boltzmann distribution’s unnormalized density and avoids constraining model quality by the quality and quantity of MD samples.Source: §1 · pp. 1–2
  • Existing molecular coupling flows use internal coordinates that depend on molecular graphs, can struggle with distant interactions, and cannot capture some atom-permutation symmetries.Cartesian-coordinate CNFs and diffusion models incorporate SE(3) equivariance, but typically require thousands of neural-network evaluations for sampling and density evaluation.Source: §1 · p. 2
  • The flow extends coupling layers with auxiliary variables and updates atom positions by transforming an SE(3)-invariant representation before mapping back to equivariant coordinates.The invariant-space transform can be a standard normalizing flow transformation.Source: §1 · p. 2
  • The proposed flow matches existing SE(3) continuous normalizing flows, diffusion models, and internal-coordinate coupling flows on molecular generation tasks, while sampling more than 10 times faster than SE(3) alternatives.The authors also report learning alanine dipeptide’s full Boltzmann distribution solely in Cartesian coordinates.Source: §1 · p. 2
  • Energy-based training on DW4 and LJ13 uses only their molecular energy functions, where internal-coordinate flows miss permutation invariance and CNFs or diffusion models are computationally intractable.The paper presents its flow as the only existing permutation- and SE(3)-equivariant method tractable for these settings.Source: §1 · p. 2

2 Background: coupling flows and invariant models

Normalizing flows transform a base density through an invertible map, and coupling transforms make this practical through conditional coordinate splits. Molecular flow models must also respect rigid-motion and atom-permutation symmetries, which standard coupling splits cannot jointly preserve.

2.1 Normalizing flows and coupling transforms

  • A coupling transform splits coordinates, transforms one set conditioned on the unchanged set, and uses an easily computed triangular Jacobian determinant for efficient sampling and density evaluation.A normalizing flow defines a density by pushing a base density through an invertible map.Source: §2.1 · pp. 2–3

2.2 Equivariance and invariance for coupling flow models of molecular conformations

  • Molecular coordinates are arranged as a 3 × n matrix, with columns indexing particles, and molecular energies are invariant to SE(3) transformations and permutations of same-type atoms.The symmetry group S_n acts by permuting the atom columns.Source: §2.2 · p. 3
  • An equivariant map commutes with a group action, whereas an invariant map returns the same output after that action.An invariant base density combined with an equivariant invertible transform yields an invariant flow density.Source: §2.2 · p. 3
  • No coupling transform can preserve both particle-permutation and rotational equivariance because splitting across particles or spatial dimensions breaks one of those symmetries.The paper frames this challenge for SO(3) × S_n invariant flow models.Source: §2.2 · p. 3
  • There is no translation-invariant probability measure on the full coordinate space, so the construction models the SO(3)-invariant zero-center-of-mass subspace instead.The probability measure is disintegrated into a translational Lebesgue measure and an SO(3)-invariant probability measure.Source: §2.2 · p. 3

3 Method: SE(3) × Sn equivariant augmented coupling flow model

The method builds an SE(3) × Sn equivariant coupling flow on zero-center-of-mass coordinates, using augmented variables to condition flexible transformations in learned invariant spaces. It combines this architecture with an invariant base distribution and sample- or energy-based training and density evaluation.

3.1 SE(3) and permutation equivariant coupling transform

  • The core coupling projects atom coordinates into a learned invariant space, applies a standard invertible flow transform, then projects the result back to Cartesian space.Equivariant reference vectors and invariant parameters are produced by a graph neural network.Source: §3.1 · p. 5
  • Three projection strategies are defined: CARTESIAN-PROJ transforms projected Cartesian coordinates, SPHERICAL-PROJ transforms spherical coordinates, and VECTOR-PROJ transforms radius alone.VECTOR-PROJ holds the polar and azimuth angles constant.Source: §3.1 · p. 5
  • ShiftCoM swaps the center of mass between observed and augmented variables, lifting zero-CoM observations for ambient-space coupling before mapping them back.The shift is invertible and has unit Jacobian determinant.Source: §3.1 · pp. 4–6

3.2 SE(3) × Sn Invariant base distribution

  • The invariant base uses a zero-CoM Gaussian for positions and a Gaussian over augmented variables centered on the positions, yielding joint translation invariance.The isotropic base density is also invariant to rotations, reflections, and permutations.Source: §3.2 · p. 6

3.3 Training and likelihood evaluation

  • With samples, training maximizes joint likelihood; with energy queries, FAB minimizes joint α-divergence at α=2 rather than using mode-seeking reverse KL.The joint-likelihood objective is a lower bound on marginal log-likelihood up to a fixed constant.Source: §3.3 · pp. 6–7
  • Importance weighting provides unbiased target-expectation estimates and marginal-density estimates, whose variance vanishes when q(a|x)=π(a|x).Minimizing the joint α-divergence with α=2 corresponds to minimizing the variance of joint importance weights.Source: §3.3 · p. 7

4 Experiments

Across sample-trained particle and molecular tasks, E-ACF is competitive with equivariant CNFs and diffusion models while sampling substantially faster. It also learns alanine dipeptide in Cartesian coordinates and approximates DW4 and LJ13 Boltzmann distributions with energy-based training.

4.1 Training with samples: DW4, LJ13 and QM9 positional

  • More than an order of magnitude faster sampling, E-ACF remains competitive with equivariant CNFs and diffusion models on DW4, LJ13, and QM9-positional.The experiments use fully permutation-invariant positional data; the CARTESIAN-PROJ variant required a lower learning rate on QM9-positional because of numerical instability.Source: §4.1 · p. 7, Table 1 · p. 8

4.2 Training with samples: Alanine dipeptide

  • E-ACF variants outperform NON-E-ACF on alanine dipeptide, and the best Cartesian-coordinate models are competitive with the internal-coordinate flow.The internal-coordinate flow is only marginally better in NLL and KLD, while E-ACF performs better on ESS; the E-ACF models are the first reported to learn the full Boltzmann distribution using only Cartesian coordinates.Source: Table 2 · p. 8, §4.2 · pp. 8–9, Figure 3 · p. 9

4.3 Energy-based training: DW4 and LJ13

  • FAB-trained E-ACF approximates the DW4 and LJ13 Boltzmann distributions with reasonably high joint ESS and NLL comparable to maximum-likelihood flows.NON-E-ACF performs worse on both ESS and NLL; maximum-likelihood models have lower ESS, and their training uses a relatively small, biased dataset.Source: Table 3 · p. 9, §4.3 · pp. 9–10

5 Discussion and Related Work

The discussion contrasts the proposed flow with alternative equivariant approaches and identifies a computational trade-off. It also points to extensions beyond full-molecule Cartesian modeling.

5 Discussion and Related Work

  • Klein et al. [2023a] use data augmentation for approximate rotation invariance, which performed significantly worse than intrinsic invariance in the authors’ experiments.The related model provides translation and permutation equivariance; the authors say their projection-based transform could make it fully invariant.Source: §5 · p. 10
  • Residual flows require fixed-point iteration for training and inversion, which is costly and may interact poorly with energy-based training methods such as FAB [Midgley et al., 2023].Bose et al. [2022] also found spectral normalization did not interact well with their equivariant CNN architecture.Source: §5 · p. 10
  • The projection framework could be extended from full-molecule Cartesian coordinates to rigid bodies, including protein backbones and many-body systems [Jumper et al., 2021, Yim et al., 2023, Köhler et al., 2023c].The authors also note recent work on faster CNF and diffusion sampling and energy-based sampling, suggesting these models may become competitive with discrete-time flows.Source: §5 · p. 10
  • Each flow layer requires an expensive EGNN pass, making this method more computationally expensive than flows on internal coordinates.The authors note that efficiency is especially important for energy-based training.Source: §5 · p. 10

6 Conclusion

The conclusion highlights the flow’s sampling and density-evaluation speed and its use for energy-based Boltzmann modeling. The appendix states tensor and measure concepts used in the paper’s mathematical setup.

6 Conclusion

  • The proposed flow matches CNF performance under maximum likelihood while enabling sampling and density evaluation more than an order of magnitude faster.It can be trained from an unnormalized target density on problems where internal coordinates are inadequate because of permutation symmetries.Source: §6 · p. 10

A.1 Tensors and representation

  • A representation assigns each group element an invertible linear map on a vector space, respecting group composition through ρ(gh) = ρ(g)ρ(h).The appendix describes tensors as objects acted on by these representations, including scalars, vectors, and higher-order tensors.Source: §A.1 · p. 15

A.2 Translational invariant densities

  • For an SE(3)-invariant measure, the appendix states a decomposition into Lebesgue measure over translations and an SO(3)-invariant probability measure on zero-center-of-mass configurations.The converse motivates constructing an SO(3)-invariant probability measure on the zero-center-of-mass subspace.Source: §A.2 · p. 15

B Method

The appendix explains why transformations on zero-center-of-mass variables can have ill-defined Jacobians and describes augmented-space operations, projections, and objectives used by the flow. It also states how these objectives relate to marginal likelihoods and divergences.

B.1 Transforming variables on a subspace

  • A transform can leave zero-center-of-mass inputs unchanged yet have an arbitrary log-determinant, permitting a degenerate maximum-likelihood solution.Scaling the center-of-mass direction by arbitrary s changes the log-determinant by log s, despite acting trivially on the zero-CoM input.Source: §B.1 · p. 15

B.2 Jacobian of ShiftCoM

  • ShiftCoM rearranges a zero-CoM vector and an augmented vector while preserving their combined dimensionality and having unit Jacobian log-determinant.The appendix says it avoids the prior issue because it does not transform the center of mass.Source: §B.2 · p. 16

B.3 Projection variants and their respective transformations

  • The layer projects atoms through equivariant local references into invariant coordinates, applies a standard flow transform, and maps the result back; Cartesian, spherical, and radial-only variants are described.Spherical coordinates use monotonic rational-quadratic splines for radius and circular splines for angles; vector projection transforms only the radial coordinate and is parity equivariant.Source: §B.3 · pp. 16, 18

B.4 Multiple augmented variables

  • With multiple augmented variables, coupling blocks update groups of variables conditionally while using local reference frames and ensuring transformed arrays have non-zero center of mass.The augmented variables follow conditional normal distributions centered at x.Source: §B.4 · p. 18

B.5 Invariant base and target distributions

  • An isotropic Gaussian augmented distribution centered at x satisfies the stated SE(3) and permutation invariance constraint.The appendix gives π(a|x) = N(a; x, η²I) as the example choice.Source: §B.5 · pp. 18–19

B.6 Maximum likelihood objective

  • The maximum-likelihood objective maximizes joint log-density over observations and augmented variables, lower-bounding expected marginal log-density, with tightness when q(a|x) = π(a|x).The appendix notes that a term constant with respect to flow parameters is dropped from training.Source: §B.6 · p. 19

B.7 Reverse KL objective; B.8 Alpha divergence objective

  • The joint reverse-KL objective upper-bounds the marginal reverse KL, and the appendix states analogous upper-bound behavior for the α=2 divergence objective.The reverse-KL bound’s looseness is the expected conditional KL; the α-divergence bound also becomes tight when q(a|x) = π(a|x).Source: §B.7 · p. 19, §B.8 · p. 19

B.9 Importance weighted estimator

  • The appendix describes an importance-weighted estimator as unbiased under its stated support condition, with a log-marginal-likelihood estimate that becomes accurate as M grows.The support condition is that q(x,a) is zero whenever π(a|x) is zero.Source: §B.9 · p. 20

B.10 Effective sample size

The appendix defines reverse and forward effective sample sizes and outlines a conversion from internal-coordinate densities to zero-center-of-mass Cartesian densities that accounts for global rotation.

B.10 Effective sample size

  • Reverse ESS uses model samples and forward ESS uses target samples; the reported joint ESS for augmented flows lower-bounds marginal ESS.The appendix gives the marginal ESS formulas and says joint estimators substitute (x, a) for x.Source: §B.10 · p. 20

B.11 Conversion of internal coordinate flow density to density on ZeroCoM Cartesian space with global rotation represented

  • Converting the resulting density to zero-CoM Cartesian space requires accounting for the Jacobians of the rotation and coordinate transformations.The appendix gives the internal-coordinate-to-zero-CoM log-density conversion, including the rotation and centering corrections.Source: §B.11 · p. 21
  • The internal-coordinate representation fixes the first atom at the origin, the second on the x-axis, and the third in the xy plane, yielding rotation-invariant Cartesian coordinates.The appendix adds global rotation degrees of freedom by sampling a rotation uniformly over SO(3).Source: §B.11 · p. 21

C Experiments: Details and Further Results

The appendix documents experimental resources and a custom optimizer intended to improve training stability. The optimizer skips or clips unusually large gradient updates and can optionally mask extreme-loss samples.

C Experiments: Details and Further Results

C.1 Custom Optimizer

  • Gradient steps are skipped when their norm exceeds 20 times the median of the previous 100 norms, and clipped at five times the median otherwise.The resulting gradients are used for an Adam update.Source: §C.1 · p. 22
  • An optional setting masks gradient contributions from batch samples with very large negative log probability to further improve training stability.Source: §C.1 · p. 22

C.2 EGNN implementation

The EGNN implementation adds stability and supports multiple equivariant position or vector channels, including augmented variables.

C.2 EGNN implementation

  • Applying a softmax to invariant features output by the EGNN helped improve stability.Source: §C.2 · p. 22
  • The implementation extends the EGNN to allow multiple position or vector channels at each node, with interactions across channels.Source: §C.2 · p. 22
  • Multiple channels are needed when local reference frames use multiple equivariant vectors or when multiple augmented-variable sets provide positions to the EGNN.Source: §C.2 · p. 22

C.3.1 DW4, LJ13 and QM9 Positional

Across DW4 and LJ13, the reported performance gap between equivariant and non-equivariant flows widens with less training data. The appendix also describes dataset properties, model configurations, and evaluation summaries.

C.3.1 DW4, LJ13 and QM9 Positional

  • DW4 and LJ13 use 1,000 training and test samples each, and both datasets show visible train–test differences.QM9-positional uses 13,831 training, 2,501 validation, and 1,813 test samples.Source: §C.3.1 · p. 22, Figure 5 · p. 22
  • Flow experiments use 12 blocks, while augmented-variable counts and transformations vary by dataset and projection variant.The appendix also specifies GNN, MLP, transformer, learning-rate, batch-size, and training-epoch settings.Source: §C.3.1 · pp. 22–23
  • The E-CNF and diffusion baselines use dataset-specific architectures and training schedules; density evaluation uses exact divergence for DW4 and LJ13 but an approximate trace estimator for QM9-positional.Diffusion sampling and density evaluation use the dopri5 ODE solver with a maximum of 1,000 evaluations.Source: §C.3.1 · p. 23
  • The appendix reports training runtimes in hours for DW4, LJ13, and QM9-positional, averaged over three seeded runs with standard errors.The runs used a GeForce RTX 2080Ti, and runtimes include intermittent evaluation.Source: §C.3.1 · p. 22, Table 4 · p. 23
  • Table 5 reports joint effective sample size (%) for flows trained by maximum likelihood with respect to the joint target distribution p(x)π(a|x), using 10,000 samples averaged over 3 seeded runs; standard error is reported as uncertainty.Source: §C.3.1 · p. 23, Table 5 · p. 24

Varying the number of training samples for DW4 and LJ13

  • The performance gap between equivariant and non-equivariant flows is larger in the lower-data regime on DW4 and LJ13.Table 6 reports test-set negative log likelihood across training-set sizes.Source: §C.3.1 · p. 24, Table 6 · p. 24

C.3.2 Alanine dipeptide

The alanine dipeptide experiments examine generalization across training-set sizes and importance reweighting at 800 K and 300 K. The appendix also specifies chirality handling and evaluation settings.

C.3.2 Alanine dipeptide

  • The study models a 22-atom alanine dipeptide at 800 K using replica-exchange MD data, and the ACF variants use 20 layers.The molecule’s ϕ and ψ dihedral angles are used in Ramachandran plots; the ACF variants were trained for 50 epochs.Source: Figure 6 · p. 24, §C.3.2 · p. 24
  • Parity-equivariant VECTOR-PROJ E-ACF, E-CNF, and E-CNF-DIFF generate both chiral forms, so evaluation filters samples to the L-form.Source: §C.3.2 · p. 25
  • Because E-CNF and E-CNF-DIFF likelihood evaluation is expensive, those models were evaluated on only 10% of the test set, or 10^6 datapoints.Source: §C.3.2 · p. 25

Limiting the number of samples

  • Equivariant models outperform the non-equivariant counterpart at every training-set size; full-data NON-E-ACF trails selected E-ACFs trained on only 1–10% of the data.Table 7 reports negative log likelihood for training subsets of different sizes.Source: §C.3.2 · p. 25, Table 7 · p. 25

Reweighting

  • At 800 K, reweighting SPHERICAL-PROJ E-ACF samples reduces Ramachandran-plot KLD from 1.88 · 10^-3 to 1.39 · 10^-3.The reweighted plot is closer to ground truth in low-density regions.Source: §C.3.2 · p. 25, Figure 7 · p. 25

Alanine dipeptide at T = 300K

  • At 300 K, reweighting reduces KLD from 2.79 · 10^-3 to 1.97 · 10^-3; the internal-coordinate flow reports (7.57 ± 3.80) · 10^-3.The internal-coordinate models had about six times fewer parameters than the authors’ model.Source: §C.3.2 · p. 26, Figure 8 · p. 26

C.4 Energy-based training

The energy-based training experiments use FAB with specified replay-buffer, importance-weight, and AIS settings, alongside separate flow-training setups for DW4 and LJ13.

C.4 Energy-based training

  • DW4 and LJ13 use batch sizes of 128 and 1024, respectively, with E-ACF training for 20,000 and 14,000 epochs and NON-E-ACF training for 60,000 and 56,000 epochs.The NON-E-ACF epoch counts correspond to roughly the same training time as SPHERICAL-PROJ E-ACF because NON-E-ACF is faster per iteration.Source: §C.4 · p. 26
  • Reverse ESS is estimated from 10,000 flow samples, while forward ESS is estimated using the test sets.Source: §C.4 · p. 26
  • FAB uses replay buffers of 64–128 batches, clips importance-weight adjustment at 10, and performs eight buffer-sampling and parameter-update steps per AIS forward pass.AIS uses HMC with five leapfrog steps and two intermediate distributions for DW4 or eight for LJ13; the HMC step size targets 65% acceptance.Source: §C.4 · p. 26
  • Table 8 reports energy-based FAB runtimes for DW4 on a GeForce RTX 2080Ti GPU and LJ13 on a V-4 TPU, averaged across three seeded runs with standard errors.The reported runtimes include evaluations performed intermittently during training.Source: Table 8 · p. 26

Related papers

Loading 2308.10364v6…