PaperlayerView exact revision on arXiv

Source-linked AI summary

Rigid Body Flows for Sampling Molecular Crystal Structures

Jonas Köhler, Michele Invernizzi, Pim de Haan, Frank Noé

arXiv:2301.11355v4cs.LGphysics.chem-phphysics.comp-phstat.ML

TL;DR

Molecular normalizing flows face representation challenges, especially when rigid bodies make the target density singular in the full coordinate space. This paper develops smooth quaternion-based flows for molecular poses and evaluates them on a multimodal tetrahedral system and ice XI, while identifying phase transitions and non-crystal phases as open challenges.

  • Problem

    Existing molecular-flow representations can be non-scalable, non-transferable, or unnatural, and rigid residues make the target density singular in the full space.

    Source: Building flows on natural molecular… · p. 2

  • Method

    The paper constructs flip-equivariant flows on S3, including smooth projective convex gradient maps, to model rigid-body molecular poses.

    Source: §3.1 · p. 5, §3.2 · p. 6

  • Results

    Projective convex gradient maps faithfully reconstruct the tetrahedral system’s multimodal rotation density, while ice XI flow samples achieve close energy-distribution overlap and nearly perfect reweighted energy and oxygen-oxygen radial-distribution overlap.

    Source: §4.1 · p. 7, §4.2 · p. 8

  • Takeaways & Limitations

    The approach models positions and orientations of interacting molecules, and its pose architecture can be extended with separate flows for molecular internal degrees of freedom.

    Source: §5 · pp. 8–9

    Handling phase transitions or extending beyond the crystal phase to liquids and gases remains a major challenge for the present architecture.

    Source: §5 · pp. 8–9

Abstract

from arXiv · show

Normalizing flows (NF) are a class of powerful generative models that have gained popularity in recent years due to their ability to model complex distributions with high flexibility and expressiveness. In this work, we introduce a new type of normalizing flow that is tailored for modeling positions and orientations of multiple objects in three-dimensional space, such as molecules in a crystal. Our approach is based on two key ideas: first, we define smooth and expressive flows on the group of unit quaternions, which allows us to capture the continuous rotational motion of rigid bodies; second, we use the double cover property of unit quaternions to define a proper density on the rotation group. This ensures that our model can be trained using standard likelihood-based methods or variational inference with respect to a thermodynamic target density. We evaluate the method by training Boltzmann generators for two molecular examples, namely the multi-modal density of a tetrahedral system in an external field and the ice XI phase in the TIP4P water model. Our flows can be combined with flows operating on the internal degrees of freedom of molecules and constitute an important step towards the modeling of distributions of many interacting molecules.

1. Introduction

Boltzmann generators use normalizing flows to sample molecular equilibrium distributions, but existing representations and rigid-body settings pose challenges for scalable, physically meaningful modeling. The paper introduces flows over molecular positions and orientations that use the quaternion double cover and keep external poses separate from internal degrees of freedom.

1. Introduction

  • Boltzmann generators sample molecular conformations at equilibrium and can connect metastable states to compute their free-energy differences.They combine maximum-likelihood training on potentially biased molecular-dynamics data with energy-based reverse-KL training.Source: §1 · p. 2

Building flows on natural molecular representations

  • Prior molecular flows use representations that can be non-scalable, non-transferable, or unnatural, while equivariant all-atom approaches can be computationally intensive and difficult to scale.In solvated systems and crystals, intermolecular interactions are especially important, while intramolecular degrees of freedom often vary predictably.Source: Building flows on natural molecular… · p. 2
  • When rigid residues make the target density singular on the full space, an all-degrees-of-freedom flow cannot be reweighted against it or trained by reverse KL.A physically meaningful flow over molecular poses must also scale to systems with many interacting poses.Source: Building flows on natural molecular… · p. 2
  • The proposed flow models the joint positions and orientations of multiple molecules, focusing here on rigid bodies while allowing internal degrees of freedom to be modeled separately.Its pose transformation maps rotations to one of two representing quaternions stochastically, applies a sign-flip-equivariant flow, and marginalizes the two paths when mapping back to rotations.Source: Figure 1 · p. 1, §1 · p. 1, Building flows on natural molecular… · p. 2
  • The architecture prescribes smooth densities on the rigid-body submanifold, a property reported as critical for faithfully modeling physical force fields.It is compatible with permutation-equivariant architectures and treats external poses independently from internal degrees of freedom.Source: Building flows on natural molecular… · p. 2
  • The method introduces symmetrized Möbius transformations and projective convex gradient maps as smooth flows for rotations.The approach is evaluated on a multimodal tetrahedral system in an external field and ice XI crystals using the TIP4P water model.Source: Building flows on natural molecular… · p. 2

