GLASS: Global Latent Aggregation with Slot-based Set Decoding for Scalable All-Atom Crystal Generation

Hendrik Kraß; Seyed Mohamad Moosavi; Mathias Niepert

cs.LGSep 29, 2026

GLASS: Global Latent Aggregation with Slot-based Set Decoding for Scalable All-Atom Crystal Generation

Authors: Hendrik Kraß, Seyed Mohamad Moosavi, Mathias Niepert

Organizations: Institute for Artificial Intelligence University of Stuttgart · NEC Labs Europe

Abstract

Generative models for crystals enable the discovery of novel structures, but scaling all-atom generation to larger systems such as metal--organic frameworks remains challenging. We connect this difficulty to the correspondence problem of particle-space generation. Even on a single fixed target set, index-free permutation-equivariant particle flows require substantially more training for reliable generation as set size and density increase, under both independent and optimal-transport couplings. To resolve this challenge, we introduce GLASS---Global Latent Aggregation with Slot-based Set Decoding, which encodes structures in a permutation-invariant global latent space and learns their distribution via flow matching. A learned-slot decoder constructs all atoms in parallel, removing atom-wise correspondence from generative transport. On MP20, GLASS is competitive with particle-space models, and flow training can reach the validity of the training data at every structure size. On a QMOF subset, GLASS generates MOFs with up to 150 atoms per unit cell without conditioning on building blocks, topology, or composition, and approaches the structural validity of the training data. On both datasets, flow training exposes a validity--novelty tradeoff, and MOF novelty remains limited by autoencoder generalization on the available data. These results show that separating correspondence assignment from generative transport provides a simple route toward high-validity generation of larger atomistic systems.

Figures & tables

Figure 1: Training requirements and site coverage in index-free permutation-equivariant particle flows. a) Target grids and generated samples from independent and optimal-transport (OT) couplings. Hollow markers indicate generated particles sharing a nearest target site. b) Left and center: training steps to whole-set failure rates below

50\,\%

and

10\,\%

. Right: mean unoccupied-site fraction at final checkpoints. The slot decoder realizes each target from learned slots with Hungarian matching and reaches every threshold at its first evaluation. Curves summarize three runs; shading denotes one standard deviation. Hollow triangles mark unreached thresholds.

Figure 2: GLASS architecture. A Transformer encoder and invariant attention pooling map a crystal to a global latent vector. Learned decoder slots predict atom types and coordinates in parallel, while a separate head predicts the lattice. Flow matching generates new latents, which are decoded into all-atom structures.

Figure 3: Validity by size and local geometric fidelity on MP20-generated structures. (a) LeMat-GenBench validity before relaxation as a function of atom count, including GLASS at 50k and 1M flow-training steps, each with

10,000

samples. Bands show 95% Wilson intervals; the vertical axis uses an arcsinh scale that expands differences near

100\,\%

. (b) Radial distribution functions for 13–16 atoms, compared with MP20 training data (grey). The red shaded interval marks interatomic distances below

0.5\,\text{\AA}

Figure 4: Structural validity of raw and relaxed all-atom MOF generations. GLASS uses the 1M-step flow checkpoint. (a) MOFChecker validity by atom count before and after relaxation (opt), compared with QMOF150 training data; bands show 95% Wilson intervals. ADiT and SinAE are absent because no released QMOF samples or checkpoints were available. (b) Raw-generation validity and individual checks (%). Arrows indicate the preferred direction; defect flags can overlap. ∗ Mofasa: precomputed MOFChecker results for 10,000 released samples with 20–150 atoms and their relaxed counterparts; trained on QMOF structures up to 170 atoms with a different split. All models are shown in Appendix G .

Appendix figures & tables15 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 5: Slot-decoder outputs on the fixed target grids. For each displayed

N

, the run with the largest final PI-RMSD among three seeds, stated relative to the minimum target-site spacing. No target site is shared.

$N$	Coupling	64 steps	1024 steps	4096 steps
4	Independent	0.00	0.00	–
4	OT	0.00	0.00	–
24	Independent	10.16	5.86	5.86
24	OT	20.31	19.92	19.92
64	Independent	80.08	58.20	57.03
64	OT	88.28	87.50	87.50

Appendix

Table 2: Failure rates (%) at different sampling resolutions for one training run per setting. Dashes denote resolutions not evaluated.

	Training split		Validation split
Latent dim.	RMSD (Å)	Type acc. (%)	RMSD (Å)	Type acc. (%)
32	0.0905	99.92	0.1098	97.05
64	0.0584	99.72	0.0637	98.85
128	0.0434	99.75	0.0460	99.49

Appendix

Table 10: LeMat-Bulk autoencoders. Reconstruction over the training (5,101,846 structures) and validation (

10,000

structures) splits. Invalid predicted cells, at most

0.05\,\%

of either split, are excluded from RMSD.

Figure 6: MP20 validity–novelty tradeoff. Raw validity, novelty among valid candidates, and their joint yield across flow checkpoints. Curves show means and sample standard deviations over three training runs, with

2,500

samples per run and checkpoint.

Figure 7: MP20 validity and radial distribution functions by size. RDFs show interatomic pair-distance distributions for generated structures and the training reference.

