# SE(3) Equivariant Augmented Coupling Flows

arXiv: https://arxiv.org/abs/2308.10364 · HTML: https://mlsanity.com/paper/2308.10364/

## Key facts

- Cited by 37 papers in the map (data through 14 Oct 2026).
- Builds on 40 of its 55 references that are in the map.
- First posted on arXiv 20 Aug 2023.

## Builds on (most cited first, 30 of 40)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Density estimation using Real NVP](https://mlsanity.com/paper/1605.08803/) (2016)
- [Consistency Models](https://mlsanity.com/paper/2303.01469/) (2023)
- [Normalizing Flows for Probabilistic Modeling and Inference](https://mlsanity.com/paper/1912.02762/) (2019)
- [Variational Diffusion Models](https://mlsanity.com/paper/2107.00630/) (2021)
- [NICE: Non-linear Independent Components Estimation](https://mlsanity.com/paper/1410.8516/) (2014)
- [E(n) Equivariant Graph Neural Networks](https://mlsanity.com/paper/2102.09844/) (2021)
- [Improving and generalizing flow-based generative models with minibatch optimal transport](https://mlsanity.com/paper/2302.00482/) (2023)
- [Tensor field networks: Rotation- and translation-equivariant neural networks for 3D point clouds](https://mlsanity.com/paper/1802.08219/) (2018)
- [E(3)-Equivariant Graph Neural Networks for Data-Efficient and Accurate Interatomic Potentials](https://mlsanity.com/paper/2101.03164/) (2021)
- [Equivariant Diffusion for Molecule Generation in 3D](https://mlsanity.com/paper/2203.17003/) (2022)
- [SE(3)-Transformers: 3D Roto-Translation Equivariant Attention Networks](https://mlsanity.com/paper/2006.10503/) (2020)
- [Neural Spline Flows](https://mlsanity.com/paper/1906.04032/) (2019)
- [FFJORD: Free-form Continuous Dynamics for Scalable Reversible Generative Models](https://mlsanity.com/paper/1810.01367/) (2018)
- [Tackling the Generative Learning Trilemma with Denoising Diffusion GANs](https://mlsanity.com/paper/2112.07804/) (2021)
- [GeoDiff: a Geometric Diffusion Model for Molecular Conformation Generation](https://mlsanity.com/paper/2203.02923/) (2022)
- [MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields](https://mlsanity.com/paper/2206.07697/) (2022)
- [Steerable CNNs](https://mlsanity.com/paper/1612.08498/) (2016)
- [Equiformer: Equivariant Graph Attention Transformer for 3D Atomistic Graphs](https://mlsanity.com/paper/2206.11990/) (2022)
- [Torsional Diffusion for Molecular Conformer Generation](https://mlsanity.com/paper/2206.01729/) (2022)
- [Learning Local Equivariant Representations for Large-Scale Atomistic Dynamics](https://mlsanity.com/paper/2204.05249/) (2022)
- [Equivariant Flows: Exact Likelihood Generative Learning for Symmetric Densities](https://mlsanity.com/paper/2006.02425/) (2020)
- [SE(3) diffusion model with application to protein backbone generation](https://mlsanity.com/paper/2302.02277/) (2023)
- [Path Integral Sampler: a stochastic control approach for sampling](https://mlsanity.com/paper/2111.15141/) (2021)
- [e3nn: Euclidean Neural Networks](https://mlsanity.com/paper/2207.09453/) (2022)
- [Spherical Message Passing for 3D Graph Networks](https://mlsanity.com/paper/2102.05013/) (2021)
- [E(n) Equivariant Normalizing Flows](https://mlsanity.com/paper/2105.09016/) (2021)
- [Stochastic Normalizing Flows](https://mlsanity.com/paper/2002.06707/) (2020)

## Built on by (most cited first, 30 of 37)

- [SE(3)-Stochastic Flow Matching for Protein Backbone Generation](https://mlsanity.com/paper/2310.02391/) (2023)
- [Iterated Denoising Energy Matching for Sampling from Boltzmann Densities](https://mlsanity.com/paper/2402.06121/) (2024)
- [Adjoint Sampling: Highly Scalable Diffusion Samplers via Adjoint Matching](https://mlsanity.com/paper/2504.11713/) (2025)
- [Beyond ELBOs: A Large-Scale Evaluation of Variational Methods for Sampling](https://mlsanity.com/paper/2406.07423/) (2024)
- [Generative Modeling of Molecular Dynamics Trajectories](https://mlsanity.com/paper/2409.17808/) (2024)
- [Transferable Boltzmann Generators](https://mlsanity.com/paper/2406.14426/) (2024)
- [Scalable Equilibrium Sampling with Sequential Boltzmann Generators](https://mlsanity.com/paper/2502.18462/) (2025)
- [A Survey of Geometric Graph Neural Networks: Data Structures, Models and Applications](https://mlsanity.com/paper/2403.00485/) (2024)
- [Training Neural Samplers with Reverse Diffusive KL Divergence](https://mlsanity.com/paper/2410.12456/) (2024)
- [EdgeBench: Unveiling Scaling Laws of Learning from Real-World Environments](https://mlsanity.com/paper/2607.05155/) (2026)
- [Temperature-Annealed Boltzmann Generators](https://mlsanity.com/paper/2501.19077/) (2025)
- [Amortized Sampling with Transferable Normalizing Flows](https://mlsanity.com/paper/2508.18175/) (2025)
- [Accelerating 3D Molecule Generation via Jointly Geometric Optimal Transport](https://mlsanity.com/paper/2405.15252/) (2024)
- [Scalable Boltzmann Generators for equilibrium sampling of large-scale materials](https://mlsanity.com/paper/2509.25486/) (2025)
- [Neural Approximate Mirror Maps for Constrained Diffusion Models](https://mlsanity.com/paper/2406.12816/) (2024)
- [Learning Boltzmann Generators via Constrained Mass Transport](https://mlsanity.com/paper/2510.18460/) (2025)
- [Energy-Weighted Flow Matching: Unlocking Continuous Normalizing Flows for Efficient and Scalable Boltzmann Sampling](https://mlsanity.com/paper/2509.03726/) (2025)
- [Efficient and Unbiased Sampling from Boltzmann Distributions via Variance-Tuned Diffusion Models](https://mlsanity.com/paper/2505.21005/) (2025)
- [MarS-FM: Generative Modeling of Molecular Dynamics via Markov State Models](https://mlsanity.com/paper/2509.24779/) (2025)
- [Boltzmann generators for amorphous particle systems](https://mlsanity.com/paper/2512.16607/) (2025)
- [Conditional Normalizing Flows for Active Learning of Coarse-Grained Molecular Representations](https://mlsanity.com/paper/2402.01195/) (2024)
- [Fast and Unified Path Gradient Estimators for Normalizing Flows](https://mlsanity.com/paper/2403.15881/) (2024)
- [Conditional Synthesis of 3D Molecules with Time Correction Sampler](https://mlsanity.com/paper/2411.00551/) (2024)
- [HollowFlow: Efficient Sample Likelihood Evaluation using Hollow Message Passing](https://mlsanity.com/paper/2510.21542/) (2025)
- [On the design space between molecular mechanics and machine learning force fields](https://mlsanity.com/paper/2409.01931/) (2024)
- [Efficient Training of Boltzmann Generators Using Off-Policy Log-Dispersion Regularization](https://mlsanity.com/paper/2602.03729/) (2026)
- [Verlet Flows: Exact-Likelihood Integrators for Flow-Based Generative Models](https://mlsanity.com/paper/2405.02805/) (2024)
- [Importance Weighted Score Matching for Diffusion Samplers with Enhanced Mode Coverage](https://mlsanity.com/paper/2505.19431/) (2025)
- [SESaMo: Symmetry-Enforcing Stochastic Modulation for Normalizing Flows](https://mlsanity.com/paper/2505.19619/) (2025)
- [Probabilistic Quantum SVM Training on Ising Machine](https://mlsanity.com/paper/2503.16363/) (2025)

## Abstract (arXiv)

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.