Source-linked AI summary

FlowMM: Generating Materials with Riemannian Flow Matching

Benjamin Kurt Miller, Ricky T. Q. Chen, Anuroop Sriram, Brandon M Wood

arXiv:2406.04713v1cs.LGcond-mat.mtrl-scics.AIphysics.comp-phstat.ML

TL;DR

Crystal generation must model a difficult mixture of discrete, continuous, periodic, and symmetry-constrained variables while finding rare stable structures for known and novel compositions. FlowMM generalizes Riemannian Flow Matching to crystals, using symmetry-aware paths and flexible base distributions for CSP and DNG. It achieves strong benchmark performance and comparable stability with fewer integration steps, including about 3x more efficient generation than previous open methods.

  • Problem

    Crystal generation must find stable structures for known compositions and novel compositions within a combinatorial space where only a vanishingly small fraction of arrangements are experimentally realizable.

    Source: §1 · pp. 1–2
  • Method

    FlowMM generalizes Riemannian Flow Matching to jointly model lattices, periodic fractional coordinates, and atom types while respecting crystal symmetries.

    Source: §1 · p. 2
  • Results

    FlowMM achieves competitive or state-of-the-art benchmark performance and comparable generated-material stability while requiring about 3x fewer integration-step resources than previous open methods.

    Source: §1 · p. 2, §5 · p. 9
  • Takeaways & Limitations

    FlowMM provides a flexible crystal-generation framework that supports both CSP and DNG with efficient inference and quantum-chemistry validation of generated structures.

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

    Generated stability has a numerical limitation because some one-element structures can be incorrectly marked stable under the Ehull < 0 criterion.

    Source: §2.3 · p. 5

Abstract

from arXiv · show

Crystalline materials are a fundamental component in next-generation technologies, yet modeling their distribution presents unique computational challenges. Of the plausible arrangements of atoms in a periodic lattice only a vanishingly small percentage are thermodynamically stable, which is a key indicator of the materials that can be experimentally realized. Two fundamental tasks in this area are to (a) predict the stable crystal structure of a known composition of elements and (b) propose novel compositions along with their stable structures. We present FlowMM, a pair of generative models that achieve state-of-the-art performance on both tasks while being more efficient and more flexible than competing methods. We generalize Riemannian Flow Matching to suit the symmetries inherent to crystals: translation, rotation, permutation, and periodic boundary conditions. Our framework enables the freedom to choose the flow base distributions, drastically simplifying the problem of learning crystal structures compared with diffusion models. In addition to standard benchmarks, we validate FlowMM's generated structures with quantum chemistry calculations, demonstrating that it is about 3x more efficient, in terms of integration steps, at finding stable materials compared to previous open methods.

1. Introduction

FlowMM addresses crystal generation for stable structure prediction and de novo composition-structure generation, where only a small fraction of the combinatorial search space is experimentally realizable. It uses symmetry-aware Riemannian Flow Matching to model crystals jointly and efficiently.

1. Introduction

  • Crystal discovery must navigate an intractable space of atomic arrangements, of which only a vanishingly small fraction are experimentally realizable.Source: §1 · p. 1
  • CSP predicts structures for a supplied composition, whereas DNG generates both compositions and crystal structures, with stability assessed by energetic distance to the convex hull.Structures with Ehull ≤0.0 eV/atom are considered stable in the stated heuristic.Source: §1 · p. 1
  • Crystal modeling jointly handles discrete atom types, continuous unit-cell and coordinate variables, periodic coordinates, bounded angles, positive lengths, and variable atom counts.Source: §1 · p. 2
  • FlowMM jointly estimates symmetric crystal distributions for CSP and DNG with Riemannian Flow Matching and a finite-time Continuous Normalizing Flow.Source: §1 · p. 2
  • FlowMM uses a symmetry-aware objective, plausible lattice base distributions, and binary atom-type representations, while diffusion baselines require multiple variable-specific frameworks and complex lattice representations.The binary representation uses 7 dimensions per atom versus 100 dimensions for a 100-class simplex.Source: §1 · p. 2
  • Quantum-chemistry validation finds comparable stability to competing methods while FlowMM is significantly faster at inference.Source: §1 · p. 2