2. Related Work

Related work covers manifold flows and density estimation on SO(3), normalizing flows for molecular equilibrium sampling, and established non-machine-learning methods for crystal sampling.

Related Work

  • Manifold normalizing flows include convex-potential flows, neural ODEs, smooth coupling flows on tori and spheres, and chart-based non-smooth flows; double-cover flows for single-pose computer vision were discussed concurrently by Liu et al. (2023).Other SO(3) density-estimation approaches use the exponential map or an implicit neural representation on S3.Source: §2 · pp. 2–3
  • Molecular normalizing flows have been used for importance sampling, free-energy estimation, and conformational generation, although Garcia Satorras et al. (2021) studied density estimation without explicit thermodynamic treatment.The cited work includes applications to equilibrium sampling and free-energy differences.Source: §2 · p. 3
  • Crystal-sampling methods without machine learning include molecular dynamics and Markov chain Monte Carlo, thermodynamic integration, and targeted free-energy perturbation combined with MBAR.The latter combination has been used to compute crystal free energies.Source: §2 · p. 3

3. Theory & Method

The method targets equilibrium distributions over multiple rigid bodies’ positions and orientations, while leaving internal degrees of freedom fixed. It uses a quaternion double cover and flip-equivariant sphere flows to obtain smooth, tractable transformations for rotations.

3. Theory & Method

  • Normalizing flows map a tractable base density toward a thermodynamic target, supporting sampling and estimation of free-energy differences; this method focuses on rigid-body poses with internal degrees fixed.The rigid-body flow can be combined conceptually with flows over internal degrees of freedom.Source: §3 · pp. 3–4

3.1. Flows on SO(3) via the S3 →SO(3) double cover.

  • Representing rotations on S3 uses the double cover of SO(3), so q and −q correspond to the same rotation and yield two paths through the pose transformation.The procedure lifts a rotation to either quaternion sign, applies a diffeomorphism on position and quaternion, then projects the result back to SO(3).Source: §3.1 · pp. 4–5
  • The construction combines a position transform with a flip-equivariant quaternion map, enabling efficient inversion and volume evaluation when the component maps are tractable.The framework is designed to support rotational equivariance and an exact density; direct rotation representations pose smoothness or expressivity challenges.Source: §3 · p. 4, §3.1 · p. 5

3.2. Flip-equivariant diffeomorphisms on S3

  • Two smooth flip-equivariant sphere-flow families are introduced: symmetrized Möbius transforms and projective convex gradient maps.Möbius maps have analytic inverses and volume changes but less expressivity for highly multimodal densities; projective convex gradient maps are more expressive but require more numerical effort, while their S3 volume change can be computed efficiently in closed form.Source: §3.2 · pp. 5–6

4. Experiments

The experiments assess rotational multimodality in a field and free-energy transfer for ice XI. Projective convex gradient maps reproduce the tetrahedral modes, while the learned flow maps ice densities with close energy and structural overlap.

4.1. Tetrahedron in external field

  • The tetrahedron test fixes methane’s internal degrees of freedom and trains three smooth rotation-flow models on 50,000 equilibrium samples.The comparison includes projective convex gradient maps, symmetrized Moebius transformations, and affine quaternion flows.Source: §4.1 · pp. 6–7
  • Projective convex gradient maps faithfully reconstruct the tetrahedron’s multimodal rotational density, while symmetrized Moebius transforms struggle and affine quaternion layers represent only one mode.Figure 2 compares the MD density with densities from the three flow architectures.Source: Figure 2 · p. 6, §4.1 · p. 7

