Masked Generative Crystal Transformer

• The paper introduces MaskGXT, a masked discrete generative transformer that predicts full crystal structures from chemical compositions, achieving 79.06% METRe accuracy on the polymorph split.
• It leverages symmetry-aware tokenization by discretizing lattice, coordinates, space group, and Wyckoff positions, and refines predictions using a novel sub-bin offset head.
• The model’s transformer architecture with confidence-ranked greedy decoding and space-group–stratified sampling enhances polymorph coverage compared to continuous diffusion and flow approaches.

Searching arXiv for MaskGXT and closely related crystal-generation papers to ground the article in current literature. Masked Generative Crystal Transformer (MaskGXT) is a discrete masked generative model for crystal structure prediction (CSP) from chemical composition. In the formulation introduced in "Discovering Crystal Structure Prediction Algorithms with an AI Co-Scientist" (Seong et al., 22 Jun 2026), the input is the set of atom types $A=(a_1,\dots,a_N)$, and the output is a crystal structure comprising the lattice $L$, the fractional atomic coordinates $F\in[0,1)^{3\times N}$, and crystallographic symmetry information in the form of a space group $g$ and Wyckoff positions $W_g$. MaskGXT was identified through the Human–AI Co-discovery system (HACO), which transferred the MaskGIT paradigm from vision to CSP, then refined it with crystallographic symmetry tokens, space-group–stratified sampling, and sub-bin coordinate refinement. On the MP-20 polymorph split, it reaches $79.06\%$ match-everyone-to-reference (METRe) accuracy, compared with $70.87\%$ for the strongest evaluated baseline (Seong et al., 22 Jun 2026).

1. Problem setting, scope, and naming

Crystal structure prediction asks for a stable crystal structure given only the chemical composition. In the MaskGXT setting, the model must predict not merely compositionally valid outputs, but a full periodic solid with lattice parameters, fractional coordinates on a torus, and symmetry labels. The task is difficult for three reasons emphasized in the original work: the domain is geometrically constrained by periodicity, symmetry-redundant because many descriptions are equivalent under origin shifts and permutations, and polymorphic because a single composition may admit multiple valid structures (Seong et al., 22 Jun 2026).

HACO selected masked discrete modeling as the operative framework because it offers parallel masked decoding in a few steps, categorical posteriors for confidence-ranked greedy selection, and a clear, finite decision space. This directly contrasts with diffusion and flow formulations that evolve continuous geometries by stochastic integration. In MaskGXT, this choice is not merely an architectural preference; it determines the representation, objective, and decoding strategy.

A common naming confusion arises from the near collision with "MaskGCT: Zero-Shot Text-to-Speech with Masked Generative Codec Transformer" (Wang et al., 2024). That model belongs to zero-shot text-to-speech rather than crystalline materials. MaskGXT, by contrast, is a CSP model over lattice, coordinates, space group, and Wyckoff structure.

2. Crystal representation and symmetry-aware tokenization

MaskGXT tokenizes a crystal $(L,F,A,g,W_g)$ into a sequence of $1+6+5N$ tokens. There is one space-group token $g$, six lattice tokens for $L$0, and five tokens per atomic site: three coordinate tokens $L$1, one Wyckoff letter token $L$2, and one atom-type token $L$3. The atom-type tokens are conditioned side information rather than prediction targets (Seong et al., 22 Jun 2026).

The geometric variables are discretized into $L$4 bins. For a fractional coordinate $L$5, quantization is defined by

$L$6

The modulo operation makes the coordinate axis circular, so bins $L$7 and $L$8 remain adjacent. Space groups use the 230 international space groups as vocabulary. Wyckoff letters use a data-derived per-space-group vocabulary.

A central design feature is the joint inclusion of symmetry tokens with geometry tokens. Space-group and Wyckoff tokens participate directly in masked prediction. Symmetry consistency is encouraged through two training augmentations. First, Euclidean-normalizer origin-shift augmentation translates all fractional coordinates modulo $L$9 and relabels Wyckoff letters using precomputed Euclidean-normalizer translation cosets. Second, intra/inter-orbit permutation augmentation permutes atoms only within symmetry-equivalent descriptions, rather than across arbitrary $F\in[0,1)^{3\times N}$0 permutations. This is a learned consistency mechanism rather than a hard constraint.