2. Preliminaries

The preliminaries represent crystals as periodic, symmetry-equivalent labeled point clouds with lattice parameters, then formulate Flow Matching on suitable manifolds. They define stability, datasets, and proxy metrics for CSP and DNG.

2.1. Representing crystals and their symmetries

  • A crystal is represented by a unit cell and a labeled periodic point cloud, with fractional coordinates capturing positions under lattice tiling.Alternative unit cells can represent the same lattice, so Niggli reduction selects a unique minimum-volume representation.Source: §2 · p. 3, §2.1 · p. 3
  • The representation accounts for permutation, translation, and rotation symmetries, using wrapped periodic coordinates and rotation-invariant lattice parameters.Source: §2.1 · p. 3
  • DNG represents atom types with ⌈log2 h⌉ binary dimensions rather than h-dimensional one-hot vectors, while CSP treats atom types as conditioning information.Source: §2.1 · p. 3

2.2. Learning distributions with Flow Matching

  • Flow Matching learns a time-dependent vector field that transports a chosen base distribution to the target distribution along a probability path, using tractable conditional vector fields.Source: §2.2 · p. 4

2.3. Crystalline Solids

  • Crystal stability is evaluated relative to a composition-dependent convex hull, with epistemic uncertainty because unknown materials are absent from the Materials Project reference hull.Source: §2.3 · pp. 4–5

2.4. Problem statements & Datasets

  • CSP models metastable structures conditional on composition, while DNG models the joint distribution of compositions and structures to generate known and novel metastable materials.Source: §2.4 · p. 5
  • The study uses MP-20 and MPTS-52 as realistic datasets, plus Perov-5 and Carbon-24 as simplified unit tests with proxy metrics but no applicable stability analysis.Source: §2.4 · p. 5
  • CSP is evaluated with match rate and coordinate RMSE, while DNG uses validity, chemical-fingerprint coverage, property distance, stability rate, S.U.N. rate, and integration-step costs.Source: §2.4 · p. 5

3. Riemannian Flow Matching for Materials

FlowMM builds a crystal-manifold generative model whose geometry, base distribution, conditional paths, and neural vector field preserve crystal symmetries. It adapts the same framework to CSP and DNG through task-specific atom-type treatment and base distributions.

3. Riemannian Flow Matching for Materials

  • FlowMM defines a generative model on the product manifold C of atom types, periodic fractional coordinates, and lattice parameters.Source: §3 · pp. 5–6
  • The construction enforces invariance of the marginal path to permutation, translation, and rotation.Source: §3 · pp. 6–7
  • Angle boundaries are handled by mapping [60, 120] to unconstrained space with a diffeomorphism before learning and integration.Source: §3 · p. 6
  • Its base distribution is factorized across components, using uniform fractional coordinates, learned positive length distributions, bounded uniform angles, and task-specific atom-type distributions.Source: §3 · p. 6
  • The periodic-coordinate construction removes average torus translation in tangent space so the marginal probability path becomes translation invariant.Source: §3 · pp. 6–7
  • A graph neural network enforces permutation equivariance and translation invariance through message passing and displacement-based edge features, while the lattice representation enforces rotation invariance.Source: §3 · p. 7
  • At inference, anti-annealing scales the predicted velocity over time; increasing fractional-coordinate velocity helped CSP, whereas increasing lattice velocity hurt CSP.Source: §3 · p. 7

4. Experiments

FlowMM is evaluated on crystal structure prediction and de novo generation, with emphasis on proxy metrics, inference efficiency, and DFT-based stability. Across these tasks, it improves sampling efficiency and matches or outperforms diffusion baselines on several distributional and stability measures.