4.2. Ice XI in the TIP4P water model

  • Ice XI tests use rigid TIP4P-Ew water from a 250 K base, with 16 molecules at 100 and 50 K and 128 molecules at 100 K; quantum effects are excluded.The authors caution that this model should not be expected to reproduce experimental measurements.Source: §4.2 · p. 8
  • A single base-state MD run supplies training and separate LFEP evaluation data for estimating the free-energy difference.The run is split to avoid systematic errors.Source: §4.2 · p. 8
  • For ice XI, the mapped density closely overlaps the target energy distribution, and reweighting yields nearly perfect overlap in energy and oxygen-oxygen g(r).The target distribution is shown as a reference and is not used for training.Source: Figure 3 · p. 7, §4.2 · p. 8
  • Table 1 reports per-molecule ∆F estimates from MBAR and LFEP, with bootstrap uncertainty at two standard deviations; LBAR would require unavailable target samples.The reported uncertainty is estimated by bootstrapping.Source: §4.2 · p. 8, Table 1 · p. 8

5. Discussion

The discussion highlights the rigid-body flow’s molecular-crystal results and identifies broader extensions and scaling questions. It also situates free-energy estimation in a setting where entropic effects matter and conventional sampling can be costly.

5. Discussion

  • The approach models positions and orientations of interacting molecules, using smooth quaternion flows and flip-equivariant layers, and was evaluated on a multimodal benchmark and ice systems.The discussion describes the method as effective for molecular-crystal density modeling.Source: §5 · p. 8
  • Rigid-body flows may also be relevant to robotics, beyond their application to molecular crystals.The discussion cites Brehmer et al. (2023) as related work.Source: §5 · p. 8
  • Phase transitions and sampling beyond crystals, including liquid and gas phases, remain major open challenges.The discussion notes these are open problems even for non-molecular systems.Source: §5 · pp. 8–9
  • Joint SE(3) equivariance and built-in handling of internal symmetries were not explored; the latter may be needed to scale to larger systems.For N water molecules, the mirror-symmetry group has order 2N.Source: §5 · p. 9
  • The work considers rigid-body molecules, while incorporating internal molecular degrees of freedom is a straightforward extension described as important for larger molecules or more accurate force fields.This extension is not part of the current rigid-body treatment.Source: §5 · p. 9
  • The scaling limits of flows for molecular systems remain untested, leaving their reliable and efficient use on much larger systems an open question.The discussion relates this question to Abbott et al. (2022), which studied physical potentials in lattice physics.Source: §5 · p. 9

A. The sampling problem in molecular crystals

  • At non-zero temperature, entropic effects make free-energy differences—not just energy differences—relevant to assessing crystal stability.The section distinguishes this task from finding all energetically stable polymorphs.Source: A. The sampling problem in molecular… · p. 13
  • Thermodynamic integration can require costly sampling at multiple intermediate states, while learned normalizing-flow maps motivate the targeted free-energy perturbation method used here.The section notes that defining such maps is nontrivial and that integration paths in molecular crystals can be complex.Source: A. The sampling problem in molecular… · p. 13

B. Proofs and derivations

The derivation frames volume change for a smooth injective map between submanifolds through the tangent spaces at an input and its image.

B. Proofs and derivations

  • The setup takes M and N as submanifolds of Euclidean space and F as a smooth injective map between them.The map is assumed extendable to open neighborhoods of M and N.Source: §B.1 · p. 13
  • Volume change is evaluated using the tangent spaces at x and F(x), with bases chosen for each space.The bases are represented as matrices in the ambient space.Source: §B.1 · p. 13
  • The derivation uses the induced change of volume for the map between the submanifolds.The passage introduces the tangent-space and basis setup for computing it.Source: §B.1 · p. 13

B.2. Density of the mixture

The mixture-density derivation accounts for the alternative paths produced by the double cover and includes each path’s induced volume change.

B.2. Density of the mixture

  • The mixture density is obtained by summing the probabilities of each path and accounting for each transformation’s induced volume change.The derivation then simplifies the resulting expression using the two alternatives associated with the double cover.Source: §B.2 · p. 14
  • The stated mixture-density expression includes Jacobian factors for the rigid-body transform and inverse flow, alongside the base density.The passage presents terms for both alternatives in the mixture.Source: §B.2 · p. 14
  • The derivation states that the frame transform contributes no volume change because changes in position or rotation after mapping to the standard frame are SE(3) actions.The stated simplification concerns the transforms into and out of that frame.Source: §B.2 · p. 14
  • The density expressions evaluate the inverse flow and base density at each of the two alternatives, A+ and A−.The passages give separate expressions for the two alternatives.Source: §B.2 · p. 14