To recover continuous precision below the discretization width, MaskGXT adds a bounded sub-bin offset head. If $F\in[0,1)^{3\times N}$1 is the bin and $F\in[0,1)^{3\times N}$2 is the learned offset, reconstruction is

$F\in[0,1)^{3\times N}$3

The offset head uses a bounded activation $F\in[0,1)^{3\times N}$4 so that the reconstructed value remains inside the bin.

3. Transformer architecture and masked learning objective

MaskGXT uses a bidirectional 34-layer Pre-LN Transformer trunk with QK-normalized multi-head self-attention and SwiGLU feed-forward blocks. The hidden dimension is $F\in[0,1)^{3\times N}$5, the number of attention heads is $F\in[0,1)^{3\times N}$6, and the parameter count is approximately $F\in[0,1)^{3\times N}$7M. Learned positional encodings are factored into an element-block encoding and a within-element encoding, reflecting the canonical site ordering used during tokenization (Seong et al., 22 Jun 2026).

Before the trunk, the five tokens associated with a site are combined through per-site linear fusion into one atom-wise feature. After the trunk, MaskGXT uses linear categorical heads over the token streams $F\in[0,1)^{3\times N}$8 together with regression heads for lattice and coordinate sub-bin offsets. The model therefore predicts both coarse discrete structure and fine continuous correction.

Training uses masked token modeling. A token sequence is corrupted by independently masking each token with probability $F\in[0,1)^{3\times N}$9, and the bidirectional transformer reconstructs all masked positions in parallel. Coordinates are supervised with ordinal circular label smoothing,

$g$0

with $g$1. Lattice tokens use non-circular ordinal smoothing. The main objective is a stream-weighted masked cross-entropy,

$g$2

with stream weights $g$3.

The sub-bin refinement head is trained with a smooth-L1 penalty on the residual from the bin center,

$g$4

with loss weights $g$5 for lattice and coordinate offsets. The paper’s ablations show that this head sharply lowers RMSE and cRMSE without materially changing match rate, indicating that it improves precision rather than coarse structural correctness.

4. Sampling, polymorph coverage, and validation protocol

Inference begins from a fully masked token sequence conditioned on the atom-type side information. From that sequence, MaskGXT reads a space-group posterior in one forward pass,

$g$6

It then selects the top-$g$7 space groups and branches generation by fixing each chosen $g$8 while decoding the remaining tokens. If no space group has posterior mass at least $g$9, the model does not clamp the space-group token and instead samples without stratification (Seong et al., 22 Jun 2026).

Within each branch, decoding uses confidence-ranked greedy masked parallel unmasking for $W_g$0 iterations. At each iteration, the model predicts posteriors for all masked tokens in parallel, fills masked positions with $W_g$1 tokens, ranks positions by confidence, fixes the highest-confidence positions, and re-masks the remainder. After decoding, the model reconstructs continuous lattice and coordinate values through dequantization combined with the learned sub-bin offsets.

The explicit purpose of the space-group branching step is polymorph coverage. The paper argues that spreading generation across distinct space groups improves coverage more than i.i.d. sampling. This is particularly important for CSP benchmarks in which the same composition may correspond to several valid reference structures.

Evaluation uses StructureMatcher with tolerances $W_g$2, $W_g$3, and $W_g$4. Two regimes are distinguished. In one-to-one CSP, a single generated structure is compared to each reference and reported with match rate (MR) and RMSE over matched pairs, both in unfiltered form and validity-filtered form. In polymorph-aware evaluation, METRe counts a reference as matched if any generated structure with the same composition matches it:

$W_g$5

The development loop deliberately relies on cheap, fast, and well-aligned validation: no DFT relaxation or energy screening is used, and invalid generated structures are counted as unmatched.

5. Benchmarks and quantitative performance

MaskGXT is evaluated on three benchmark settings: MP-20, MPTS-52, and the MP-20 polymorph split. MP-20 contains Materials Project crystals with at most $W_g$6 atoms per cell and formation energies within $W_g$7 eV/atom of the hull; MPTS-52 extends to at most $W_g$8 atoms per cell with a chronological split; and the MP-20 polymorph split assigns all polymorphs of a composition to the same subset, preserving multi-reference compositions for polymorph-aware evaluation (Seong et al., 22 Jun 2026).

