Zatom-2: Multitask Pretraining on Atomistic Data for Generative Modeling across Domains

Miruna Cretu; Alex Abrudan; Antonia Panescu; Tynan Perez; Rishabh Anand; N. Benjamin Erichson; Michael W. Mahoney; Samuel Blau

cs.LGOct 8, 2026

Zatom-2: Multitask Pretraining on Atomistic Data for Generative Modeling across Domains

Authors: Miruna Cretu, Alex Abrudan, Antonia Panescu, Tynan Perez, Rishabh Anand, N. Benjamin Erichson, Michael W. Mahoney, Samuel Blau, +6 more

Organizations: University of Cambridge · Yale University · MIT · Lawrence Berkeley National Lab · International Computer Science Institute · UC Berkeley

Abstract

Unified atomistic modeling has the potential to accelerate discovery in chemistry, materials science, and biology by bridging data-rich chemical domains and data-scarce biological contexts. However, existing generative approaches to atomistic modeling remain highly specialized to scientific disciplines (chemistry vs. biology) or do not leverage both high-volume organic (molecule) and inorganic (material) data for general-purpose pretraining. To this end, we introduce Zatom-2, an atomistic generative model pretrained on approximately five million structures from the OMol25 and OMat24 electronic structure datasets. Zatom-2 features a multiscale Transformer architecture coupled with conditional flow matching that supports force conditioning and foundational pretraining tasks such as generation, structure prediction, and prediction of molecular and material energies and forces. Empirically, Zatom-2 achieves better molecular distribution fidelity than Zatom-1 and achieves strong performance on existing molecule and material generation benchmarks. Zatom-2 demonstrates the ability to control sample generation across low- and high-force regimes, and enhances protein generation in a low-data setting through joint generative-predictive pretraining and transfer learning, increasing protein backbone designability in a length extrapolation setting from 67.8% without pretraining to 74.8% after finetuning on 2,000 protein domains.

Figures & tables

Panel of Figure 1 — Figure 1: Zatom-2 multitask pretraining across domains. Zatom-2 is jointly trained on OMol25 (4M) and OMat24 (1M) as a generative model. Its input initializer receives task-dependent conditioning, summarized by task strips: checkmarks indicate enabled inputs and dashes indicate disabled inputs. MLIP task supervision (force and energy prediction) uses clean coordinates at $t=1$ and known atom types, without bond or system force conditioning. The initializer provides atom $(C_{L},P_{LL})$ and token $(S_{I},Z_{II})$ conditioning through the dotted connections. Noisy coordinates $\mathbf{x}_{t}$ and flow time $t$ enter the U-shaped atom encoder—token trunk—atom decoder. Atom features $Q_{L}$ and token features $Q_{I}$ feed task-specific readout heads. The recycling path converts the clean-coordinate estimate obtained from $\widehat{\mathbf{v}}_{\theta}$ into a token-level distogram, updating token-pair conditioning for a second forward pass. In this work, we experiment with $N=9/18/36$ trunk layers.

		Datasets		Tasks
Model	Size	OMol25	OMat24	Generation	Structure prediction	Force & energy prediction	Force conditioning
Zatom-1 (released)	Large	$\checkmark$	–	$\checkmark$	–	–	–
Zatom-2 (I)	Small	$\checkmark$	–	$\checkmark$	–	–	–
Zatom-2 (II)	Small	$\checkmark$	–	$\checkmark$	–	–	$\checkmark$
Zatom-2 (III)	Small	$\checkmark$	–	$\checkmark$	$\checkmark$	–	$\checkmark$
Zatom-2 (IV)	Base	$\checkmark$	$\checkmark$	$\checkmark$	$\checkmark$	–	$\checkmark$

Table 1: Pretraining configurations and generation quality. Panel (a) shows multiple training configurations for OMol25 and OMat24 mixed training. Panel (b) reports their performance, where we also compare to 10,000 dataset samples and Zatom-1. All models are trained for 80 epochs. Forces are calculated using UMA ( Wood et al., 2025 ) and metrics are reported in eV Å -1 . AUROC uses a conditioning mean force norm cutoff of 1 eV Å -1 . Structure prediction results are reported in Appendix Table 9 .