B.3. Derivations for symmetrized Moebius transform

The appendix derives properties of the symmetrized Möbius transformation on the three-sphere, including its sign symmetry, inverse, and volume change.

B.3. Derivations for symmetrized Moebius transform

  • A generalized Möbius transform maps the sphere to itself and is parameterized by q in the unit ball.The derivation focuses on the d = 3 case for volume change.Source: §B.3 · pp. 14–15
  • The underlying Möbius map is an involutive diffeomorphism, so applying it twice with the same parameter returns the original point.The passage states ΦM(ΦM(p; q); q) = p.Source: §B.3 · p. 15
  • The symmetrized map averages the transformations using q and −q, then normalizes their sum, making it invariant to the sign of q.The resulting map also has sign-equivariance in its input point.Source: §B.3 · p. 15
  • The inverse derivation uses an orthogonal coordinate change to reduce the calculation to a convenient coordinate arrangement.The passages describe transforming q and p′ before applying the inverse.Source: §B.3 · p. 15
  • For the volume calculation, the tangent space of S3 is embedded in R4, and the Jacobian is formed from directional changes along tangent vectors.The derivation then simplifies the change of volume and states that the general case uses r = |q|.Source: §B.3 · p. 15

B.4. Derivations for projective convex gradient maps

This appendix establishes that projective convex gradient maps are diffeomorphisms under strict convexity and describes their parameterization, volume calculation, and numerical inversion. It also identifies affine quaternion flows as a special case.

Invertibility

  • The proof assumes a strictly convex potential with its minimum at 0; positive definiteness of its Hessian rules out a singular tangent map.The argument also uses the transformed potential ψ, which remains strictly convex with minimum at 0.Source: §B.4 · pp. 15–16
  • Strict convexity makes the projective convex gradient map a diffeomorphism of the sphere.The proof shows its tangent-space Jacobian is nonsingular, using positive definiteness of the Hessian.Source: §B.4 · pp. 15–16

Parameterizing the potential ϕ

  • The chosen potential is a softsign neural expression plus a quadratic term, with H controlling its complexity and multimodal modeling capacity.The appendix notes that this simple implementation is fast to evaluate and worked well in practice.Source: §B.4 · p. 16

Computing the volume element

  • For d = 3, the volume element uses a 3×3 tangent-space Jacobian determinant computed analytically and numerically stably; for d > 3, numeric computation can be expensive and unstable.The tangent bases are constructed with a three-step Gram-Schmidt procedure.Source: §B.4 · p. 16

Numerical inverse

  • Numerical inversion minimizes a residual with LBFGS; for d = 3 and the potentials used here, it reaches absolute error below 0.00001 in 10–15 iterations.For a quadratic potential ϕ(p) = pᵀWᵀWp, the map is exactly the affine quaternion map of Liu et al. (2023).Source: §B.4 · p. 17

B.5. Bundle flows

The appendix constructs maps on SO(3) from sign-equivariant maps on its quaternion double cover S3. A stochastic lift to S3 followed by mapping and projection gives the same pushforward regardless of the chosen section.

Bundle flows

  • A sign-equivariant map on S3 induces a well-defined map on SO(3), independent of the local trivialization of the double cover.The construction uses the two-point fibers {−1, 1} of the projection.Source: §B.5 · p. 17
  • The induced SO(3) pushforward equals stochastic lifting to S3, applying the sign-equivariant map, and projecting back, independently of the stochastic section chosen.A deterministic choice of either point in each fiber is one possible section.Source: §B.5 · p. 17
  • The construction extends to other bundles when the total-space map is a bundle morphism and a stochastic section to the projection can be constructed.Figure 4 identifies the coupling used for the tetrahedron experiment.Source: §B.5 · p. 17, Figure 4 · p. 18

C. Details on tetrahedron experiments

