Source-linked AI summary
uFlowCSP: Crystal Structure Prediction using Mean flow generative models
Sourin Dey, Dipannoy Das Gupta, Lai Wei, Sadman Sadeed Omee, Jianjun Hu
TL;DR
Crystal structure prediction needs efficient generation from composition because existing diffusion and flow-matching samplers use many sequential evaluations. uFlowCSP learns interval-averaged velocity with a chemistry- and symmetry-aware Transformer, achieving comparable or better accuracy in one to five evaluations while inference remains formula-only. Its gains are strongest in accuracy per evaluation, although the reported training data and relaxation analyses have defined scope boundaries.
Problem
Diffusion and flow-matching CSP inference requires tens to thousands of sequential network evaluations per candidate, limiting sampling efficiency.
Method
uFlowCSP learns average probability-flow velocity with a chemistry- and symmetry-aware Transformer and uses one-step or few-step MeanFlow updates conditioned at inference only on composition.
Results
Five-step uFlowCSP reaches 83.64% on MP-20, exceeding CrystalFlow at 78.34% and DiffCSP at 77.93% with 20x fewer evaluations than CrystalFlow and about 20,000 evaluations for DiffCSP.
Takeaways & Limitations
uFlowCSP improves accuracy per network evaluation rather than only peak accuracy, with one-step performance matching CrystalFlow while using 100x fewer evaluations.
Takeaways & Limitations
Training on larger, carefully cleaned data is left to future work because added sources include mostly DFT-hypothetical and two-dimensional entries.
Abstract
from arXiv · showhide
Crystal structure prediction (CSP) is fundamental to computational materials discovery. Generative models including CDVAE, DiffCSP, FlowMM, and CrystalFlow learn stable-crystal distributions directly, but diffusion and flow-matching inference requires tens to thousands of sequential network evaluations per candidate. We introduce uFlowCSP, a MeanFlow-based CSP model that learns the average, rather than instantaneous, probability-flow velocity. It generates a complete structure in one to five evaluations, delivering 5x-58x faster inference with equal or better performance. A chemistry- and symmetry-aware Transformer uses canonical atom ordering, global composition, and per-token chemistry embeddings. A coarse crystal-system token is used only during training; it provides additive gains, particularly improving space-group agreement despite being absent at inference, which remains formula-only. On MP-20 with 20 candidates per target, one step matches CrystalFlow (78.38% vs. 78.34%) with 100x fewer evaluations and about 10x lower wall-clock time. Five steps reach 83.64%, exceeding CrystalFlow (78.34% at 2,000 evaluations) and DiffCSP (77.93% at about 20,000), while using 20x fewer evaluations. uFlowCSP generates 10,000 structures in 0.39-1.31 minutes, versus 6.5 for CrystalFlow and 76.1 for DiffCSP. Under CSPBench's energy-ranked top-five structure-and-space-group criterion, five-step uFlowCSP reaches 72%/72%/65% structure, space-group, and consensus match rates. CrystalFlow reaches 78%/73%/68% at 100 steps but falls to 49%/32%/31% at five. Thus, uFlowCSP improves accuracy per network evaluation, not merely peak accuracy.
1 INTRODUCTION
CSP determines periodic atomic arrangements from composition, enabling downstream materials-property prediction and discovery. Generative models learn stable-crystal distributions directly, but their sequential sampling still requires many network evaluations; uFlowCSP targets this bottleneck with one-step or few-step MeanFlow generation.
- Motivation: CSP determines how atoms arrange in periodic solids from chemical composition alone.The resulting structure links composition to predicted stability, electronic, transport, catalytic, and mechanical properties.
- Prior approaches: Generative CSP replaces per-compound energy-landscape searches with models that learn stable-structure distributions from materials databases.CDVAE, DiffCSP, symmetry-aware models, and flow-based methods exemplify this shift.
- Limitation: Many diffusion and flow-based methods still require numerical integration over sequential network evaluations during generation.This inference cost remains even when flow-based models use fewer steps than diffusion models.
- uFlowCSP: One-step or few-step generation is the motivation for reducing per-candidate inference cost in CSP.uFlowCSP adapts MeanFlow to produce a complete periodic structure in one network evaluation and uses a chemistry- and symmetry-aware Transformer.
- Contribution: uFlowCSP reports higher match rates than multi-step diffusion and flow-matching baselines at 20 to 200 times fewer network evaluations and 5 to 58 times lower measured wall-clock time.Its auxiliary symmetry token is used only during training yet improves space-group agreement while inference remains formula-only.
2.1 PROBLEM SETTING AND REPRESENTATION
uFlowCSP represents a crystal using the composition-fixed atom identities together with continuous coordinates and lattice parameters. The model therefore predicts geometry and cell information while conditioning on, rather than generating, atomic numbers.
- Problem setting: Both MP20 and CSPBench provide the primitive-cell atom multiset as the composition condition without supplying a space-group label.The number of atoms is therefore known at inference in both protocols.
- Representation: A crystal with N atoms is represented by atomic numbers, Cartesian coordinates X ∈ R^N×3, and lattice matrix L ∈ R^3×3.The lattice matrix is flattened into nine continuous parameters.
- Prediction target: Atomic numbers are fixed by composition, while the model predicts continuous coordinates and lattice parameters.The representation separates discrete chemical identity from the continuous state transported by the flow.
2.2 FEW-STEP GENERATIVE BACKBONE
uFlowCSP uses MeanFlow to learn interval-averaged probability-flow velocity and generate crystals through one or a few direct updates. Its atom-token state carries coordinates and replicated lattice information, while sampling runs from noise to data in a fixed number of steps.
- MeanFlow backbone: MeanFlow generates crystal structures in a small number of network evaluations by learning average rather than instantaneous velocity.This avoids the many-step numerical integration required by ordinary flow matching.
- Crystal state: Each of N atom tokens concatenates the atom position with a flattened copy of the unit cell, producing a state x ∈ R^N×12.Atomic numbers remain composition conditioning, and the flow transports only coordinates and lattice channels.
- Velocity prediction: The network outputs a same-shaped per-atom displacement whose first three channels update coordinates and remaining nine channels update each atom’s lattice copy.The prediction is informed jointly by all atoms, the shared cell, composition, and the two MeanFlow times.
- Sampling: Sampling starts from Gaussian noise at t = 1 and integrates toward data at t = 0 using S uniform mean-velocity jumps.The final lattice is recovered by averaging the per-atom lattice channels.
- Sampling: Setting S = 1 enables single-step generation, while S = 5 is the main configuration used to improve quality at a still-small fixed cost.The same update rule supports one-step and few-step generation.
2.3 CHEMISTRY- AND SYMMETRY-AWARE TRANSFORMER
The chemistry- and symmetry-aware Transformer combines canonical atom ordering and local chemical descriptors with global formula and coarse-symmetry conditioning. The crystal-system signal is training-only, so inference remains formula-only while benefiting from auxiliary symmetry learning.
- Architecture: A sequence Transformer treats each atom as a token and uses canonical ordering plus composition-derived conditioning to make attention chemically and structurally informative.The conditioning includes per-atom chemistry, whole-formula chemistry, and a coarse symmetry hint.
- Canonical atom ordering: Atoms are sorted by Pauling electronegativity and, for symmetry-aware training examples, grouped by space-group orbits and ordered by Wyckoff letter.This removes arbitrary permutations and supplies a consistent sequence structure.
- Per-atom chemistry: Each token receives four fixed elemental descriptors—electronegativity, atomic radius, period, and group—in addition to its learned atom-type embedding.These per-atom features provide chemical information locally.
- Whole-formula chemistry: A whole-formula embedding summarizes element fractions and conveys stoichiometric context to the model.It complements token-level features by describing the composition as a whole.
- Coarse symmetry: The coarse symmetry signal maps space groups to seven crystal-system classes plus a null class rather than predicting an exact 230-way space group.The signal supplies a broad symmetry prior while remaining compatible with formula-only inference.
- Inference: The crystal-system token is used only during training and is always null at inference, where atom order, per-atom chemistry, and formula embedding come from composition alone.Its benefit is described as regularization from an auxiliary symmetry-aware task rather than supplied symmetry information.
2.4 IMPLEMENTATION AND TRAINING
The study trains and evaluates uFlowCSP on MP-20 under controlled benchmark settings, while identifying data composition as an important scope boundary for generative CSP.
- Training setup: MP-20 contains structures with at most 20 atoms per unit cell, and the 180 CSPBench evaluation structures are excluded from training.The denoiser is a 12-layer Transformer with hidden size 768 and 12 attention heads.
- Training setup: The study uses two evaluation protocols and trains a separate model for each, with both models trained on MP-20 and evaluation structures withheld.The protocols differ in input format and scoring procedure.
- Controlled comparison: All generative baselines are trained on the same MP-20 split, isolating comparisons around the one-step objective rather than training-data size.This design also preserves comparability with published results.
- Evaluation data: The CSPBench targets form a curated, difficulty-stratified set of 60 binary, 60 ternary, and 60 quaternary compounds.They are graded by symmetry, prototype, element count, and cell size.
- Scope boundary: Generative models may not benefit from additional data when sources are poorly matched to bulk crystals, so larger carefully cleaned training data remains future work.GNoME is mostly DFT-hypothetical and C2DB is two-dimensional.
2.5 EVALUATION PROTOCOLS
The paper evaluates uFlowCSP under separate MP-20 and CSPBench protocols that differ in inputs, candidate processing, and match criteria. Their absolute rates are therefore reported separately rather than compared directly.
- Protocol design: The two protocols test different settings: Protocol A scores relaxed, energy-ranked candidates, whereas Protocol B scores raw generated candidates on structure match alone.Their absolute rates are not directly comparable.
- MP-20 evaluation: Protocol A uses full primitive-cell composition and generates k raw candidates for each of 9,046 MP-20 test structures without relaxation or energy-based selection.The protocol follows prior generative CSP evaluations.
- MP-20 evaluation: A Protocol A match requires at least one candidate within ltol = 0.3, stol = 0.5, and angle_tol = 10° under StructureMatcher.RMSE is averaged over matched structures using the minimum candidate displacement.
- CSPBench evaluation: Protocol B receives only the primitive-cell chemical formula, generates 50 candidates, relaxes them with ORB-v3, and ranks them by CHGNet energy per atom.The five lowest-energy candidates form the scored shortlist, while the reference is used only for scoring.
- CSPBench evaluation: CSPBench counts structure and space-group matches separately, while consensus requires one top-five candidate to satisfy both criteria.Formulas without a valid candidate count as failures.
3 RESULTS & DISCUSSION
uFlowCSP’s conditioning and average-velocity Transformer improve crystal-structure prediction under low evaluation budgets. Across MP-20 and CSPBench, it provides strong accuracy and competitive structure quality with substantially fewer evaluations than flow-matching baselines.
- Component ablation: +14 space-group and +12 consensus are contributed by canonical atom ordering, the largest ablation gain.The ordering is described as a preprocessing choice that removes an order-related nuisance for the Transformer.
- Component ablation: 72/72/65 is reached after adding canonical ordering, formula and per-token chemistry features, and coarse crystal-system conditioning under the same five-step budget.The ablation attributes gains to conditioning rather than additional integration steps.
- Velocity-field parameterization: 72/72/65 versus 70/… is reported for the chemistry- and symmetry-aware Transformer versus the equivariant-GNN backbone at matched MeanFlow objective and five-step cost.Both predictors use the same training objective and sampling budget, isolating the backbone comparison.
- Comparison with CrystalFlow: 49/32/31 at five steps for CrystalFlow versus uFlowCSP’s 72/72/65 shows a large matched-cost advantage, while CrystalFlow reaches 78/73/68 at 100 steps using 20× more evaluations.The comparison concerns structure, space-group, and consensus rates on CSPBench.
- Comparison across method families: Five-step uFlowCSP is the strongest generative method at equal cost on CSPBench and remains close behind full-budget CrystalFlow, while TCSP 2.0 is strongest overall with much broader data access.The template and generative comparisons are not matched on data access.
- Qualitative examples: All three Figure 4 examples—U3Sb4Ir3, Ba2YIrO6, and CeAl2BRu2—match ground-truth structure and space group.The examples are shown in an aligned coordinate frame for visual comparison, while match statistics use unaligned structures.
- MP-20 benchmark: 83.64% at k = 20 exceeds CrystalFlow’s 82.49 at k = 100 on MP-20 with 100× fewer network evaluations.Table 3 defines NFE as S × k and reports wall-time for generating 10,000 structures.
- Relaxation behavior: Five-step uFlowCSP relaxes about as easily as much longer CrystalFlow integration on the ORB-v3 SiO2 proxy, with all generators reaching the same low-energy peak.This proxy is defined by a machine-learned potential and optimizer rather than DFT.
4 CONCLUSION
uFlowCSP uses one-step MeanFlow to generate complete crystal structures in a handful of evaluations, addressing the sequential integration bottleneck. Its simple Euclidean geometric representation remains a stated implementation limitation.
- uFlowCSP generates complete periodic structures in a handful of network evaluations by learning average rather than instantaneous probability-flow velocity.
- 83.64% match rate at k = 20 uses 100 network evaluations per target, surpassing CrystalFlow’s 78.34% at 2,000 and DiffCSP’s 77.93% at ∼20,000.
- Five-step uFlowCSP reaches 72/72/65% structure, space-group, and consensus match rates on the 180-formula CSPBench suite.
- The best configuration transports Cartesian positions and a raw 3 × 3 lattice along Euclidean paths, ignoring fractional-coordinate periodicity and unreduced-cell over-completeness.
5 CONTRIBUTION
The contribution statement assigns responsibility across conceptualization, methodology, software, resources, writing, visualization, supervision, and funding acquisition.
- Conceptualization is assigned to J.H.
- Methodology and original-draft writing are assigned to J.H., S.D., L.W., S.O., and D.D., with the listed authors contributing across those roles.
- J.H. is assigned resources, supervision, and funding acquisition, while S.D. is assigned software and visualization.