Predicting Atomistic Transitions with Transformers
Authors: Henry Tischler, Wenting Li, Qi Tang, Danny Perez, Thomas Vogel
Organizations: Computing and Artificial Intelligence Division, Los Alamos National Laboratory, Los Alamos, NM 87545 · School of Engineering and Computer Science, University of Denver, Denver, CO 80210 · Department of Physics and Astronomy, University of Denver, Denver, CO 80210 · Department of Electrical and Computer Engineering, University of Texas at Austin, Austin, TX 78712 · School of Computational Science and Engineering, Georgia Institute of Technology, Atlanta, GA 30332 · Theoretical Division, Los Alamos National Laboratory, Los Alamos, NM 87545 · X Computational Physics Division, Los Alamos National Laboratory, Los Alamos, NM 87545
Accurate knowledge of the atomistic transition pathways in materials and material surfaces is crucial for many material science problems. However, conventional simulation techniques used to find these transitions are extremely computationally intensive. Even with large-scale, accelerated material simulations, the computational cost constrains the applicable domain in practice. Machine learning models, with the potential to learn the complex emergent behaviors governing atomistic transitions as a fast surrogate model, have great promise to predict transitions with a vastly reduced computational cost. Here, we demonstrate how transformers can be trained to predict atomistic transitions in nano-clusters. We show how we evaluate physical validity of the predictions and how a multitude of additional, different microstates can be generated by slightly varying the data provided to the model.
Figures & tables
Figure 1: Three exemplar atomistic transitions originating from the same initial state of a 147-atom Pt nanocluster. Note that while all atoms are of the same element (Pt), we highlight atoms that are involved in each transition, that is, that significantly change their location during the transition, in red. Top row: The initial state, which is the same in each transition, but duplicated to show different colorings for the different atoms involved in each transition. Bottom row: Three different final states.
Figure 2: Diagram of our transformer model. The architecture is nearly identical to the one introduced in [ 26 ] , though the input encoding scheme is modified to allow continuous input. Dotted lines indicate steps not involved in training.
Figure 3: A bi-partite connectivity graph for an illustrative movement in a cluster of 7 atoms. Atom indices give arbitrary indexings.
Figure 4: The cumulative radial distribution function, which aggregates the pair separation distances across all configurations. The position of the valley separating the first and second peak, 3.2 Å, is chosen as the connectivity threshold for creating a connectivity graph for each Pt cluster.
Hyperparameter
Search Space
Learning Rate
[0.0001, 0.1]
Number of Encoder Layers
{1, 2, 3}
Number of Decoder Layers
{1, 2, 3}
Number of Multi-Attention Heads
{2, 4, 5, 10, 20}
Batch Size
[16, 4096]
Feed Forward (FF) Hidden-Layer Size
{1024, 2048}
Table 1: The search space for hyperparameter scans. Hyperparameters that vary continuously were sampled uniformly.
Hyperparameter
Value
PPH
IDM
Learning Rate
0.0248
0.0777
Number of Encoder Layers
1
3
Number of Decoder Layers
2
2
Number of Multi-Attention Heads
10
10
Batch Size
389
197
Table 2: Optimized hyperparameter values for the selected “partial-position-hinted” model (PPH) and “individual displacement” model (IDM). See Table 1 for the search space for each parameter.
Figure 5: Losses from individual trials in our hyperparameter scan. Top: Results our model not provided with hints during training (partial-position hinted). Bottom: Results for our model provided with hints during training (individual magnitude hinted). Plots were generated using Weights & Biases [ 4 ] .
Figure 6: A diagram of our partial-position hinted model. Red colored atoms are those that change their position most during the transition. Top left: an initial state. Bottom left: positions of atoms in its final state that are provided to our model as a “hint”. Right: the final state that is correctly predicted from the initial state and “hint”.
Figure 7: Five examples of output from the partial-position hinted model. (Left) the initial state provided to the model. (Middle) the partial-position hint given to the transformer. (Right) The correctly predicted final state. “Hint Sizes” ( sh ) from top to bottom are: 7%, 11%, 53%, 15%, and 13%.
Figure 8: The amount of additional information required for our partial-position hinted model to correctly predict transitions known from our dataset. Top: The percentage of correct predictions depending on the “hint size” ( Sh ). Bottom: The percentage of correct predictions based on the number of provided final positions. The dotted line separates the atoms in the surface layer (to the right of the dotted line), where transitions are most likely to occur, from the inside of the cluster (on the left of the dotted line).
Figure 9: Transitions predicted by our partial-position hinted model for a single initial state as our hint increases in length. Top: The distinct transition predicted for each possible length of hint provided to our model. Below: Visualizations of these transitions. The initial state (left) and final state (right) are shown alongside a reconstructed minium energy pathway (MEP), calculated using a nudged elastic band (NEB) test. The approximate transition rates are provided in the top right of each plot. Red atoms are those that rearrange during the transition (note that the coloring of the initial state varies between transitions). Hyperdistance gives the cumulative displacement for each frame in the transition pathway.
Correct
Incorrect
Fail
96 %
2 %
2 %
Table 3: Model prediction outcomes for the individual displacement hinted model. “Correct” refers to states with an identical connectivity graph compared to the known final state. “Incorrect” refers to the case when the model predicts a final state different from the expected one in that regard. “Fail” means the predicted final state is equal to the initial state (that is, no transition is predicted).
Figure 10: Selected examples of predictions made by our model without any additional hints. Left: Visualizations of the initial state. Right: Visualizations of the final sate. Red atoms emphasize those that rearrange during the transition. Plots contain the Minium Energy Pathway (MEP) for our prediction, calculated using a Nudged Elastic Band (NEB) test.
Figure 11: The transitions predicted with our perturbation method as the magnitude of the noise added to the core of the initial state is increased. Solid colored bars represent the number of times a specific transition was predicted, the red area indicates the proportion of final states that failed to converge to a metastable state. Visualizations of the most common predicted final states using noise with ηmax=0.5 Å (highlighted) are shown in Fig 12 .
Figure 12: The four most common transitions predicted by perturbing the innermost 20 atoms in the initial state with random noise of ηmax=0.5 Å (cf. Fig. 11 ). In the top left corner of each plot we show the number of times the transition was observed and the average energy of the predicted final state before optimization. In the top right corner of each plot the approximate transition rate is given. None of these transitions were previously known from our simulations. Note that transitions 1 and 3 are very similar, but not identical.
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
General-purpose 3D modeling in chemistry encompasses molecules and materials, requiring both generative and predictive capabilities. However, most existing AI approaches are optimized for a single domain (molecules or materials) and a single task (generation or prediction), which limits representation sharing and transfer. We introduce Zatom-1, a cross-domain, general-purpose model architecture that unifies generative and predictive learning of 3D molecules and materials. Zatom-1 is a deliberately simplified Transformer trained with a multimodal flow matching objective that jointly models discrete atom types and continuous 3D geometries. This approach supports scalable pretraining with predictable gains as model capacity increases, while enabling fast and stable sampling. We use cross-domain generative pretraining as a universal initialization for downstream multi-task prediction of properties, energies, and forces. Empirically, Zatom-1 outperforms or competes with specialized baselines on both multi-task generative and predictive benchmarks in data-controlled settings, while improving generative inference speed by more than an order of magnitude. Our experiments demonstrate positive predictive transfer between data domains from joint generative pretraining: modeling materials during generative pretraining improves molecular property prediction accuracy. Open-source code and model weights are freely available at https://github.com/Zatom-AI/zatom.
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.
Miruna Cretu, Alex Abrudan, Antonia Panescu +11
University of Cambridge · Yale University · MIT +3