PaperlayerView exact revision on arXiv

Source-linked AI summary

Crystalite: A Lightweight Transformer for Efficient Crystal Modeling

Tin Hadži Veljković, Joshua Rosenthal, Ivor Lončarić, Jan-Willem van de Meent

arXiv:2604.02270v2cs.LGcs.AI

TL;DR

Crystal generation must search a vast space of possible structures, while common equivariant GNN approaches can be computationally demanding and recent diffusion Transformers often encode geometry weakly. Crystalite combines compact Subatomic Tokenization with GEM, which injects periodic pair geometry into Transformer attention. It reports state-of-the-art crystal structure prediction, the best S.U.N. score among evaluated baselines, and substantially faster sampling than geometry-heavy alternatives.

  • Problem

    Crystal discovery searches a combinatorially vast space, and common equivariant GNN approaches can be computationally demanding while some diffusion Transformers incorporate crystal geometry only weakly or indirectly.

    Source: §1 · p. 1

  • Method

    Crystalite is a lightweight diffusion Transformer that combines Subatomic Tokenization with GEM’s periodic geometric attention biases.

    Source: §1 · p. 2, §3.4 · p. 7

  • Results

    Crystalite achieves state-of-the-art crystal structure prediction and strong de novo generation results, the best S.U.N. score among evaluated baselines, and substantially faster sampling than geometry-heavy alternatives.

    Source: §1 · p. 2, §6 · p. 12

  • Takeaways & Limitations

    The results indicate that strong crystal modeling performance does not necessarily require full equivariance when periodic geometry and chemical structure are incorporated.

    Source: §6 · p. 12

Abstract

from arXiv · show

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.

1 Introduction

Crystalite targets a simpler crystal generator that retains useful geometric structure without the complexity and cost of equivariant message passing. It combines Subatomic Tokenization with GEM and reports strong prediction and generation performance at higher sampling speed than geometry-heavy baselines.

1 Introduction

  • The search space of crystalline materials is vast, and candidate evaluation typically relies on computationally expensive DFT.Only a small fraction of possible candidates is thermodynamically stable.Source: §1 · p. 1
  • Equivariant GNNs capture crystal geometry effectively but can be complex and computationally demanding, motivating simpler competitive models (Luo et al., 2025; Jiao et al., 2024a; Zeni et al., 2025; Miller et al., 2024; Yang et al., 2024).The central question is whether a lightweight Transformer can capture enough geometry without explicit equivariant message passing.Source: §1 · p. 1
  • Crystalite replaces one-hot atom types with Subatomic Tokenization and injects periodic geometric biases into standard multi-head attention.GEM computes periodic minimum-image pair geometry and converts it to additive attention biases; recent diffusion Transformers often incorporate geometry only weakly or indirectly (Joshi et al., 2025).Source: §1 · pp. 1–2, Figure 1 · p. 2
  • Crystalite achieves state-of-the-art crystal structure prediction and de novo generation performance while sampling much faster than geometry-heavy baselines.Its reported contributions also include the best S.U.N. discovery score among evaluated baselines.Source: §1 · p. 2

2 Related Work

Crystal-generation research spans geometry-rich equivariant architectures and lighter models with weaker geometric inductive bias. Crystalite is positioned within the diffusion-Transformer line, injecting periodic pair geometry directly into attention.

Equivariant and geometry-aware crystal generators

  • Equivariant and geometry-aware generators encode symmetry and periodicity directly, often with increasing architectural and computational complexity (Zeni et al., 2025; Gasteiger et al., 2021; Satorras et al., 2021; Miller et al., 2024; Hoellmer et al., 2025; Cornet et al., 2025; Luo et al., 2025).Examples include MatterGen, DiffCSP, FlowMM, and KLDM, which use distinct geometric or generative formulations.Source: §2 · p. 3

Lightweight alternatives to full equivariance

  • Lighter alternatives include UniMat’s diffusion 3D U-Net, which can remain competitive with equivariant baselines, and Transformer or hybrid crystal-generation approaches (Cao et al., 2025).These works leave open how much geometry a crystal generator should encode directly.Source: §2 · p. 3

Diffusion transformers for atomistic and crystal generation

  • Recent diffusion Transformers span minimal-bias molecule and material generation, crystal structure prediction, and crystalline generation (Joshi et al., 2025).Crystalite differs by injecting periodic pairwise geometry directly into attention rather than relying only on augmentation or latent-space structure.Source: §2 · p. 3

3 Methodology

Crystalite keeps a standard diffusion Transformer backbone while adapting its representation, attention, and sampling procedure to crystal structure. It jointly models atom types, fractional coordinates, and lattice variables, adding periodic geometry through GEM.