The tetrahedron experiments use an augmented flow with rotation and auxiliary-noise coupling, flip-invariant quaternion conditioning, and several candidate rotation transforms. The appendix also specifies the methane simulation and training setup.

Setup of the simulation and data generation

  • Methane simulations fix the carbon position and internal geometry, then collect 50,000 samples from a 100 K Langevin simulation.The CH bonds are constrained and the HCH angles restrained to 109.47122°; every 500th frame is retained.Source: C. Details on tetrahedron experiments · p. 18

Flow model

  • The augmented flow couples rotation and auxiliary noise in two layers, using a two-dimensional unit-normal noise density and an antipodal mixture of von Mises–Fisher base densities with concentration 2.5.Auxiliary noise is added because the dataset otherwise contains only one quaternion, which cannot be modeled with coupling layers.Source: C. Details on tetrahedron experiments · p. 18
  • Quaternion conditioning uses a flip-invariant embedding before dense conditioning networks, with linear layers F and S shared across inputs.The figure passage identifies the embedding as the input to the conditioning network.Source: C. Details on tetrahedron experiments · pp. 18–19, Figure 5 · p. 19

Rotation transformations

  • The rotation-transform comparison includes affine flows, symmetrized Moebius layers, and projective convex gradient maps with H = 8, 32, or 128; the main-text result uses H = 128.The appendix reports that quality degrades when H is varied and refers to Figure 6 for the ablation.Source: C. Details on tetrahedron experiments · p. 19

Training

  • Each candidate flow is trained for 50,000 steps with ADAM, batch size 32, and learning rate 0.0005.The objective is the transformed likelihood objective described after augmenting the data and base densities.Source: C. Details on tetrahedron experiments · p. 19

D. Details on the Ice XI experiments

The Ice XI experiments evaluate energy-trained flows across three system and temperature setups, using alternating position and rotation couplings. The appendix reports Kish effective sampling sizes and describes the evaluation and computational-cost context.

Details of the simulations and data generation process

  • The simulations use rigid TIP4P-Ew water and treat Ice XI as sufficiently stable for long runs, although it is likely metastable under the considered conditions.The authors note that similar results can be obtained with other rigid water models such as TIP4P-ice.Source: D. Details on the Ice XI experiments · p. 20

Flow model

  • Across the three experiments, the flow alternates position and rotation couplings in round-robin order over four iterations.Positions conditioned on rotations use Moebius layers; rotations conditioned on positions use symmetrized Moebius transforms.Source: D. Details on the Ice XI experiments · p. 20, Figure 7 · p. 23
  • Because global translation invariance leaves 3(N − 1) positional degrees of freedom, one molecule’s position is fixed and the flow acts on the others.The fixed molecule’s rotation remains variable.Source: D. Details on the Ice XI experiments · p. 20
  • Transformer conditioners share parameters and reuse local structure, encode molecule indices as one-hot vectors, and use two blocks with 8 heads and 32 channels per flow layer.A modified attention mechanism squares rotation logits to cancel their sign and preserve flip-invariance for position conditioning.Source: D. Details on the Ice XI experiments · p. 20, Figure 8 · p. 23, Figure 9 · p. 24

Details on the evaluation

  • The evaluation compares potential-energy profiles and oxygen–oxygen radial distribution functions; target configurations were not seen during training because the flows were trained purely on energies.Figure 11 shows 2D projections for N = 16, with base samples, mapped target configurations, and target MD samples in separate rows.Source: D. Details on the Ice XI experiments · pp. 20–21, Figure 11 · p. 26
  • 22.74 ± 1.58% for setup (a), 6.90 ± 1.36% for (b), and 0.39 ± 0.06% for (c) are the Kish effective sampling sizes of generated samples.These are averages and standard deviations over 10 estimates, each based on 10,000 random evaluation-dataset points.Source: D. Details on the Ice XI experiments · pp. 20–21

Computational cost

  • Computational efficiency was not optimized; the authors expect speedups over MD when energy evaluation is more expensive, and report training times of about 10 minutes for small systems and 30 minutes for the larger one.Training times were measured on a GeForce GTX 1080 Ti.Source: D. Details on the Ice XI experiments · p. 21

Related papers

Loading 2301.11355v4…