4.1. Crystal Structure Prediction

  • FlowMM reaches maximum match rate in about 50 integration steps, at least an order of magnitude fewer than DiffCSP’s 1000 steps on CSP.The models use the same underlying neural network for an apples-to-apples comparison.Source: §4.1 · p. 8
  • FlowMM outperforms competing models on Match Rate and RMSE on the challenging realistic MP-20 and MPTS-52 CSP datasets.CSP evaluation covers Perov-5, Carbon-24, MP-20, and MPTS-52 using StructureMatcher-based proxy metrics.Source: §4.1 · p. 8

4.2. De novo generation

  • FlowMM matches the MP-20 distribution of unique elements more closely than diffusion models, which generate too many materials with N-ary ≥5.The N-ary comparison is supported by both Figure 3 and the Wasserstein distance metric N_el in Table 2.Source: Figure 3 · p. 8, §4.2 · p. 9
  • FlowMM significantly outperforms CDVAE and DiffCSP on Cost and S.U.N. Cost while remaining competitive on Stability Rate and S.U.N. Rate for DNG.The comparison focuses on inference-stage integration steps, although training and relaxation costs are also acknowledged.Source: §4.2 · p. 8
  • FlowMM generates lower-energy relaxed structures than CDVAE and is competitive with DiffCSP on the DNG Ehull distribution.Ehull is computed after relaxation using CHGNet and density functional theory.Source: §4.2 · p. 8, Figure 4 · p. 9

5. Conclusion

The conclusion presents FlowMM as a Riemannian Flow Matching method for periodic crystal generation that performs strongly on CSP and DNG. Its main practical advantage is reduced inference cost, while the number of integration steps needed for comparable stable-structure generation remains open.

5. Conclusion

  • FlowMM outperforms DiffCSP on MP-20 Match Rate with as few as 50 integration steps, an improvement of at least an order of magnitude.The CSP comparison uses the same network architecture for both methods.Source: §5 · p. 9
  • FlowMM significantly outperforms CDVAE and DiffCSP on Cost and S.U.N. Cost while remaining competitive on Stability Rate and S.U.N. Rate.The authors attribute this to 3x more efficient inference in integration steps.Source: §5 · p. 9
  • The authors did not investigate whether FlowMM can generate a similar number of stable structures using only a handful of integration steps.They identify this as future work based on the efficient CSP results.Source: §5 · p. 9

Impact Statement

The paper describes FlowMM as a model for predicting compositions and structures of stable materials, with possible relevance to materials discovery. The authors state that downstream societal effects are difficult to judge and likely limited to new research directions.

Impact Statement

  • FlowMM may aid discovery of novel materials relevant to catalysis, higher-energy-density batteries, and other areas of materials science and chemistry.The statement concerns potential downstream applications rather than demonstrated impacts.Source: Impact Statement · p. 9
  • The authors state that downstream effects are difficult to judge and likely limited to new research directions because synthesizing computational predictions remains challenging.Source: Impact Statement · p. 9

A.1. Datasets

The evaluation uses two realistic Materials Project datasets and two simplified unit-test datasets with specified sample-selection and splitting protocols. The datasets vary in unit-cell size, chronology, elemental filtering, and structural constraints.

A.1. Datasets

  • All four datasets use 60% training, 20% validation, and 20% test splits, following the splits used by Xie et al. (2021) and Jiao et al. (2023).Source: §A.1 · p. 13
  • MP-20 contains 45,231 materials with at most 20 atoms per unit cell and an energy threshold of 0.08 eV/atom.Radioactive elements are removed, and the dataset is drawn from the Materials Project database around July 2021.Source: §A.1 · p. 13
  • MPTS-52 contains 40,476 materials with up to 52 atoms per unit cell and chronological train, validation, and test splits.It applies no elemental filtering, and missing publication timestamps exclude some database entries.Source: §A.1 · p. 13
  • Perovskite-5 fixes five atoms per unit cell, fractional coordinates, and lattice angles while varying lattice lengths and atomic types.Source: §A.1 · p. 13
  • Carbon-24 contains only carbon atoms but permits variable geometries, producing a challenging CSP problem with depressed match rates.Source: §A.1 · p. 13

A.2. Proxy metrics

This section defines evaluation metrics and implementation details for FlowMM’s crystal generation, including manifold geometry, symmetry handling, variable atom counts, and quantum-chemistry validation.

