AtomSteer: Test-Time Constraint Steering for Atomistic Generative Models

Pengfei Cai; Katharina S. C. Jäger; Utkarsh Utkarsh; Nofit Segal; Bowen Deng; Akshay Subramanian; Rafael Gomez-Bombarelli

cs.LGOct 28, 2026

AtomSteer: Test-Time Constraint Steering for Atomistic Generative Models

Authors: Pengfei Cai, Katharina S. C. Jäger, Utkarsh Utkarsh, Nofit Segal, Bowen Deng, Akshay Subramanian, Rafael Gomez-Bombarelli

Organizations: Massachusetts Institute of Technology

Abstract

Diffusion and flow matching models offer scalable approaches to generating crystalline materials, molecular conformers, and protein-ligand complexes. However, satisfying crystallographic symmetry, structural geometry, and accurate stereochemistry remains challenging. Satisfying these requirements calls for corrections that preserve the geometry already captured by the pretrained model. In this work, we propose AtomSteer, a test-time inference framework for steering pretrained atomistic generative models with structural constraints. During sampling, AtomSteer corrects intermediate structure predictions with bounded local projections, graph-rigid stereochemical transformations, and projections over periodic symmetries, accepting updates that improve constraint satisfaction while limiting changes to the rest of the structure. We demonstrate AtomSteer across three pretrained backbones, FlowMM, ET-Flow, and Boltz-2, spanning de novo crystal generation and crystal structure prediction, molecular conformer generation, and protein-ligand complex structure prediction. AtomSteer improves space-group distribution coverage in de novo crystal generation and improves target space-group recovery in crystal structure prediction. AtomSteer improves molecular validity and E/Z stereochemical accuracy while retaining conformer ensemble fidelity. For protein-ligand complexes, AtomSteer reduces ligand E/Z errors relative to experimental reference structures while achieving competitive docking success and PoseBusters validity. Our method provides a unified framework for adding chemical and crystallographic control to existing atomistic generative models without retraining or architectural changes.

Figures & tables

Figure 1: AtomSteer applies structure-dependent projections during sampling. (1) At state $x_{t}$ , the pre-trained model predicts the endpoint $\hat{x}_{t}^{\rm end}$ (the state at $t=1$ for flows and at noise level $\sigma=0$ for EDM). (2) AtomSteer projects this prediction towards the structural constraint set $\mathcal{C}_{c,a}$ (Section 3.1 ), giving $\tilde{x}_{t}^{\rm end}$ . (3) The native sampler continues from $\tilde{x}_{t}^{\rm end}$ to $x_{t^{\prime}}$ (green); the gray line shows the unsteered trajectory. Fragments show an E / Z double bond in the uncorrected and corrected conformers. Right: vanilla and AtomSteer samples for a molecular conformer, a crystal with a requested Archimedean lattice, and a protein–ligand complex; the gray overlay shows the ground truth.

Projection family	Structural constraint examples	Coordinates corrected	Projection
Chemical Gauss–Newton (ChemGN)	Local chemical geometry and sterics: bond and 1–3 geometry, designated planarity, intramolecular clashes, protein–ligand exclusion, and E / Z dihedral error	Selected molecular or ligand coordinates	Cartesian ChemGN
	Local R / S handedness	Ligand coordinates with first-order bond protection	Bond-tangent ChemGN
Rigid-fragment and rigid-body projections	E / Z configuration	Connected substituent around the double bond	E / Z fragment flips and refinement
	Component packing	Ligand or eligible protein component	Rigid-body $SE(3)$
Crystal symmetry and motif projections	Space-group symmetry	Fractional coordinates and lattice	Hierarchy-aware symmetry selection, orbit projection, and lattice-metric projection
	Lattice-metric relations	Lattice	Metric-template projection

Table 1: Structural constraints grouped by projection family. One requirement can combine families: e.g., E / Z fragment flips and fragment refinement are followed by E / Z ChemGN (Cartesian). Appendices A.2 , A.1 , and D.1 define the residuals and details.