3.1 Subatomic Tokenization

  • Subatomic Tokenization replaces high-dimensional one-hot atom types with low-dimensional continuous descriptors, giving diffusion a chemically meaningful geometry and reducing capacity to memorize frequent compositions.The descriptors encode element properties and may be compressed with PCA; the supplied Figure 2 passage describes compression to 16 dimensions.Source: §3.1 · pp. 4–5, Figure 2 · p. 4

3.2 Diffusion formulation for crystals

  • Crystalite diffuses atom-type, coordinate, and lattice channels jointly, using a lower-triangular lattice latent and wrapped fractional-coordinate residuals to respect coordinate periodicity.The lattice representation remains basis-dependent, so structures are preprocessed into Niggli-reduced cells and a fixed lattice-parameter convention.Source: §3.2 · pp. 5–6

3.3 Crystalite architecture

  • One token per atom combines atom-type and coordinate embeddings, while a separate global lattice token joins the sequence processed by an AdaLN-conditioned Transformer.Output heads predict denoised atom tokens, coordinates, and the lattice latent.Source: §3.3 · pp. 6–7, Figure 3 · p. 7

3.4 Geometry Enhancement Module (GEM)

  • GEM recomputes periodic minimum-image geometry and adds a geometry-dependent bias to attention, injecting pair structure without equivariant message passing.Its distance selection minimizes the Cartesian quadratic form induced by the lattice metric; the bias acts only on atom–atom interactions.Source: §3.4 · pp. 7–8
  • Figure 4 contrasts two atom-loss settings: the higher-stability setting loses UN faster, while the more diversity-preserving setting has a flatter, more sustained SUN trajectory.The caption frames this as a stability–novelty/uniqueness trade-off that makes checkpoint selection and loss balancing important.Source: Figure 4 · p. 8

3.5 Channel-wise anti-annealing during sampling.

  • Channel-wise anti-annealing rescales the reverse-time update separately for atom-token, coordinate, and lattice channels, intensifying refinement without changing the learned denoiser.The channel-specific factor comes from an auxiliary Karras schedule; a factor of 1 recovers the standard EDM sampler.Source: §3.5 · p. 8

4 Experimental Setup

The experiments evaluate Crystalite on de novo generation and crystal structure prediction using three Materials Project- or Alexandria-derived datasets. The setup uses task-specific model configurations and benchmarks sampling throughput under matched hardware and batch-size conditions.

4.1 Datasets

  • The evaluation uses MP-20, chronologically split MPTS-52, and Alex-MP-20, covering structures with up to 20 or 50 atoms per unit cell.MP-20 contains 45,231 materials and 89 atom types; MPTS-52 contains 40,476 structures and has a temporal distribution shift; Alex-MP-20 contains 675,204 structures.Source: §4.1 · pp. 8–9

4.2 Task setup

  • De novo generation jointly generates atom types, fractional coordinates, and lattice parameters from noise; crystal structure prediction fixes composition and predicts geometry.For CSP, atom-type features are fixed to the known composition and the type loss is masked during training and sampling.Source: §4.2 · p. 9
  • DNG uses a 14-layer, width-512 Transformer with PCA-compressed Subatomic Tokenization, while CSP uses a same-depth, width-1024 model with fixed atomic-number features.Sampling speed is timed for 1,000 crystals on one H100 GPU at each model's largest feasible batch size; Crystalite also has a separate optimized-bfloat16 timing.Source: §4.2 · p. 9

5 Results and Discussion

Across CSP benchmarks, Crystalite outperforms prior methods on Match Rate and RMSE, with especially pronounced RMSE gains. In de novo generation, it leads the compared methods on SUN rate and sampling speed and preserves diversity as sampling scales up.

5.1 CSP Results

  • Crystalite outperforms prior methods across all CSP benchmarks on Match Rate and RMSE.The reported improvement is especially pronounced in RMSE.Source: §5.1 · p. 9, Table 1 · p. 10
  • 34.8%, 44.3%, and 74.7% RMSE reductions on MP-20, MPTS-52, and Alex-MP-20 yield an approximately 51% average reduction over prior bests.Source: §5.1 · p. 9
  • About 20% RMSE improvement from GEM, with little effect on Match Rate, indicates its gains chiefly refine geometric accuracy.Source: §5.1 · p. 9