A.2. Proxy metrics

  • CSP evaluation uses StructureMatcher match rate and normalized RMSE, with RMSE computed only for matched generated and held-out structures.The reported tolerances are stol = 0.5, angle_tol = 10, and ltol = 0.3.Source: §A.2 · p. 13
  • De novo evaluation measures structural and compositional validity, CrystalNN/Magpie coverage, and Wasserstein distances for atomic density and unique-element counts.Validity requires pairwise distances above 0.5 Å and charge neutrality under SMACT.Source: §A.2 · p. 13

A.3. Riemannian Manifolds

  • FlowMM models crystals on a product of Riemannian manifolds whose component tangent spaces are orthogonal and weighted by hyperparameters λa, λf, and λl.The associated measure is a product measure over the component spaces.Source: §A.3 · pp. 13–14

A.4. Specifics for De Novo Generation

  • For DNG, the model first samples the number of atoms from the training-set empirical distribution before generating the conditional crystal variables.Stability is assessed against a later Matbench Discovery convex hull, while novelty compares generated structures with element-matching training examples using StructureMatcher.Source: §A.4 · p. 14
  • Stable, novel structures are grouped by pairwise StructureMatcher comparisons so equivalent structures count once in S.U.N. computations, but matcher errors can inflate estimates.The procedure treats approximate matching as an equivalence relation despite threshold-related failures and the lack of guaranteed reflexivity.Source: §A.4 · p. 14

A.5. Symmetry

  • The method enforces invariance to permutation, translation, and rotation, while flow-matching visualizations show base samples, targets, manifold paths, intermediate points, and conditional vectors.DFT validation uses VASP at 0 K and 0 atm with Materials Project settings, PBE, and Hubbard U corrections.Source: §A.5 · p. 14, Figure 5 · p. 15, §A.6 · p. 15

A.8. Limitations of quantifying a computational approach to materials discovery

The paper identifies important limits on computational materials-discovery evaluations, spanning nonphysical conditions, simplified representations, model coverage, and imperfect novelty criteria.

A.8. Limitations of quantifying a computational approach to materials discovery

  • Stability and energy calculations use nonphysical zero-temperature and zero-pressure conditions, limiting their direct correspondence to experimental settings.Source: §A.8 · p. 15
  • The representation assumes complete homogeneity and an infinite crystal without disorder, while DFT also has basis-set, functional, and computational-cost inaccuracies.Source: §A.8 · pp. 15–16
  • FlowMM is not expected to discover the most novel structures far outside the empirical training distribution because it is trained to fit that distribution.Source: §A.8 · p. 16
  • A structure can satisfy the computational stability or novelty criteria while remaining implausible to domain experts, and StructureMatcher may not reliably define equivalence.Source: §A.8 · p. 16
  • Integration steps capture only one cost dimension because the comparison omits training and relaxation time.Source: §A.8 · pp. 16–17

B.1. Crystal Structure Prediction (CSP)

The CSP ablations examine inference-time velocity anti-annealing, lattice-parameter inductive biases, and inference variability, while DNG analysis compares the stability of generated crystals across compositional complexity.

B.1. Crystal Structure Prediction (CSP)

  • FlowMM generally benefits from applying velocity anti-annealing to fractional coordinates but not lattice parameters, while often outperforming competitors without anti-annealing.The match-rate curve without Inference Anti-Annealing remains above DiffCSP at every tested integration-step count.Source: §B.1 · p. 17, Figure 8 · p. 17
  • FlowMM’s custom lattice-parameter base distribution and unconstrained transformation provide a significant performance boost across velocity anti-annealing schemes.Ablating these inductive biases can also produce representations that do not correspond to real crystals, which are counted as failed generations.Source: §B.1 · p. 17
  • Inference variability was estimated by rerunning each trained model three times with identical hyperparameters, measuring variation during inference rather than training.Some unit-test datasets produced identical match rates across three reconstructions at n = 20, yielding ±0.0.Source: §B.1 · p. 17