Figure 2: Crystallographic symmetry and requested structures in FlowMM. a, Detected space-group distributions of the MP-20 test set, vanilla FlowMM, and AtomSteer in de novo generation. b, Symmetry success rate of paired de novo samples by number of atoms in the generated unit cell. c, Target-space-group recovery (SG@1) for lattice-informed MP-20 CSP by number of atoms. d, Requested triangular, honeycomb, and Kagome Archimedean lattices (ALs) beside matched vanilla FlowMM and AtomSteer structures; dashed guides in the vanilla panels mark where AtomSteer realizes the AL. Metrics, evaluation subsets, and tolerances: Appendix B.2 .

	As generated					After UMA relaxation
Method	Symmetry success rate (%) $\uparrow$	$d_{\mathrm{SG}}$ $\downarrow$	Unique SGs $\uparrow$	Valid (%) $\uparrow$	Paired $\Delta F_{\max}$ (eV/Å) $\downarrow$	Symmetry success rate (%) $\uparrow$	$d_{\mathrm{SG}}$ $\downarrow$
FlowMM	71.47	0.5541	89	97.82	–	64.19	0.5363
FlowMM + AtomSteer	77.91	0.5083	105	97.79	$-0.0394$	71.50	0.5095

Table 2: Crystallographic symmetry in MP-20 de novo generation. Each method uses 10,000 paired samples. Symmetry success is the fraction of all samples assigned a space-group number

\geq 3

;

d_{\mathrm{SG}}

is the Jensen–Shannon distance to the MP-20 space-group distribution; Unique SGs counts distinct detected groups; and Valid measures as-generated geometric validity.

\Delta F_{\max}

is the median paired change in maximum atomic force relative to FlowMM. Relaxed columns use full-cell UMA relaxation ( Wood et al., 2026 ) . Best results are bolded , and second-best underlined . Details: Appendices B.1 , B.2.1 , and B.3 .

	MR $\uparrow$		SG $\uparrow$		Joint $\uparrow$
Method	@1	@20	@1	@20	@1	@20
CDVAE †	33.90	66.95	–	–	–	–
DiffCSP	52.21	77.77	47.56	75.41	41.50	65.49
CrystalFlow	62.15	78.27	3.06	15.03	1.64	12.77
OMatG	63.50	74.38	1.86	2.30	0.63	1.35
MCFlow-L †	64.08	76.08	–	–	–	–

Table 3: Crystal structure prediction on MP-20 (

N=9{,}046

). Values (%) measure structure match (MR@K), target-space-group recovery (SG@K), and both on one candidate (Joint@K) within

K

samples.

\dagger

marks published MR values ( Xie et al., 2022 ; Jiao et al., 2024a ; Luo et al., 2025 ; Höllmer et al., 2025 ; Seong et al., 2026 ; Miller et al., 2024 ) ; other rows use the common evaluator. AtomSteer (lattice-informed) also receives a reference-derived lattice metric template and integer setting; other rows receive composition only. Best bold , second-best underlined . Evaluation details: Appendix B.2.2 .

AL	Method	AL constraint satisfied (%) $\uparrow$	Prescreening survival (%) $\uparrow$	Joint (%) $\uparrow$	AL detected after relaxation (%) $\uparrow$
Kagome	SCIGEN	98.3	8.0	8.0	10.3
	FlowMM (vanilla)	0.1	45.2	0.1	0.3
	AtomSteer	98.8	11.4	10.9	10.5
Honeycomb	SCIGEN	98.4	7.4	7.4	13.1
	FlowMM (vanilla)	0.0	47.3	0.0	0.0
	AtomSteer	99.2	8.5	7.7	9.4

Table 4: Motif-constrained generation with AtomSteer on FlowMM against SCIGEN. Each row uses all 1,000 planned conditions for its Archimedean lattice (AL), scored by the common evaluator. Joint requires AL constraint satisfaction and prescreening survival on one sample; AL detection after relaxation requires a valid, converged UMA-relaxed structure. Best within each AL and metric bold , second-best underlined . Protocols: Appendices B.2.3 and B.1 .

	a. Conformer generation						b. GFN2-xTB ensemble properties
