Crystalite: A Lightweight Transformer for Efficient Crystal Modeling

Abstract: 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.

• The paper presents a novel diffusion Transformer that incorporates continuous atom tokenization and periodic geometric attention biases to effectively model crystal structures.
• The paper demonstrates state-of-the-art results in predicting crystal structures and generating de novo crystals while reducing computational overhead compared to equivariant GNNs.
• The paper introduces channel-wise anti-annealing and calibrated loss balancing to navigate the trade-off between stability and diversity, ensuring fast and scalable generation.

Crystalite: Lightweight Diffusion Transformer for Efficient Crystal Modeling

Introduction

Crystalite proposes a diffusion Transformer architecture for generative modeling of crystalline materials, targeting both crystal structure prediction (CSP) and de novo crystal generation (DNG). Traditionally, state-of-the-art performance in crystal generation has relied on equivariant GNNs that encode geometric and symmetry constraints, but these models are computationally intensive and exhibit slow sampling rates. The fundamental question addressed is whether a standard Transformer, when augmented with appropriate inductive biases, can achieve competitive geometric fidelity for crystals at lower computational overhead.

Architecture Overview

Crystalite operates on a continuous tuple $(\mathbf{H},\mathbf{F},\mathbf{y})$ associated with each crystal, representing atom types (via chemically structured tokens), fractional coordinates, and a parameterized lattice. The backbone is a Transformer composed of a stack of AdaLN-conditioned attention blocks. Each atom is represented as a token summing embeddings of its chemical identity and spatial location, with a separate token for the global unit cell lattice.

Figure 1: Crystalite's architecture—atoms and lattice embedded as tokens, processed through AdaLN-conditioned Transformer, outputting denoised crystal states.

Key innovations are:

• Subatomic Tokenization: Instead of high-dimensional one-hot vectors for atom identity, chemical elements are embedded as continuous 34-dimensional descriptors encoding period, group, block, and valence shell occupancies, compressed to 16 dimensions with PCA. This delivers a chemically meaningful embedding space, regularizes atom-type encoding, and aligns the diffusion process with continuous variables.

Figure 2: Subatomic tokenization—element descriptors compressed to 16D tokens for continuous diffusion, examples shown for Oxygen and Titanium.

• Geometry Enhancement Module (GEM): GEM computes pairwise minimum-image geometry under periodic boundary conditions using fractional coordinates and lattice parameters. Two attention biases (distance penalty and learned edge features) are constructed and injected additively into the Transformer attention logits, directly encoding crystal geometry without resorting to equivariant message passing.

  Figure 3: GEM's architecture overview—periodic geometry is mapped to additive attention biases for multi-head attention.

• Channel-wise Anti-Annealing for Sampling: EDM reverse-time updates are rescaled per channel (atom-type, coordinates, lattice), dynamically warping denoising trajectories to improve geometric refinement.

Subatomic Tokenization and Chemical Embedding

Classic approaches use one-hot encoding for atomic species, which is suboptimal due to its high dimensionality and lack of chemical structure. Crystalite's tokenization produces a continuous, low-dimensional embedding reflecting the chemical stratification of the periodic table. Atom tokens are built from fundamental features and compressed via PCA, producing a latent space where chemically similar atoms are close—increasing model generalization and reducing compositional memorization.

Figure 4: PCA projection of atom tokens preserves chemical organization after dimensionality reduction.

Neighbor relationships in token space align with chemical similarity, e.g., Fe is close to its chemically related elements.

Figure 5: Fe's local neighborhood in token PCA space—chemically related elements cluster together.

Periodic Geometry Module and Attention Biasing

GEM leverages minimum-image conventions under the lattice metric to compute periodic atom pair geometry. The resulting biases (distance and learned edge features) are head-specific, noise-dependent, and modulate attention scores for atom-atom pairs while ignoring lattice-token interactions. This introduces periodic geometric structure directly into the attention mechanism, trading off explicit equivariance for modular efficiency.

Figure 6: Detailed GEM workflow—minimum-image geometry computed, edge features extracted, biases injected into attention logits.

Diffusion Formulation and Training Dynamics

Crystalite utilizes joint EDM diffusion over atom types, coordinates, and lattice parameters, treating all channels as continuous variables. During training, noise is injected channel-wise and denoising losses are computed with appropriate wrapped residuals (for coordinates) and Euclidean metrics (for atom and lattice). Checkpoint selection and loss balancing are critical for maintaining diversity and stability; strong atom-type losses induce compositional memorization and reduce uniqueness.

Figure 7: DNG trade-off—novelty decreases as stability improves; loss balancing controls diversity/stability trajectory.

Experimental Results

Crystal Structure Prediction (CSP)

Crystalite achieves state-of-the-art CSP results on MP-20, MPTS-52, and Alex-MP-20 benchmarks, outperforming prior GNN-based and flow-matching models across match rate and RMSE criteria. In ablation, GEM yields significant reductions in RMSE (~20%) but limited impact on match rate, indicating that geometric attention biases primarily enhance local atomic fidelity.

Figure 8: GEM's effect—RMSE improvement in CSP, match rate unchanged.

De Novo Crystal Generation (DNG)

Generative performance is evaluated via validity, uniqueness, novelty, and stability metrics. Crystalite attains the highest stable unique-and-novel (SUN) discovery rate among evaluated baselines, with strong compositional and structural validity. Notably, it achieves markedly faster sampling speeds than geometry-heavy models, making it suitable for large-scale screening.

Figure 9: Large-scale generation—Crystalite preserves diversity at $10^6$ samples, outperforming ADiT in UN rate.

Training induces a stability-diversity trade-off, mitigated via downweighting atom-type losses. GEM increases stability and SUN rates consistently in DNG.

Figure 10: GEM effect on DNG—higher stability and SUN rates maintained throughout training.

Efficiency and Practicality

Crystalite is computationally efficient: a single $6.7\times 10^7$-parameter model (14 layers, 16 heads) samples crystals substantially faster than equivariant GNN-based baselines. Optimized inference further reduces generation latency, expanding practical accessibility.

Theoretical and Practical Implications

Crystalite demonstrates that full geometric equivariance is not strictly necessary for high-fidelity crystal modeling. Injecting chemically informed atom tokenization and periodic geometric attention bias into a standard Transformer backbone achieves strong geometric and chemical structure recovery, while markedly reducing architectural and computational complexity. This raises critical implications for scalable materials discovery: lightweight Transformer architectures, properly biased, can enable fast, broad exploration of crystal phase space.

Future Directions

• Scalable Materials Generation: The modularity and throughput of Crystalite recommend it for extensive screening and inverse design workflows.
• Hybrid Architectures: Integrating lightweight geometry-aware Transformers with LLM-based generative priors may further enhance flexibility and conditional capability.
• Property-Guided Generation: Extension to property-conditional generative tasks is straightforward via additional conditional embeddings.
• Transfer and Multi-domain Modeling: Given the success of joint embedding spaces, combining molecular and materials modeling in a unified architecture is promising, as shown in works like Zatom-1 (Morehead et al., 24 Feb 2026).
• Fine-grained Geometry Inductive Bias: Further improvements may be realized by incorporating finer symmetry constraints or explicit space group conditioning.

Conclusion

Crystalite establishes that appropriately biased diffusion Transformers can achieve state-of-the-art crystal structure generation and prediction, with strong chemical and geometric fidelity, efficient sampling, and modular simplicity. By eschewing full equivariant message passing in favor of structured tokenization and periodic geometric attention modulation, it achieves both speed and diversity—supporting scalable, practical applications in computational materials discovery.