B.2. De Novo Generation (DNG)

  • Stable structures become much less frequent as N-ary increases, so models generating many high-N-ary materials often produce structures that do not relax to stability.This pattern is reported for the generated stable-crystal distributions in Figure 9b.Source: §B.2 · p. 18
  • The DNG analysis compares normalized N-ary distributions for generated crystals and for structures that remain stable after density functional theory relaxation.The figures include CDVAE, DiffCSP, and FlowMM generations, with FlowMM using 1000 integration steps in the stable-crystal comparison.Source: Figure 9 · p. 18

C. Neural network

The neural network represents crystal components with graph-based message passing and predicts vector-field velocities, using separate training and selection procedures for CSP and DNG.

C. Neural network

  • The network uses a graph neural network adapted to fractional coordinates, with hidden node representations and parametric functions producing learned velocity components.The velocity components are denoted by dotted variables such as f and l and are parameterized jointly by θ.Source: C. Neural network · p. 20
  • The experiments standardize network inputs and predicted tangent-vector outputs, use AdamW optimization, apply fixed 60 −20−20 training, validation, and test splits, and report hyperparameters in Tables 5–8.General and network hyperparameters appear in Tables 5 and 6, while CSP and DNG hyperparameters appear in Tables 7 and 8.Source: C. Neural network · p. 20, Table 5 · p. 20, Table 6 · p. 20
  • Training enforces an affine combination of lattice, fractional-coordinate, and atom-type loss weights, with an additional DNG loss included when applicable.The normalized weights satisfy λ_l = ˜λ_l/˜λ, λ_f = ˜λ_f/˜λ, and λ_a = ˜λ_a/˜λ.Source: C. Neural network · p. 20
  • CSP model selection uses match rate on 2,000 validation reconstructions, whereas DNG selection filters models by validity and Wasserstein metrics before choosing the best Stability Rate after CHGNet prerelaxation.The DNG procedure selects among approximately 50 models in the top 86th percentile on the proxy metrics.Source: C. Neural network · pp. 20, 22
  • DNG adds architectural and loss components for unconditional generation, including atom-type prediction represented through analog bits and a sigmoid cross-entropy term.The atom-type representation uses a binary encoding rather than a one-hot simplex, and the added loss is combined with the other objectives.Source: C. Neural network · p. 21

D. Enforcing G-invariance of marginal probability path

FlowMM enforces crystal symmetries by constructing pairwise G-invariant conditional paths, combining equivariant flows, invariant representations, and flows between equivalence classes.

D. Enforcing G-invariance of marginal probability path

  • Pairwise G-invariant conditional probability paths yield G-invariant marginal distributions, assuming the target distribution q is G-invariant.The symmetry group includes atom permutations, lattice rotations, and translations of fractional coordinates.Source: D. Enforcing G-invariance of marginal… · p. 22
  • The framework uses three routes to construct invariant paths: equivariant vector fields, G-invariant representations, and a novel equivalence-class generalization of Riemannian Flow Matching.These approaches support the crystal symmetries through different representations or flow constraints.Source: D. Enforcing G-invariance of marginal… · p. 22
  • Regular geodesic paths with equivariant vector fields enforce permutation invariance, while directly modeling lengths and angles enforces rotation invariance through an equivalence-class representation.The first approach requires a G-invariant base distribution and equivariant vector fields; the second uses a representation bijective with the full equivalence class.Source: D. Enforcing G-invariance of marginal… · p. 22
  • The generalized construction permits the endpoint conditional flow to land on any member of the target’s equivalence class rather than requiring the exact target configuration.This relaxation enables invariant conditional vector fields and makes the flow depend only on equivalence classes while still recovering p1 = q.Source: D. Enforcing G-invariance of marginal… · pp. 22–23
  • Periodic translation invariance cannot use a uniquely defined point-cloud mean on a flat torus, so FlowMM removes the translation component from the geodesic conditional vector field.The resulting mean-free conditional vector field flows fractional coordinates to an equivalent target configuration.Source: D. Enforcing G-invariance of marginal… · p. 23
Loading 2406.04713v1…