In one-to-one CSP, MaskGXT achieves the best match rate and lowest RMSE among the compared methods on both MP-20 and MPTS-52. In polymorph-aware evaluation, it attains the best reported match rate on standard MP-20 and on the polymorph split. Under a uniform sampling budget $W_g$9, it also leads the baselines.

Benchmark	Setting	MaskGXT result
MP-20	One-to-one MR / RMSE	$79.06\%$0 / $79.06\%$1 unfiltered; $79.06\%$2 / $79.06\%$3 validity-filtered
MPTS-52	One-to-one MR / RMSE	$79.06\%$4 / $79.06\%$5 unfiltered; $79.06\%$6 / $79.06\%$7 validity-filtered
MP-20	METRe / cRMSE	$79.06\%$8 / $79.06\%$9
MP-20 polymorph split	METRe / cRMSE	$70.87\%$0 / $70.87\%$1
MP-20, budget $70.87\%$2	METRe / cRMSE	$70.87\%$3 / $70.87\%$4
Polymorph split, budget $70.87\%$5	METRe / cRMSE	$70.87\%$6 / $70.87\%$7

The strongest baseline on the polymorph split is Crystalite at $70.87\%$8 METRe with cRMSE $70.87\%$9, so MaskGXT improves METRe by $(L,F,A,g,W_g)$0 percentage points. Decoder ablations identify two especially consequential components. Removing space-group stratification lowers polymorph-split METRe from $(L,F,A,g,W_g)$1 to $(L,F,A,g,W_g)$2. Removing the offset head raises MP-20 matched-pair RMSE from $(L,F,A,g,W_g)$3 to $(L,F,A,g,W_g)$4 unfiltered while leaving match rate essentially unchanged. These results separate the model’s two main gains: stratification improves polymorph coverage, while offset regression improves geometric precision.

6. Relation to earlier crystal generative models, limitations, and significance

Within CSP, MaskGXT differs from diffusion and flow baselines such as DiffCSP, FlowMM, and OMatG, and from autoregressive or multimodal baselines such as Crystalite and MCFlow, by representing geometry and symmetry as a single discrete token sequence decoded by a masked bidirectional transformer (Seong et al., 22 Jun 2026). Most continuous models evolve lattice and fractional coordinates directly in continuous space and often require lengthy stochastic integration at sampling time. MaskGXT instead exploits discrete tokens, confidence-ranked greedy decoding, and explicit space-group stratification.

An earlier but distinct masked generative crystal model appears in "Crystal Transformer: Self-learning neural LLM for Generative and Tinkering Design of Materials" (Wei et al., 2022). That BLMM Crystal Transformer is a masked generative transformer for inorganic compositions rather than full crystal structures. It uses a transformer encoder plus three policy heads for blank selection, element selection, and insertion options, enabling de novo composition generation and tinkering or doping through probabilistic blank filling. BLMM therefore belongs to the broader family of masked generative transformers for crystalline materials, but it does not solve CSP over lattices, coordinates, space groups, and Wyckoff positions.

MaskGXT’s limitations are explicit. Symmetry consistency is learned rather than hard enforced, because lattice, coordinates, space group, and Wyckoff tokens are predicted independently and only coupled through augmentation and joint training. HACO’s empirical selection is compute-budget sensitive, so methods that converge more slowly may be underexplored in the discovery loop. The approach also benefits from a domain with fast, aligned validation; the paper suggests that domains with long training cycles, weak proxy objectives, or physical experiments may require stronger human steering, surrogate models, or structured evaluation. A plausible implication is that MaskGXT is as much a demonstration of a discovery methodology as of a single CSP architecture: the reported gains derive from the interaction between cross-domain transfer, sparse high-level human objectives, and crystallography-specific refinements.

In sum, MaskGXT defines a masked discrete formulation of crystal structure prediction in which symmetry tokens, quantized geometry, bounded sub-bin refinement, and space-group–stratified decoding are integrated into a single bidirectional transformer. Its results on MP-20, MPTS-52, and especially the MP-20 polymorph split establish masked generative crystal modeling as a competitive alternative to continuous diffusion and flow approaches for benchmark CSP (Seong et al., 22 Jun 2026).