Discovering Crystal Structure Prediction Algorithms with an AI Co-Scientist

Abstract: We introduce Human-AI Co-discovery system (HACO) for scientific algorithm discovery through cross-domain search and sparse human steering. Starting from the goal of generating crystal structures from chemical compositions, HACO searched across generative modeling methodologies from multiple fields and identified MaskGIT, a masked generative model from vision, as a promising framework for crystal structure prediction (CSP). HACO instantiated this masked formulation as a discrete token model of crystal structure; guided by sparse high-level human objectives, it then added crystallographic symmetry tokens, space group stratified sampling for polymorph coverage, and sub-bin coordinate refinement, yielding the Masked Generative Crystal Transformer (MaskGXT). On the MP-20 polymorph split, MaskGXT reaches 79.06% match-everyone-to-reference (METRe) accuracy, compared with 70.87% for the strongest evaluated baseline. MaskGXT also attains the best match rate on standard MP-20 and MPTS-52 CSP benchmarks. These results provide evidence that, in domains offering cheap, fast, and well-aligned validation, transfer-guided interactive AI co-scientists can contribute to scientific algorithm discovery by identifying transferable modeling principles and combining them with targeted human domain guidance.

• The paper introduces the HACO system that autonomously discovers CSP algorithms using a tree-structured search, achieving high METRe accuracy on standard benchmarks.
• The paper demonstrates that cross-domain methodology transfer, via adapting MaskGIT to crystal systems, leads to superior polymorph prediction and reduced RMSE values.
• The paper highlights that agentic search paired with high-level human steering enables scalable discovery of innovative algorithms for crystal structure prediction.

MaskGXT: AI-Driven Algorithm Discovery for Crystal Structure Prediction

System Design: HACO and Cross-Domain Scientific Algorithm Search

The Human–AI Co-discovery (HACO) system operationalizes scientific algorithm discovery via agentic, tree-structured empirical search and sparse high-level human intervention. In this paradigm, the agent autonomously proposes, implements, and empirically validates generative modeling frameworks for crystal structure prediction (CSP), targeting match-everyone-to-reference (METRe) accuracy on MP-20 polymorph splits. HACO’s orchestrator (implemented with Claude Opus 4.7) manages node expansions, operator selection, parallel GPU training, METRe-based evaluation, and human steering integration, facilitating three staged search phases: methodology discovery, scale-up optimization, and sampling refinement. The critical innovation lies in the cross-domain “idea” operator, which surveys generative modeling literature beyond material science, enabling transfer of architectures from domains such as vision and language.

Figure 1: Empirical research trajectory visualizing METRe improvements across agent-proposed candidates and escalating compute budgets.

Human steering is strictly high-level—supplying domain mechanisms (e.g., symmetry, sampling diversity) or objectives (e.g., sub-bin precision)—leaving technical adaptation and implementation to the agent. This ensures the agent’s autonomy in empirical search and code synthesis while allowing domain-guided algorithmic refinement. The search yields MaskGIT, a masked generative model from vision, as the foundational architecture for the CSP target.

Figure 2: HACO agentic search tree for CSP, with colored branches tracing methodology transfer and refinement, including sparse human interventions.

MaskGXT: Masked Generative Crystal Transformer

MaskGXT represents crystals as discrete token sequences—lattice parameters, fractional coordinates, space group, Wyckoff positions, atom types—and adapts MaskGIT-style masked parallel decoding to the symmetry-redundant, periodic domain of crystals. The model employs coordinate discretization on the fractional torus, periodic label smoothing to preserve geometric information, symmetry-consistent augmentations, and confidence-ranked greedy decoding. All crystal attributes are predicted as tokens, with symmetry directly encoded in the sequence, enabling space group-stratified sampling for polymorph coverage.

Figure 3: MaskGXT tokenization, training via masked cross-entropy, and sampling conditional on space group posterior.

Discrete coordinate bins ($K=64$) are used, with sub-bin offset regression for geometric refinement, augmenting categorical predictions with periodic precision via bounded $\tanh$ offsets. The architecture fuses site-wise token embeddings and passes them through a QK-normalized, SwiGLU-based bidirectional Transformer trunk, supporting categorical and regression heads for all token streams.