Method	COV-R $\uparrow$	COV-P $\uparrow$	AMR-R $\downarrow$	AMR-P $\downarrow$	PB-valid $\uparrow$	E / Z $W_{1}$ $\downarrow$ ( ∘ )	$E$ $\downarrow$ (kcal/mol)	$\mu$ $\downarrow$ (D)	$\Delta\epsilon$ $\downarrow$ (kcal/mol)	$E_{\min}$ $\downarrow$ (kcal/mol)
ET-Flow	79.92	73.40	0.458	0.562	93.13	9.04	0.1795	0.2245	0.3397	0.0217
ET-Flow-SS	79.93	73.52	0.456	0.559	93.17	8.94	0.1812	0.2193	0.3413	0.0197
AtomSteer	79.01	75.49	0.457	0.529	96.93	2.15	0.1610	0.1558	0.2354	0.0126

Table 5: Conformer generation and GFN2-xTB ensemble properties on GEOM-Drugs. a, Equal-molecule averages with

2K

generated conformers for

K

references. COV (0.75 Å threshold) and AMR (Å) measure coverage and mean minimum RMSD over references (R) or generated conformers (P). PB-valid is the PoseBusters validity rate of the first

\min(2K,32)

conformers. E / Z

W_{1}

is the Wasserstein distance to GEOM dihedral distributions, median over 153 molecules with an evaluable stereogenic double bond. b, Median absolute errors of Boltzmann-weighted GFN2-xTB ( Bannwarth et al., 2019 ) ensemble properties (mean energy

E

, dipole moment

\mu

, HOMO–LUMO gap

\Delta\epsilon

, minimum energy

E_{\min}

) on a fixed random 100-molecule panel. ET-Flow and ET-Flow-SS are our reruns of the released deterministic and stochastic 50-step samplers of Hassan et al. (2024) ; AtomSteer uses the same pre-trained checkpoint. Best bold , second-best underlined . Protocols: Appendices C.2 and C.1 .

	PoseBusters benchmark			StereoBench
Metric	Vanilla	Boltz- steering	AtomSteer	Vanilla	Guidance only	Boltz- steering	Ishitani & Moriwaki	Chen et al.	AtomSteer
PB-valid $\uparrow$	79.1	99.3	99.3	51.0	90.7	91.7	67.6	87.1	93.4
RMSD $<2$ Å $\uparrow$	83.7	84.1	84.7	63.2	64.0	64.5	62.2	64.2	64.7
Success@1 $\uparrow$	66.1	83.4	84.1	35.7	59.1	59.7	45.7	56.8	60.7
Success@5 $\uparrow$	78.3	87.3	88.7	46.9	66.0	65.8	57.9	63.3	67.5
E / Z median error ( ∘ ) $\downarrow$	6.12	5.26	2.98	6.27	6.44	6.13	11.13	6.73	2.61

Table 6: Protein–ligand steering on the PoseBusters benchmark (301 targets) and StereoBench (1,103 targets). Entries are percentages except median E / Z error. Methods apply inference-time steering to the same pre-trained Boltz-2. Docking success requires RMSD

<2

Å; Success@1 and Success@5 require a PB-valid pose with RMSD

<2

Å (confidence top-1; any of five). Best bold , second-best underlined . Definitions: Appendix D.3 .

Figure 3: Ligand stereochemistry and docking success on StereoBench. (a) Docking success (ligand RMSD $<2$ Å after pocket alignment, light) and joint PB-valid docking success (Success@1, dark). (b) Median per-ligand E / Z dihedral error to the experimental ligand. (c) Example complex (PDB 8ZQ2 ): each method’s pose overlaid on the experimental ligand (gray), with an inset on the stereogenic bond.

Appendix figures & tables19 assets

Supplementary material from the paper’s appendix.

Appendix

\hat{x}_{t}^{\rm end}=\begin{cases}x_{t}+(1-t)v_{\theta}(x_{t},t),&\text{ET-Flow},\\ \operatorname{projx}[x_{t}+(1-t)v_{\theta}(x_{t},t)],&\text{FlowMM},\\ D_{\theta}(x_{\sigma},\sigma),&\text{Boltz-2}.\end{cases}

Appendix

Algorithm 1 AtomSteer sampling from noise to structure (one trajectory)

