Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

SE(3) Equivariant Augmented Coupling Flows

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

TL;DR

This work proposes a coupling flow that preserves SE(3) and permutation equivariance by performing coordinate splits along additional augmented dimensions, and is the first to learn the full Boltzmann distribution of alanine dipeptide by only modeling the Cartesian positions of its atoms.

Abstract

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.

SE(3) Equivariant Augmented Coupling Flows

TL;DR

This work proposes a coupling flow that preserves SE(3) and permutation equivariance by performing coordinate splits along additional augmented dimensions, and is the first to learn the full Boltzmann distribution of alanine dipeptide by only modeling the Cartesian positions of its atoms.

Abstract

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.
Paper Structure (46 sections, 3 theorems, 36 equations, 8 figures, 8 tables)

This paper contains 46 sections, 3 theorems, 36 equations, 8 figures, 8 tables.

Table of Contents

  1. Introduction
  2. Background: coupling flows and invariant models
  3. Normalizing flows and coupling transforms
  4. Equivariance and invariance for coupling flow models of molecular conformations
  5. Method: SE3xSn equivariant augmented coupling flow model
  6. SE3 and permutation equivariant coupling transform
  7. SE3xSn Invariant base distribution
  8. Training and likelihood evaluation
  9. Experiments
  10. Training with samples: DW4, LJ13 and QM9 positional
  11. Training with samples: Alanine dipeptide
  12. Energy-based training: DW4 and LJ13
  13. Discussion and Related Work
  14. Limitations
  15. Conclusion
  16. ...and 31 more sections

Key Result

Proposition 3.1

Assume $q: \mathcal{X} \times \mathcal{A} \rightarrow \mathbb{R}_+$ is a G-invariant density over the probability space $(\mathcal{X} \times \mathcal{A}, \lambda_\mathcal{X} \otimes \lambda_\mathcal{A})$, then $q_{\bm{x}}\triangleq \int_\mathcal{A} q(\cdot, \bm{a}) \lambda_\mathcal{A}(\mathrm{d} \bm

Figures (8)

  • Figure 1: Illustration of the equivariant coupling layer of our augmented normalizing flow, where our variable with zero center of mass (CoM) $\tilde{\bm{x}}$ is transformed with the augmented variable $\bm{a}$.
  • Figure 2: Inter-atomic distances for samples from the train-data and Spherical-projE-ACF.
  • Figure 3: Ramachandran plots, i.e. marginal distribution of the dihedral angles $\phi$ and $\psi$ (see \ref{['sec:app_aldp_exp']}), obtained with MD (ground truth) and various models.
  • Figure 4: Illustration of the Cartesian and spherical projection.
  • Figure 5: Inter-atomic distances for samples from the train and test data on DW4 and LJ13. Both datasets are biased with visible differences between the train and test sets.
  • ...and 3 more figures

Theorems & Definitions (5)

  • Proposition 3.1: Invariant marginal
  • proof
  • Proposition 3.2: Equivariant augmented coupling flow
  • proof
  • Proposition A.1: Disintegration of measures for translation invariance pollard2002useryim2023SE