Figure 4: MaskGXT architecture with per-site token fusion, positional and mask encodings, transformer trunk, and categorical/regression prediction heads.

Canonicalization uses Euclidean-normalizer origin shifts and intra/inter-orbit permutation augmentations during training, ensuring symmetry-consistent tokenization across arbitrary descriptions. Ordinal label smoothing with circular neighborhoods regularizes masked learning objectives, directly encoding toroidal adjacency.

Confidence-ranked greedy decoding and space group stratified sampling produce high-accuracy, diverse polymorph candidates. Space group posterior is used to branch sampling, allocating candidates to distinct high-posterior space groups, countering i.i.d. noise-based diversity with targeted polymorph recovery.

Numerical Results: Benchmark Performance and Ablative Analyses

MaskGXT demonstrates superior empirical performance across all standard CSP benchmarks. On MP-20, MaskGXT achieves a match rate (MR) of 73.79% (unfiltered) and 67.06% (validity-filtered), and achieves the lowest RMSE values among all methods, including strong baselines such as MCFlow, OMatG, DiffCSP, FlowMM, and Crystalite. On the more challenging MPTS-52 split, MaskGXT maintains best-in-class match rate and geometric precision (RMSE 0.1004/0.0975).

For polymorph-aware evaluation, MaskGXT substantially outperforms competitors on METRe, reaching 79.06% accuracy on the MP-20 polymorph split, an 8.2% absolute improvement over the strongest baseline (Crystalite 70.87%). Corrected RMSE is also lowest, validating both polymorph and geometric fidelity. Uniform sampling analyses, with fixed $S=2$ structures per composition, confirm MaskGXT’s polymorph coverage advantage with stratified sampling, holding the lead for all values of $S$.

Figure 5: METRe sweep versus sampling budget $S$, MaskGXT stratified sampling dominating on both MP-20 and polymorph splits.

Ablations reveal critical contributions: periodic label smoothing, explicit symmetry tokens, and offset regression each improve accuracy and polymorph coverage. Removing symmetry tokens particularly erodes METRe on polymorph splits due to loss of stratified sampling. Decoder ablations validate the necessity of confidence-ranked greedy decoding and stratification—removing both reduces METRe and cRMSE, confirming their complementary roles.

Implications, Limitations, and Prospects for AI-Driven Algorithm Discovery

MaskGXT, discovered in the agentic HACO loop, substantiates the viability of cross-domain methodology transfer under empirical feedback and high-level human guidance for scientific algorithm discovery. The discrete token representation with direct symmetry tokenization, stratified sampling, and sub-bin offset regression all reflect algorithmic principles unattainable by mere code optimization. The strong numerical results validate the combined architecture and sampling innovations.

Practically, such co-discovery systems are applicable in domains with rapid, aligned empirical validation (e.g., CSP), and may generalize to related crystal generation tasks (e.g., de novo, structure-conditioned atom-type generation). Theoretically, the approach illustrates that search over modeling principles—not only code or hyperparameters—when guided by empirical metrics and domain objectives, enables state-of-the-art algorithm production.

Current limitations relate to domain constraints: empirical loops require fast, reliable validation metrics and controllable compute budgets. In domains with slow experiments, weak proxies, or ambiguous objectives, HACO’s efficacy may diminish, necessitating surrogate modeling and richer human intervention.

Future directions involve: (1) developing agent-managed research state for scalable hypothesis tracking and human judgment delegation; (2) improving proxy evaluation signals to extrapolate short-budget results via scaling laws; (3) expanding beyond metric optimization to exploratory, analytic, and mechanistic experimentation within agentic loops.

Conclusion

MaskGXT leverages masked generative modeling, direct symmetry tokenization, label smoothing, and stratified decoding to achieve state-of-the-art performance in crystal structure prediction. The HACO agentic framework proves that AI co-scientists, empowered by cross-domain search and sparse human guidance, can discover competitive scientific algorithms. The methodology sets a precedent for agent-driven scientific algorithm discovery, suggesting scalable future generalization to wider scientific domains under well-aligned empirical validation.