Task	Endpoint prediction	Active projection families	Projection schedule	Continuation
FlowMM de novo	Velocity extrapolation	Crystal symmetry: inferred orbits and lattice metric	80–100%: every 2.5%	Flow continuation ( 25 )
FlowMM CSP	Velocity extrapolation	Crystal symmetry: lattice-metric template (inferred or supplied) and inferred orbits	Composition-only: Lattice and coordinates: 90%, 100% Lattice-informed: Lattice: every 1% Coordinates: 90%, 100%	Flow continuation ( 25 )
FlowMM motif	Velocity extrapolation	Crystal motif: species, positions, and in-plane cell	Every step	Initial-noise interpolation (Appendix B.1.4 )
ET-Flow conformers	Velocity extrapolation	ChemGN; rigid-fragment E / Z ; global parity selection	E / Z fragment flips: 16–80% E / Z fragment refinement: 32–100% Cartesian ChemGN: 16–100% Final: ChemGN, then parity selection	Initial noise interpolation ( 30 )
Boltz-2 complexes	Denoised EDM estimate	ChemGN (Cartesian and bond-tangent R / S ); rigid-fragment E / Z ; rigid-body ligand and protein components	E / Z fragment flips: 15–80% E / Z fragment refinement: 32.5–100% Cartesian ChemGN: 15–100% Ligand $SE(3)$ : 67.5–100% Protein-component $SE(3)$ : 60–100% Bond-tangent ChemGN ( R / S ): every step Final: ligand, protein components, then R / S	EDM update ( 31 )

Appendix

Table 7: AtomSteer instantiations on the three backbones. Every task runs Algorithm 1 . Percentages give the share of sampling steps completed when a projection acts (FlowMM 1,000, ET-Flow 50, and Boltz-2 200 steps). Projections act at every step within each range unless a cadence is specified. Ranges ending at 100% include the last sampling update. “Final” means after the last sampling update. The schedule column lists molecular projections in their per-step order. Note that a projection happens only when a constraint is applicable to the structure. In FlowMM, the native sampler’s velocity scaling is adopted so AtomSteer applies on top of what FlowMM intended:

Q(A,F,\ell)=(1,1+5t,1+5t)

for de novo and motif-constrained generation and

(1,1+10t,1)

for both CSP variants.

Panel of Table 8 — Table 8: Design choices in constrained generative modeling. Comparison of representative approaches to inference-time control and structure-aware generation. 3 denotes a demonstrated capability, a partial or variant-dependent capability, and 7 a capability not demonstrated in the reported method.

Metric family	Entries of the constraint residual
Triclinic	No additional metric equalities
Monoclinic ( $b$ unique)	$G_{12},\ G_{23}$
Orthorhombic	$G_{12},\ G_{13},\ G_{23}$
Tetragonal ( $c$ unique)	$G_{12},\ G_{13},\ G_{23},\ G_{11}-G_{22}$
Cubic	$G_{12},\ G_{13},\ G_{23},\ G_{11}-G_{22},\ G_{22}-G_{33}$
Hexagonal ( $c$ unique)	$G_{13},\ G_{23},\ G_{11}-G_{22},\ 2G_{12}+G_{11}$

Appendix

Table 9: Lattice-metric constraints. Residual components measure deviations from the metric relations of each family in the stated setting.

Setting	Value
Symmetry detection	Distance tolerances $0.05,0.075,0.10,0.15,0.20,0.30$ Å; angle tolerance $5^{\circ}$
Scales for comparing corrections	Atomic displacement scale $0.20$ Å; cell deformation scale $0.10$ (logarithmic strain)
Symmetry-candidate search	Maximum atomic displacement $0.10$ Å in CSP
Endpoint correction	Maximum atomic displacement $0.20$ Å in de novo generation and CSP
CSP next-state correction	Maximum atomic displacement $0.10$ Å; maximum absolute logarithmic stretch $0.10$
Periodic contacts	$d_{0}=0.5$ Å, with the CSP acceptance rule above

Appendix

Table 10: Numerical settings for AtomSteer projections in FlowMM.

Figure 4: Target-space-group recovery of composition-only AtomSteer by crystal size. Target SG@1 is evaluated on all 9,046 MP-20 targets and grouped by the number of atoms in the generated unit cell.