Figure 2: Force-conditioning consistency. Panel (a) compares the conditioned mean force norm with the mean force norm evaluated by UMA for unconditional OMol25 molecule generation. Panels (b) and (c) give the corresponding comparison for OMol25 structure prediction on the GEOM ORCA6 and ANI2x subsets, respectively. Red dotted lines mark the AUROC cutoff (1 eV Å -1 in (a)–(b), and 2.5 eV Å -1 in (c)).

Model	Valid $\uparrow$	Unique $\uparrow$	PB-valid $\uparrow$	Valid & unique $\uparrow$	Valid & PB-valid $\uparrow$
EQGAT-diff	94.6	100.0	59.7	94.60	56.48
SemlaFlow	93.9	100.0	87.5	93.90	82.16
ADiT	95.3	100.0	85.3	95.30	81.29
TABASCO	97.6	99.03	91.6	96.65	89.40
Zatom-1	93.6	99.93	94.1	93.53	88.08
GEOM-Drugs (reference)	–	–	94.0	–	–

Table 2: Generation benchmarks on GEOM-Drugs and MP20. Panel (a) reports unconditional GEOM-Drugs molecule generation metrics, and panel (b) reports MetaSUN yield for unrelaxed MP20 material generation.

Figure 3: Model and data scaling across training epochs. We evaluate Zatom-2 models trained either on 500k examples per dataset or on the full OMol25 and OMat24 datasets. We report (a) OMol25 molecular generation quality measured by AFD, (b) OMat24 structure prediction measured by top-1 match rate, and force prediction MAE on (c) OMol25 and (d) OMat24. Training on the full dataset with more model parameters consistently improves validation performance across tasks.