5.2 DNG Results

  • Crystalite achieves the highest SUN rate and fastest sampling speed among the compared de novo generation methods.Table 2 metrics use 10,000 generated crystals per model.Source: §5.2 · p. 10, Table 2 · p. 10
  • Downweighting atom-type loss yields a flatter, more sustained SUN trajectory, while GEM mainly improves the stability side of the trade-off.The atom-type loss is downweighted because compositional memorization is difficult to control independently of structural quality when atom types, coordinates, and lattice parameters are modeled jointly.Source: §5.2 · p. 10
  • Crystalite preserves diversity better than ADiT as sampling scales up, while uniqueness and UN rates should be reported with total sample count.Source: Figure 5 · p. 11, §5.2 · p. 11

6 Conclusion

Crystalite combines Subatomic Tokenization and GEM in a lightweight Transformer for crystal prediction and generation. Its reported results include state-of-the-art crystal structure prediction, the best SUN score among evaluated baselines, and substantially faster sampling than geometry-heavy alternatives.

6 Conclusion

  • Crystalite achieves state-of-the-art crystal structure prediction and the best SUN score among evaluated baselines while sampling substantially faster than geometry-heavy alternatives.The conclusion says these results show strong crystal modeling performance does not necessarily require full equivariance when periodic geometry and chemical structure are incorporated.Source: §6 · p. 12

A.1 Unit-cell representation of crystals

  • A crystal is represented by a unit cell and its lattice, with fractional coordinates specifying basis-atom positions and periodic repetition defining the full material.The finite unit-cell description implicitly defines the infinite periodic crystal.Source: §A.1 · pp. 16–17

A.2 Symmetries and representation non-uniqueness

  • Equivalent crystal representations motivate fractional coordinates, wrapped periodic residuals, metric-aware minimum-image geometry in GEM, and random global translations during training.For orthogonal cells, componentwise wrapping and metric-aware minimum-image geometry coincide; for general non-orthogonal cells, they need not.Source: §A.2 · p. 18, §A.1 · p. 18

B.1 Subatomic Tokenization

Subatomic Tokenization replaces one-hot atom identities with compact continuous tokens built from chemical descriptors. Crystalite embeds atom and lattice information in a standard Transformer and adds periodic pair geometry to its attention logits through GEM.

B.1 Subatomic Tokenization

  • Subatomic Tokenization builds element descriptors from period, group, block, and valence-shell occupancies, then standardizes, rebalances, and normalizes them into continuous tokens.The raw descriptor has 34 dimensions; PCA can compress the token to dimension d, and nearest-prototype matching deterministically decodes predicted tokens to elements.Source: §B.1 · pp. 19–21

C.1 Tokenization and input embeddings

  • Crystalite represents each atomic site with a token and the unit-cell geometry with a global token, processing them through a standard Transformer trunk conditioned on diffusion noise.Output heads predict denoised atom-type and coordinate quantities from site tokens and lattice parameters from the lattice token.Source: §C.1 · pp. 22–23

C.2 Geometry Enhancement Module (GEM)

  • GEM computes periodic minimum-image pair geometry and injects distance and edge-aware biases additively into attention logits, without changing tokenization or prediction heads.When GEM is disabled, its bias is zero and the model retains the same tokenization, diffusion objective, and output heads.Source: §C.2 · pp. 24–25, Figure 10 · p. 24

C.3 Task-specific model configurations

The paper uses a shared Crystalite model family with task-specific configurations for its main de novo generation and crystal structure prediction experiments.

C.3 Task-specific model configurations

  • Table 4 summarizes task-specific configurations for the main DNG and CSP experiments; Alex-MP-20 CSP uses the MP-20 configuration, and CSP has zero effective atom-type loss weight because composition is fixed.The table marks whether PCA-compressed Subatomic Tokenization is used or atomic-number features are used instead.Source: Table 4 · p. 26, §C.3 · p. 26

D EDM Training Details

Crystalite applies EDM to a continuous crystal state containing atom-type features, fractional coordinates, and a lattice latent. The appendix also describes channel-wise anti-annealing during sampling and the evaluation procedures for DNG and CSP.

D EDM Training Details

  • EDM noises atom-token and lattice channels in Euclidean space, while coordinates use a centered representation that is wrapped back into the unit cube before embedding.The coordinate loss uses a wrapped fractional residual, distinct from GEM's metric-aware minimum-image displacement.Source: D EDM Training Details · p. 27

D.1 Channel-wise anti-annealing during sampling

  • Channel-wise anti-annealing rescales selected EDM drifts with a channel-dependent time warp while leaving the denoiser and underlying EDM schedule unchanged.Channels with α(q) greater than 1 are driven more aggressively toward their denoised predictions.Source: §D.1 · p. 29