Figure 5: Motif-constrained generation examples. Each row fixes the target element, nearest-neighbor distance $d$ , and atom count $N$ , then shows the requested AL, matched vanilla FlowMM and AtomSteer structures, and the AtomSteer structure after UMA relaxation. Structures use a shared tilt and length scale within each row.

Setting	Value
Relative distance padding	Bonds and 1–3 pairs: 0.125. Intramolecular clashes: 0.10
Planarity	Allowed distance from the plane: 0.25 Å for ring and double bond groups. Cartesian residuals are excess distances divided by this tolerance
E / Z parameters	Angular scale $s_{\phi}=20^{\circ}$ , Cartesian residual $r^{\rm EZ}=e/s_{\phi}$ , flip acceptance band $5^{\circ}$
E / Z fragment refinement	Eight ADAM steps, learning rate 0.05, maximum applied fraction 0.75. Rotation $\leq 30^{\circ}$ and RMS displacement $\leq 0.75$ Å over moved fragment atoms. Acceptance: increase $\leq 0.01$ in each checked constraint type’s maximum absolute residual and residual $L_{2}$ norm, including R / S
ChemGN	One step and up to 64 largest absolute residuals per stage, damping $10^{-6}$ . During sampling, bound $\\|\Delta x\\|_{2}\leq(0.45/\tau)$ Å, then apply 30% of the bounded correction
ChemGN (terminal)	Combined coordinate norm $\\|\Delta x\\|_{2}\leq 0.04$ Å. Acceptance: increase $\leq 10^{-6}$ in each checked constraint type’s maximum squared residual

Appendix

Table 11: Numerical settings for AtomSteer projections in ET-Flow.

\Delta x

is the bounded ChemGN correction over all atom coordinates, and

\tau

is the next sampling time.

	COV-R (%) $\uparrow$		AMR-R (Å) $\downarrow$		COV-P (%) $\uparrow$		AMR-P (Å) $\downarrow$
Method	Mean	Median	Mean	Median	Mean	Median	Mean	Median
ET-Flow	79.92	84.52	0.458	0.436	73.40	80.31	0.562	0.500
ET-Flow-SS	79.93	84.57	0.456	0.435	73.52	80.19	0.559	0.497
AtomSteer	79.01	82.90	0.457	0.440	75.49	82.65	0.529	0.463

Appendix

Table 12: Conformer coverage and matching across 998 GEOM-Drugs test molecules. Each metric is computed per molecule, with means and medians reported across molecules. COV uses a 0.75 Å threshold. Best results are bold and second-best results are underlined .

Figure 6: Examples of constraint corrections on GEOM-Drugs. ET-Flow (purple) and AtomSteer (green) conformers are aligned to a common GEOM reference (gray). Highlighted atoms indicate the evaluated constraint. Annotations report PoseBusters validity (green = valid, red = invalid) and the E / Z dihedral error, nonbonded distance, or bond length. PoseBusters settings here do not check E / Z configuration, so stereochemical violations may still pass its checks. Each example pairs an ET-Flow violation with a nearby PB-valid AtomSteer conformer satisfying the highlighted constraint.

Figure 7: AtomSteer conformer ensembles for five GEOM-Drugs molecules. Each row shows the molecular graph, three distinct AtomSteer conformers, and the GEOM reference nearest to the first. Annotations report PoseBusters validity and symmetry-corrected heavy-atom RMSD to each conformer’s nearest reference. All shown conformers are PB-valid, with nearest-reference RMSDs of 0.25–0.75 Å.

Setting	Value
Relative distance padding	Bonds and 1–3 pairs: 0.125. Intramolecular clashes: 0.10. Protein–ligand exclusion: 0.225
Planarity	Allowed distance from the plane: 0.25 Å for ring and double bond groups, 0.02 Å for terminal rings. Cartesian residuals are excess distances in Å
E / Z parameters	Angular scale $s_{\phi}=20^{\circ}$ , Cartesian residual $r^{\rm EZ}=2e/s_{\phi}$ , flip acceptance band $5^{\circ}$
E / Z fragment refinement	Eight ADAM steps, learning rate 0.05, maximum applied fraction 0.75. Rotation $\leq 30^{\circ}$ and RMS displacement $\leq 0.75$ Å over moved fragment atoms. Acceptance: increase $\leq 0.01$ in each checked constraint type’s maximum absolute residual and residual $L_{2}$ norm, including R / S
Cartesian ChemGN	One step and up to 64 largest absolute residuals per stage, damping $10^{-6}$ . During sampling, bound $\\|\Delta x\\|_{2}\leq(0.45/\alpha_{k})$ Å, then apply 30% of the bounded correction
Rigid ligand	Translation $\leq 0.45$ Å, rotation $\leq 5^{\circ}$ , RMS displacement $\leq 0.50$ Å