(a) In-distribution: 50–128 residues
Model / Zatom-2 initialization	Designability (%) $\uparrow$	Per-sequence success (%) $\uparrow$	Diversity (#clusters) $\uparrow$	Novelty (%) $\uparrow$
RFdiffusion3	$\mathbf{91.07\pm 0.63}$	$\mathbf{67.61\pm 1.18}$	$435.0\pm 11.5$	$13.33\pm 1.25$
None	$88.22\pm 0.26$	$61.35\pm 0.14$	$400.3\pm 5.0$	$11.41\pm 0.77$
Zatom-2 (IV): Generation + structure (w/o FC)	$86.30\pm 1.12$	$59.99\pm 1.61$	$432.0\pm 6.1$	$16.11\pm 1.32$
Zatom-2 (IV): Generation + structure	$89.68\pm 0.84$	$62.40\pm 2.08$	$439.0\pm 5.0$	$14.26\pm 0.71$
Zatom-2 (V): Generation + structure + F/E pred.	$89.09\pm 1.61$	$63.39\pm 1.98$	$\mathbf{453.7\pm 9.0}$	$\mathbf{16.77\pm 1.04}$

Table 3: SCOPe-2k protein generation and transfer. (a): We evaluate 1,027 protein backbones for each configuration, 13 corresponding to each length between 50 and 128 (with means

\pm

standard deviations over 3 random seeds). (b): The same protocol is repeated for 129–256 residue lengths. FC denotes force-conditioning, and F/E pred. refers to a model trained with force and energy prediction as auxiliary tasks (Zatom (V)). The Zatom (V)-finetuned model substantially outperforms all configurations on the length extrapolation task. Metrics are explained in Appendix D.9 .

Figure 4: Protein latent space organization and localized decoding responses. Left two: t-SNE views of final token features, colored by amino acid. Right two: increases in sequence and structure reconstruction losses after final-token perturbation; solid/dashed lines denote the target/other amino acids. Bands show 95% intervals from 2,000 protein bootstrap resamples. Pretrained Zatom-2 is notably more sensitive to local amino acid-specific perturbations than without pretraining.

Appendix figures & tables25 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 5: Mean force norm distributions across OMol25-4M subsets.

Figure 6: Mean force norm distributions across OMat24-1M subsets. Data subsets are grouped by generation protocol (rattled and then re-relaxed, rattled at various temperatures, or ab initio molecular dynamics trajectories; see Barroso-Luque et al. (2024) for more details). The legend shows subset medians in eV Å -1 .

Figure 7: Atom count distributions of OMol25-4M & OMat24-1M. Hydrogens are included.

Dataset	Source data terms
OMol25	Creative Commons Attribution 4.0 (CC BY 4.0)
OMat24	CC BY 4.0
QM9	CC0 1.0 (original Figshare data deposit)
MP20	MIT for the processed ADiT release used by Zatom-1; underlying Materials Project data are CC BY 4.0
GEOM-Drugs	CC0 1.0 (Harvard Dataverse deposit)
SCOPe / ASTRAL	Provider states that all data are freely available to all users; no standard named license is specified

Appendix

Table 4: Dataset licenses and access terms. Dataset names link to the corresponding release or the provider’s usage statement.

	RFdiffusion3	Emyx	Zatom-2
1	Build local atom neighborhoods (128 keys); use global token attention.	Build sparse atom/token graphs prioritizing sequence, bonds, ligand, and motif edges, then spatial neighbors.	Build local atom neighborhoods (128 keys, token-index/spatial neighbors); use global token attention.
2	Embed atom/token metadata; mix dense pair features with two triangle-free Pairformer blocks; pool atom pairs into token pairs.	Bottleneck-embed atom14 atom/token and global features; project relative positions and bonds to edge biases.	Bottleneck-embed atom1/atom14 metadata, domain, and charge/spin; residually-embed mean force norm; map bonds, relative positions, and any available fixed/reference geometry to head biases.
3	Add linear embeddings of $x_{\sigma}/\sqrt{\sigma^{2}+\sigma_{\rm data}^{2}}$ ; condition on Fourier features of $\frac{1}{4}\log(\sigma/\sigma_{\rm data})$ .	Add Fourier coordinate embeddings; average geometry-aware conditioning with sinusoidal time features.	Add linear atom and mean-pooled linear token embeddings of $x_{t}/\sigma_{\rm data}$ ; average each conditioning stream with Fourier features of $t$ .
4	Cache $q_{e}=\mathrm{Enc}(q_{0};c,p)$ and $a_{e}=\mathrm{Down}(q_{e},a_{0};s)$ before recycling.	Begin recycling before the atom encoder.	Begin recycling before the atom encoder.
5	Reset $q,a\leftarrow q_{e},a_{e}$ ; concatenate $z_{0}$ , noisy-input and recycled distograms; apply transitions and two triangle-free Pairformer blocks to obtain $s,z$ .	Reset $q,a$ to initial states plus projected, normalized $\operatorname{sg}(q_{-},a_{-})$ ; set $z=z_{0}+WD(\hat{x}_{-})$ on sparse edges.	Reset $q,a\leftarrow q_{0},a_{0}$ ; set $z=z_{0}+WD(\hat{x}_{-})$ for all token pairs. No node-state feedback.
6	Reuse cached encoder/downcast outputs.	$q\leftarrow\mathrm{Enc}(q;c,p)$ ; $a\leftarrow\mathrm{Down}(q,a)$ .	$q\leftarrow\mathrm{Enc}(q;c,p)$ ; $a\leftarrow\mathrm{Down}(q,a)$ .

Appendix

Table 5: Aligned forward pass operations. Enc, Trunk, and Dec are attention stacks; Down/Up are gated cross-attention between input resolutions. By definition, no recycling feedback is used in the first pass.

Hyperparameter	Setting
Approximate model size (S/base/L)	70M / 140M / 270M parameters
Token-trunk blocks (S/base/L)	9 / 18 / 36
Atom encoder / decoder blocks	3 / 3
Atom / token / token-pair widths	128 / 768 / 256
Atom / token attention heads	4 / 12
Attention	Local atom attention (128 keys); global token attention

Appendix

Table 6: Key Zatom-2 hyperparameters. S/base/L vary the token trunk depth while retaining the same feature widths. Token budgets are per GPU before diffusion batch/view augmentation; a smaller token budget is paired with twice as much gradient accumulation.

Model training configuration	Budget / endpoint	GPUs	GPU hours
Small, OMol25 G, no FC	80 epochs	16	1,830
Small, OMol25 G, FC	80 epochs	16	1,820
Small, OMol25 G+S, FC	80 epochs	64	2,280
Base, OMol25+OMat24 G+S, FC	80 epochs	64	2,740
Base, OMol25+OMat24 G+S+P, FC	80 epochs	64	2,850
Large, OMol25+OMat24 G+S+P, FC	80 epochs	64	4,090

Appendix

Table 7: Approximate 80GB A100 GPU hours by training configuration. G/S/P denote generation, structure prediction, and MLIP prediction; FC denotes force conditioning. Epoch budgets count completed epochs; MLIP finetuning runs end at raw epoch 119. All finetuning costs exclude pretraining the corresponding initial model weights.

$N$	OMol25 4m / ct-scd-omol25	OMat24 1m / ct-scd-amp20
100	$0.02321\pm 0.00021$	$0.0776\pm 0.0042$
500	$0.005461\pm 0.000078$	$0.01691\pm 0.00091$
1,000	$0.002527\pm 0.000032$	$0.008268\pm 0.00034$
5,000	$0.0000916\pm 0.0000047$	$0.001236\pm 0.000027$

Appendix

Table 8: AFD baselines for encoder dataset pairs used in this work (finite-

N

floor): Each entry is the mean

\pm

standard deviation of the AFD score between two non-overlapping ground-truth reference sets of size

N

, over 5 replicates (no structure is shared within or across replicates at a given

N

	OMol25
Model	Mean RMSD $\downarrow$	Median RMSD $\downarrow$	Mean $\overline{\\|F\\|}$ $\downarrow$	Target MAE $\downarrow$	AUROC $\uparrow$
Zatom-2 (III)	1.739	1.714	2.271	0.725	0.798
Zatom-2 (IV)	1.708	1.675	2.138	0.787	0.787
Zatom-2 (V)	1.762	1.729	1.953	0.637	0.801
Zatom-2 (VI)	1.723	1.743	1.211	1.214	0.818

Appendix

Table 9: Molecule/material top-1 structure prediction after pretraining. All models use the 80-epoch pretrained checkpoints described in Table 1 . OMol25 RMSDs are reported in Å and forces are calculated with UMA. Compared to OMol25, OMat24 structure prediction is more sensitive to outliers, so we report median metrics for OMat24. OMol25 AUROC uses a target mean force norm cutoff of 1 eV Å -1 .

	OMol25
	Generation				Structure prediction
Model	AFD $\downarrow$	Mean $\overline{\\|F\\|}$ $\downarrow$	Target MAE $\downarrow$	AUROC $\uparrow$	Mean RMSD $\downarrow$	Median RMSD $\downarrow$	Mean $\overline{\\|F\\|}$ $\downarrow$	Target MAE $\downarrow$	AUROC $\uparrow$
Zatom-2 (IV) + FT	0.017	1.205	0.540	0.807	1.771	1.747	1.570	0.632	0.890
Zatom-2 (V) + FT	0.016	1.137	0.530	0.809	1.729	1.706	1.458	0.690	0.828
Zatom-2 (V) + FT (extended)	0.015	1.040	0.526	0.840	1.720	1.713	1.178	0.871	0.882

Appendix

Table 10: MLIP accuracy and retention of generative capabilities. “+ FT” denotes 40 epochs of finetuning with 80% of training steps allocated to MLIP prediction and 10% for generation and structure prediction each; here, “extended” means adding 120 additional MLIP-heavy finetuning epochs. Model names follow Table 1 . Panel (a) reports generation and top-1 structure prediction results after finetuning, with forces calculated using UMA; corresponding pretrained results appear in Tables 1 and 9 . For OMat24, we observe that generation and structure prediction are sensitive to outliers and we report median values. Panels (b)–(c) report validation MAEs (Appendix D.8 ); external baselines are published values with partially differing training and evaluation protocols.

Figure 8: SCOPe-2k convergence From left to right, the panels show best-of-eight designability (CA scRMSD

<1.5

Å), the mean maximum exact, query-normalized TM-align score against the canonical SCOPe-2k training set among designable backbones (lower is more novel), and distinct joint sequence–structure clusters among designable backbones. Joint clusters require TM

\geq 0.5

together with at least 10% sequence identity and 70% coverage. Each point uses 96 generated backbones and eight refolds per backbone.

Model	Validity $\uparrow$
Equivariant Diffusion	91.90
Symphony	83.50
GeoLDM	93.80
ADiT (QM9-only)	92.19
ADiT (joint QM9+MP20)	94.45
Zatom-1 (QM9-only, 80M)	92.88

Appendix

Table 11: QM9 generation results. Zatom-2 is trained jointly on QM9 and MP20 (Appendix B ). Published baselines are from Table 2 of Morehead et al. (2026) (arXiv v4).

Model	Overall valid $\uparrow$	Unique $\uparrow$	Novel $\uparrow$	MetaSUN $\uparrow$
ADiT (joint QM9+MP20)	90.6	87.8	26.0	1.0
Crystal-GFN	51.7	51.7	51.7	0.0
Crystalformer	69.9	69.4	31.8	3.1
LLaMat2-CIF	84.4	81.4	30.0	2.1
LLaMat3-CIF	15.4	15.2	10.5	0.2
SymmCD	73.4	73.0	47.0	2.4

Appendix

Table 12: MP20 generation results (unrelaxed).

Model	Atoms connected	Bond angles	Bond lengths	Aromatic ring flat	Double bond flat	Internal energy	No steric clash
EQGAT-diff	84.4	86.9	87.0	87.0	87.0	86.8	82.9
SemlaFlow	92.3	94.8	94.6	94.9	94.2	94.8	92.0
ADiT	93.0	92.3	92.5	95.4	95.3	91.3	91.8
TABASCO	99.9	99.2	99.4	100.0	99.8	99.4	94.3
Zatom-1	99.6	99.6	99.7	100.0	100.0	99.3	96.4
GEOM-Drugs (reference)	–	–	–	–	–	–	–

Appendix

Table 13: Individual PoseBusters checks for GEOM-Drugs generation. Published baseline values are reported from Table 3 of Morehead et al. (2026) (arXiv v4). Results are percentages among RDKit-valid molecules.

Model	COV-R mean $\uparrow$	MAT-R mean $\downarrow$	COV-P mean $\uparrow$	MAT-P mean $\downarrow$
GO-Flow	94.82	0.7971	70.13	1.1068
Zatom-2	86.09	0.8459	83.07	0.9083

Appendix

Table 14: GEOM-Drugs conformer generation evaluation. Values are means across test molecules; COV is reported as a percentage and MAT as RMSD in Å. GO-Flow values are from Table 1 of Liu et al. (2026) .

Model	MR (%) $\uparrow$	RMSE $\downarrow$
Crystalite	66.09	0.0337
Zatom-2	64.00	0.1156

Appendix

Table 15: MP20 structure prediction evaluation. We report the Crystalite result from Hadži Veljković et al. (2026) .

Figure 9: Generated QM9 molecules. MolStar ball-and-stick views of nine Zatom-2 molecules. Carbon is gray, hydrogen white, nitrogen blue, and oxygen red. Bonds are perceived by MolStar from the saved structures (note that coordinates are unchanged).

Figure 10: Generated OMol25 molecules. MolStar ball-and-stick views spanning biomolecules, electrolytes, metal complexes, reactivity, and community datasets (labels identify the source subset). Notably, Zatom-2’s source category and chemical composition for each sample are not fixed during generation but rather are (re)produced. Note that Transition-1x and RGD1 supply metadata for known reaction path examples, but Zatom-2’s generated samples for these subsets are not verified transition states. Atoms are colored by element type, and bonds are perceived by MolStar.

Figure 11: Generated MP20 materials. MolStar views of unrelaxed Zatom-2 materials, with atoms colored by element. Each panel displays a $2\!\times\!2\!\times\!2$ periodic repetition of the saved cell; a gray outline marks one unit cell. Labels report the composition of a single saved cell (not a reduced formula) and its sample ID.

Figure 12: Generated OMat24 materials. MolStar views of nine force-conditioned Zatom-2 samples, with atoms colored by element. Each panel displays a $2\!\times\!2\!\times\!2$ periodic repetition of the generated cell, with one cell outlined in gray. Labels give the composition of the saved cell and an abbreviated sample ID. No relaxation, space group refinement, or symmetrization is applied.

Figure 13: Generated SCOPe-2k protein backbones. MolStar cartoons colored from the N-terminus (blue) to the C-terminus (red). Rows show helix-rich, mixed, and sheet-rich examples, respectively. Labels display backbone length, best-of-eight C $\alpha$ scRMSD in Å, and abbreviated sample IDs. The coordinates rendered for each sample are the generated backbones, not their ESMFold2 refolded structures.

Figure 14: Base Zatom-2 system embeddings across OMol25 and OMat24. Left: Zatom-2 (IV), without energy/force prediction; right: Zatom-2 (V), with it. Each point represents the same validation example in both columns, mean-pooled over all atoms. Rows show results for the joint datasets (a,b), OMol25 subsets (c,d), and OMat24 subsets (e,f), with matching colors across checkpoints. Each panel has an independent PCA fit; axis percentages report variance. Note that axes are not aligned between models and that “(sub)” denotes subsampled OMat24 data.

Figure 15: Base Zatom-2 atom embeddings across OMol25 and OMat24. Atom-level counterpart of Figure 14 (a,b), separately for O, P, N, and Br, selected by shared structure coverage. Left: Zatom-2 (IV), without MLIP supervision; right: Zatom-2 (V), with it. Each point is one atom, colored by domain, with identical atoms in both columns. Titles give sampled atom and structure counts. Each panel fits centered PCA independently; axis percentages report variance, and axes are not aligned between models.

Figure 16: Base Zatom-2 atom embeddings within OMol25. Atom-level counterpart of Figure 14 (c,d) for H, C, N, and O, the four elements present in the most sampled molecules. Left: without MLIP supervision; right: with it. Colors indicate the same OMol25 subsets as in Figure 14 . Each point is an individual atom, with an independent PCA fit per panel.

Figure 17: Base Zatom-2 atom embeddings within OMat24. Atom-level counterpart of Figure 14 (e,f) for Hg, La, Ag, and K, selected by structure coverage in the sampled validation materials. Left: without MLIP supervision; right: with it. Each point is an atom, colored by the corresponding OMat24 subset; “(sub)” denotes a subsampled subset. Identical atoms appear in both columns, with independent PCA fits. MLIP supervision concentrates more variance in PC1.

Explore similar work

Oct 6, 2026cs.LG

Atom-JEPA: Joint-Embedding Predictive Architecture for 3D Atomistic Systems

Large-scale self-supervised pretraining has reshaped modern machine learning, substantially advancing the ability of language and vision models to generalize across downstream tasks. While deep learning has driven considerable progress in modeling atomistic systems in recent years, self-supervised pretraining in this domain has not yet achieved comparable downstream generalization. To address this, we introduce Atom-JEPA, a self-supervised pretraining framework that learns latent representations from unlabeled 3D structures through complementary atom-level and substructure-level objectives inspired by joint-embedding predictive architectures. We pretrain Atom-JEPA on large-scale molecular and crystalline datasets and evaluate its transfer performance by fine-tuning on a diverse set of downstream property prediction tasks. Atom-JEPA achieves state-of-the-art performance on molecular ADMET and quantum-chemical property prediction tasks, and is highly competitive in predicting the physical properties of crystalline materials. These results demonstrate the potential of latent-space predictive pretraining to support broad downstream generalization from structural data alone. Code and pretrained model checkpoints are publicly available at https://github.com/khelverskovp/atom-jepa

Kasper Helverskov Petersen, Rasmus Hannibal Tirsgaard, François R J Cornet +2

Department of Applied Mathematics and Computer Science Technical University of Denmark

Joint-Embedding Predictive Architecture Materials Property Prediction 3D Representation Learning Molecular Property Prediction Molecular Representation Learning

Ask PDF

Jun 19, 2026cs.LG

Atomistic Language Models Understand and Generate Materials

Atomistic structure and natural language have long been modeled separately, with language models either calling atomistic models as tools or being fine-tuned on lossy textual encodings that discard atomistic information. We introduce Atomistic Language Models (ALMs) to pursue native multimodality, in which a single language backbone understands atomistic structures, generates materials from natural language, and optimizes crystal structures as instructed by text. By unifying a pretrained atomistic encoder, large language model, and denoising diffusion model through purely continuous projectors and staged training, ALMs achieve state-of-the-art results on crystal structure prediction and de novo generation. ALMs are enabled by a continuous bridge that maps language model embeddings directly into the steering space of atomistic diffusion, and are assisted by Text-to-Crystal Feynman-Kac (T2C-FK), a particle-based sampler that scores partial denoising trajectories to enforce stoichiometric targets at inference time. To evaluate the ability of ALMs to optimize and generate materials from natural-language prompts and 3D atom-coordinate inputs, we introduce ALM Bench, the first benchmark for text-conditioned crystal generation and optimization. Code, training data, and model weights will be released soon.

Sathya Edamadaka, Krithik Ramesh, Ju Li +1

Massachusetts Institute of Technology · Lyra Labs · Lila Sciences

Crystal Structure Prediction Language Modeling Crystal Structure Generation Conditional Diffusion Models

Ask PDF

Apr 2, 2026cs.LG

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