E.1 De novo generation (DNG)

  • DNG evaluation reports validity, uniqueness and novelty, distribution matching, and thermodynamic metrics, with relaxation and thermodynamic-processing failures counted against the corresponding rates.The unique-and-novel set is formed by restricting to novel structures and then deduplicating that subset.Source: §E.1 · pp. 30–32

E.2 Crystal structure prediction (CSP)

  • CSP evaluates one composition-conditioned sample per test composition, counting a prediction correct when StructureMatcher finds a valid match to its ground-truth structure.For matched pairs, the evaluation also reports the matcher RMS displacement after alignment.Source: §E.2 · pp. 32–33

E.3 Sample-size intensive and extensive metrics

The section distinguishes metrics whose targets remain meaningful on subsets from discovery-yield metrics that depend on the full generation budget. It recommends reporting extensive metrics at matched budgets while estimating intensive metrics from random subsets when needed.

E.3 Sample-size intensive and extensive metrics

  • Sample-intensive metrics can be estimated from random subsets without changing their target, including validity, stability, per-sample properties, and distribution metrics.These describe typical generated crystals rather than the discovery yield of the whole set.Source: §E.3 · p. 33
  • Sample-extensive metrics depend on generation budget because duplicate collisions grow with sampling, so subset evaluation measures a different budget rather than estimating the full-budget value.The section includes uniqueness, distinct discoveries, discovery-yield novelty, UN, and SUN in this category.Source: §E.3 · p. 33
  • Uniqueness, UN, and SUN should be reported with the number of generated samples and compared only at matched sample budgets.Small subsets can make uniqueness and related discovery metrics look artificially better.Source: §E.3 · p. 33
  • Novelty is intensive when tested per sample against a fixed reference set, but is treated as extensive when used in the deduplicated discovery pipeline.Its classification depends on how novelty is defined and reported.Source: §E.3 · p. 34
  • SUN combines full-batch UN with a subset-based estimate of the stable fraction within the UN structures.The practical recommendation is to report uniqueness, UN, and SUN with budget n, while estimating validity, stability, and distribution metrics from random subsets when needed.Source: §E.3 · pp. 33–34

F.1 Ablation of Subatomic Tokenization

With stability trajectories aligned, Subatomic Tokenization maintains a stronger stability-diversity balance than one-hot encoding. The section also reports evaluation under the MatterGen pipeline, but the supplied passage gives no comparative outcomes.

F.1 Ablation of Subatomic Tokenization

  • Subatomic Tokenization sustains higher SUN and UN rates in later training than one-hot encoding after stability curves are aligned.The comparison uses an identical Crystalite model and balances the one-hot atom-type loss weight to align stability trajectories.Source: §F.1 · p. 35, Figure 11 · p. 35

F.2 DNG MatterGen evaluation pipeline results

  • The MatterGen evaluation uses its evaluation pipeline to compare models on a setup not designed by the authors.Table 5 covers validity, uniqueness, novelty, stability, and relaxation metrics for MP-20; no values are supplied here.Source: §F.2 · p. 35, Table 5 · p. 35

F.3 GEM effect on DNG Results

GEM is associated with higher stability and SUN during de novo generation, and clearer RMSE gains than Match Rate gains in structure prediction. The anti-annealing ablations report robust performance, with task-specific trade-offs and preferred settings.

F.3 GEM effect on DNG Results

  • With GEM, de novo generation achieves higher stability throughout training and consistently higher SUN, despite the shared decline in UN as training progresses.The authors report that this is achieved without a disproportionate loss in generative diversity.Source: §F.3 · p. 36, Figure 12 · p. 36

F.4 GEM effect on CSP Results

  • In CSP, GEM improves RMSE more clearly and consistently than Match Rate, suggesting a stronger effect on refinement than structural-mode recovery.The authors interpret the geometric biases as improving local atomic placement and structural fidelity during denoising.Source: §F.4 · p. 36, Figure 13 · p. 36

F.5 DNG Sensitivity to anti-annealing

  • DNG performance is fairly insensitive to anti-annealing, but the highest-SUN configuration has worse Wasserstein distances, so a more balanced setting is reported.This preserves strong discovery performance with better agreement with the reference distribution.Source: §F.5 · p. 37

F.6 CSP Sensitivity to anti-annealing

  • For CSP, coordinate and lattice anti-annealing at 4 each yields both the highest Match Rate and lowest RMSE, while performance remains broadly stable across settings.The authors report a measurable boost without inducing instability.Source: §F.6 · p. 38, Table 7 · p. 38

G Crystalite S.U.N. Crystals

  • Figure 14 presents stable, unique, and novel crystals generated by Crystalite trained on MP-20.Source: Figure 14 · p. 39
Loading 2604.02270v2…