Appendix

Table 13: Numerical settings for AtomSteer projections in Boltz-2.

Metric	Vanilla	Boltz-steering	AtomSteer
lDDT-PLI $\uparrow$	0.73	0.74	0.73
Complex lDDT $\uparrow$	0.95	0.95	0.95
Backbone lDDT $\uparrow$	0.98	0.98	0.98
TM-score $\uparrow$	0.98	0.98	0.98
GDT-TS $\uparrow$	0.97	0.97	0.97
C $\alpha$ RMSD (Å) $\downarrow$	0.96	0.96	0.97

Appendix

Table 14: OpenStructure evaluation of the same predictions on StereoBench. Computed with OpenStructure 2.12.0. The interface rows use the StereoBench multi-chain subset. OpenStructure contains no stereochemical term, so these scores are a preservation check on protein and interface quality. Best in bold .

		All E / Z bonds $\leq 10^{\circ}$ (%) $\uparrow$			Median E / Z error ( ∘ ) $\downarrow$
Number of E / Z bonds	$N$	Vanilla	Boltz-steering	AtomSteer	Vanilla	Boltz-steering	AtomSteer
One bond	813	64.6	66.2	81.2	5.63	5.53	2.61
Two bonds	192	53.7	48.4	75.0	6.46	6.96	2.83
Three or more	98	29.6	25.5	70.4	11.78	8.06	2.19

Appendix

Table 15: E / Z accuracy by number of stereogenic double bonds on StereoBench. Best in bold , second-best underlined .

Threshold $\uparrow$	Vanilla	Boltz-steering	AtomSteer
$\leq 2.5^{\circ}$	21.7	21.8	47.1
$\leq 5^{\circ}$	39.1	39.1	63.4
$\leq 10^{\circ}$	59.6	59.5	79.1
$\leq 15^{\circ}$	69.2	70.8	85.9
$\leq 20^{\circ}$	73.3	78.6	89.8
$\leq 30^{\circ}$	77.2	92.4	93.0

Appendix

Table 16: E / Z threshold sensitivity on StereoBench. Percentage of the 1,103 ligands for which every evaluable stereogenic double bond falls within the given threshold of the experimental ligand, the criterion of Table 6 at

10^{\circ}

. Best per row in bold , second-best underlined .

Split	Metric	Vanilla	Boltz-steering	AtomSteer
Before 2023-06-01 (N=926)	PB-valid $\uparrow$	51.2	90.6	92.5
	RMSD $<2$ Å $\uparrow$	67.9	69.7	69.9
	Success@1 $\uparrow$	38.1	64.0	65.1
	E / Z median error ( ∘ ) $\downarrow$	6.02	5.80	2.48
	All E / Z bonds $\leq 10^{\circ}$ $\uparrow$	61.0	60.7	80.5
	Success@1+stereo $\uparrow$	27.8	39.8	49.7

Appendix

Table 17: StereoBench split by PDB release date at the Boltz-2 training cutoff of 2023-06-01. Same protocol, selector, and metric definitions as Table 6 , recomputed within each split of the 1,103 targets. Best in bold , second-best underlined .

Figure 8: Examples of AtomSteer correcting E / Z stereochemistry and pose validity. Each panel overlays a confidence top-1 ligand pose with the experimental ligand (gray), after protein alignment. PDB IDs identify the targets, and circular insets compare all three methods in the binding pocket. Annotations report PoseBusters validity (green = valid, red = invalid), ligand RMSD to the experimental reference (Å), and E / Z dihedral error ( ∘ ). Methods shown: Vanilla (purple), Boltz-steering (blue), and AtomSteer (green).