Sampler	Churn	Anti-ann.	Overall	1–4	5–8	9–12	13–16	17–20
Vanilla	No	No	67.03	98.17	92.65	73.25	37.29	16.67
Anti-annealing only	No	Yes	66.82	98.11	92.26	73.13	37.59	15.90
Churn only	Yes	No	94.39	98.17	95.79	94.86	92.40	89.48
Production	Yes	Yes	94.44	98.11	96.00	94.86	92.40	89.54

Appendix

Table 14: Crystalite sampler settings on MP20. LeMat-GenBench total validity (%) for the four Crystalite sampler settings on 10,000 matched MP-20 samples per setting. Churn uses

S_{\mathrm{churn}}=60

; anti-annealing is applied to coordinates and lattices with

\rho=10

. Bold values are best in each column.

Figure 8: QMOF150 validity–novelty tradeoff. Raw validity, MOFid novelty, and valid-and-novel yield under permissive (left) and strict (right) identifiers. Means and sample standard deviations over three training runs, with

10,000

samples per run and checkpoint; all rates use requested-sample denominators.

Figure 9: QMOF150 novel yield by size at 1M steps. Valid-and-novel yield and validity among identified novel candidates under each MOFid convention. Identifier novelty can include close variants of training frameworks.

Flow step	Valid (%)	Matched $\mid$ valid (%)	Valid, unmatched (%)
250k	22.57	90.27	2.20
500k	51.93	97.00	1.56
1M	74.79	98.29	1.28
Mofasa ∗	54.03	0.15	53.95

Appendix

Table 17: StructureMatcher comparison of valid generations with QMOF150 training structures. GLASS values are means over three training runs with

10,000

raw samples each. Validity and valid-unmatched yield use all requested samples; the match rate is conditional on validity. ∗ 10,000 released samples; different training split and size range.

Figure 10: MP20 latent-space visualization. PCA of MP20 autoencoder training latents, colored by space group (a) and atom count (b). Axis labels report explained variance.

Figure 11: QMOF150 training latents. PCA colored by the most frequent metal (a), permissive MOFid net (b), and atom count (c); grey denotes unassigned nets. One region contains predominantly Zn/ pcu and Al/ rna frameworks, while other classes overlap.

Figure 12: Learned MP20 slot correspondences. Atoms assigned to each decoder slot across $1,000$ training structures, showing the fractional coordinates. Empty slots are excluded. Spatial patterns indicate specialization within the supplied coordinate frame.

Dataset	Batch size	Total (s)	Flow (s)	Decoder (s)	Peak memory (GiB)
MP20	32	1.74	1.72	0.022	0.43
	128	0.67	0.66	0.012	0.46
	256	0.39	0.38	0.009	0.50
	512	0.27	0.26	0.009	0.57
	1,000	0.21	0.20	0.008	0.71
QMOF150	32	1.86	1.71	0.139	0.53

Appendix

Table 18: Sampling cost per 1,000 structures on one full H100. Median seconds over five repetitions; peak memory is allocated GPU memory.

Figure 13: Generated MP20 structures and training analogues. Two valid examples per size bin (1–5, 6–10, 11–15, and 16–20 atoms), each beside its nearest training structure under the composition and pair-distance fingerprint. Bonds are distance-based visual guides.

Figure 14: Generated QMOF150 structures and training analogues. Each row covers one size range. Left: a valid generation and its StructureMatcher training match. Right: a valid unmatched generation and its nearest fingerprint analogue. These unmatched examples retain the training framework with substituted metal sites.

Crystal Structure Generation

Crystal structure generation uses computational models to create periodic atomic arrangements for predicting and designing materials.

See all 61 papers

uFlowCSP: Crystal Structure Prediction using Mean flow generative models
Sep 9, 2026Sourin Dey, Dipannoy Das Gupta, Lai Wei +2Crystal Structure Prediction One-Step Generative Modeling
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
Multimodal Crystal Flow: Any-to-Any Modality Generation for Unified Crystal Modeling
Feb 23, 2026Kiyoung Seong, Sungsoo Ahn, Sehui Han +1Crystal Structure Generation Unified Multimodal Models
Graduate School of AI, KAIST, Seoul, South Korea · Materials Intelligence Lab, LG AI Research, Seoul, South Korea
MolCrystalFlow: Molecular Crystal Structure Prediction via Flow Matching
Feb 17, 2026Cheng Zeng, Harry W. Sullivan, Thomas Egg +8Crystal Structure Prediction Flow-Based Generative Modeling
Department of Chemistry, University of Florida, Gainesville, FL 32611, USA · Quantum Theory Project, University of Florida, Gainesville, FL 32611, USA · Department of Chemical Engineering and Material Science, University of Minnesota - Twin Cities, Minneapolis, MN 55455, USA +6

Sep 9, 2026J/K move · Enter open · S save

uFlowCSP: Crystal Structure Prediction using Mean flow generative models

Sourin Dey, Dipannoy Das Gupta, Lai Wei, Sadman Sadeed Omee+1

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

Open paper AskPDF arXiv

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.

Abstract

Crystal Structure Generation

Explore similar work

uFlowCSP: Crystal Structure Prediction using Mean flow generative models

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

MolCrystalFlow: Molecular Crystal Structure Prediction via Flow Matching