Figure 9: Examples where AtomSteer preserves accurate E / Z geometry and ligand placement. All methods have ligand RMSD $\leq 0.6$ Å, and most poses are PB-valid. Layout and annotations follow Figure 8 .

Figure 10: Examples where Boltz-steering achieves lower ligand RMSD than AtomSteer. All methods have E / Z dihedral errors within $10^{\circ}$ . For 7G4D, ligand RMSDs are 1.2 Å (Boltz-steering) and 11.5 Å (AtomSteer); for 7G6A, they are 1.1 Å and 5.3 Å, respectively. Layout and annotations follow Figure 8 .

Constrained Generative Modeling

Constrained generative modeling develops methods that produce samples satisfying structural, physical, semantic, or user-defined requirements.

See all 191 papers

Sep 9, 2026

uFlowCSP: Crystal Structure Prediction using Mean flow generative models

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.

Sourin Dey, Dipannoy Das Gupta, Lai Wei +2

Department of Computer Science and Engineering University of South Carolina Columbia, SC 29201 · Department of Computer Science and Cybersecurity University of North Georgia Columbia, SC 29201

Crystal Structure Prediction One-Step Generative Modeling Crystal Structure Generation Flow-Based Generative Modeling Conditional Generative Modeling Computational Materials Science

Ask PDF

Feb 23, 2026

Multimodal Crystal Flow: Any-to-Any Modality Generation for Unified Crystal Modeling

Crystal modeling spans a family of conditional and unconditional generation tasks, including crystal structure prediction (CSP) and de novo generation (DNG). While recent deep generative models have shown promising performance, they remain largely task-specific, lacking a unified framework that shares crystal representations across tasks. To address this limitation, we propose Multimodal Crystal Flow (MCFlow), a unified multimodal flow model that realizes multiple crystal generation tasks as distinct inference trajectories via independent time variables for atom types and crystal structures. To enable multimodal flow in a standard transformer model, we introduce a composition- and symmetry-aware atom ordering with hierarchical permutation augmentation, injecting compositional and crystallographic priors without explicit structural templates. Experiments on the MP-20 and MPTS-52 benchmarks show that a single MCFlow model is competitive with task-specific baselines across CSP, DNG, and structure-conditioned atom type generation.

Kiyoung Seong, Sungsoo Ahn, Sehui Han +1

Graduate School of AI, KAIST, Seoul, South Korea · Materials Intelligence Lab, LG AI Research, Seoul, South Korea

Crystal Structure Generation Unified Multimodal Models Crystal Structure Prediction Flow-Based Generative Modeling

Ask PDF

Apr 2, 2026

Crystalite: A Lightweight Transformer for Efficient Crystal Modeling

Generative models for crystalline materials often rely on equivariant graph neural networks, which capture geometric structure well but are costly to train and slow to sample. We present Crystalite, a lightweight diffusion Transformer for crystal modeling built around two simple inductive biases. The first is Subatomic Tokenization, a compact chemically structured atom representation that replaces high-dimensional one-hot encodings and is better suited to continuous diffusion. The second is the Geometry Enhancement Module (GEM), which injects periodic minimum-image pair geometry directly into attention through additive geometric biases. Together, these components preserve the simplicity and efficiency of a standard Transformer while making it better matched to the structure of crystalline materials. Crystalite achieves state-of-the-art results on crystal structure prediction benchmarks, and de novo generation performance, attaining the best S.U.N. discovery score among the evaluated baselines while sampling substantially faster than geometry-heavy alternatives.

Tin Hadži Veljković, Joshua Rosenthal, Ivor Lončarić +1

1UvA-Bosch Delta Lab · University of Amsterdam · 3Ruđer Bošković Institute

Diffusion Transformer Transformer Crystal Structure Prediction Crystal Structure Generation

Ask PDF

Abstract

Constrained Generative Modeling

Explore similar work

uFlowCSP: Crystal Structure Prediction using Mean flow generative models

Multimodal Crystal Flow: Any-to-Any Modality Generation for Unified Crystal Modeling

Crystalite: A Lightweight Transformer for Efficient Crystal Modeling