ATOM: A Pretrained Neural Operator for Multitask Molecular Dynamics
Authors: Luke Thompson
Davy Guan
Slade Matthews
Dai Shi
Junbin Gao
Andi Han
Venue: ICLR 2026 Poster (2026)
Keywords: Molecular dynamics
neural operator
transformer
ai for science
equivariant
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
TL;DR: We present the first neural operator to our knowledge that generalizes to out-of-domain molecular dynamics trajectories and performs strongly on large molecules; we release TG80, a multitask MD dataset (>2.5M fs, 80 compounds) to support future work.

Abstract: Molecular dynamics (MD) simulations underpin modern computational drug discovery, materials science, and biochemistry. Recent machine learning models provide high-fidelity MD predictions without the need for repeated quantum-mechanical force calculations, enabling significant speedups over conventional pipelines. Yet many such methods typically enforce strict equivariance and rely on sequential rollouts, thus limiting their flexibility and simulation efficiency. They are also commonly single-task, trained on individual molecules and fixed time frames, which restricts generalization to unseen compounds and extended timesteps. To address these issues, we propose Atomistic Transformer Operator for Molecules (ATOM), a pretrained transformer neural operator for multi-task molecular dynamics. ATOM adopts a quasi-equivariant design that does not require an explicit molecular graph and employs a temporal attention mechanism to enable accurate, parallel decoding of multiple future states. To support operator pretraining across chemicals and timescales, we curate TG80, a large, diverse, and numerically stable MD dataset with over 2.5 million femtoseconds of trajectories across 80 compounds. ATOM achieves state-of-the-art performance on established single-task benchmarks, such as MD17, RMD17, and MD22. After multi-task pretraining on TG80, ATOM shows exceptional zero-shot and robust generalization to unseen molecules across varying time horizons. We believe ATOM represents a significant step toward accurate, efficient, and transferable molecular dynamics models.

Published as a conference paper at ICLR 2026
A PRETRAINED NEURAL OPERATOR FOR MUL- TITASK MOLECULAR DYNAMICS
ATOM:
Luke Thompson
Davy Guan Data61 CSIRO
Dai Shi University of Cambridge
The University of Sydney
Slade Matthews*
Junbin Gao* The University of Sydney
Andi Han*1
The University of Sydney
The University of Sydney
ABSTRACT
Molecular dynamics (MD) simulations underpin modern computational drug dis- covery, materials science, and biochemistry. Recent machine learning models provide high-fidelity MD predictions without the need to repeatedly solve quantum mechanical forces, enabling significant speedups over conventional pipelines. Yet many such methods typically enforce strict equivariance and rely on sequential rollouts, thus limiting their flexibility and simulation efficiency. They are also com- monly single-task, trained on individual molecules and fixed timeframes, which restricts generalization unseen compounds and extended timesteps. To address these issues, we propose Atomistic Transformer Operator for Molecules (ATOM), pretrained transformer neural operator for multitask molecular dynamics. ATOM adopts a quasi-equivariant design that requires no explicit molecular graph and employs a temporal attention mechanism, enabling accurate parallel decoding of multiple future states. To support operator pretraining across chemicals and timescales, we curate TG80, large, diverse, and numerically stable MD dataset with over 2.5 million femtoseconds of trajectories across 80 compounds. ATOM achieves state-of-the-an performance on established Linle-taskk benchmarks, such MD17, RMD17 and MD22. After multitask pretraining on TG80, ATOM shows exceptional zero-shot generalization to unseen molecules across varying time hori- zons. We believe ATOM represents significant step toward accurate, efficient, and transferable molecular dynamics models.
INTRODUCTION
Molecular dynamics (MD) serves as a computational microscope of atomic motion and is now integral drug discovery and materials science pipelines (Dror et al., 2012: De Vivo al., 2016). In ab initio molecular dynamics, quantum mectnnncechaiccll density functional theory (DFT) is used to compute atomic forces, and the resulting equations of motion are integrated to generate high-fidelity trajectories. However, DFT's s computational complexity scales at least cubically with the number of atoms, and relies on double-precision arithmetic that limits GPU acceleration (Kresse & Furthmüllerr, 1996; Stein et al., 2020; Li et al., 2024).
Neural approaches have recently emerged as a promising solution to the scalability bottleneck. Equivariant architectures, in particular, encode physical symmetries model interatomic dynamics, achieving ab initio-level accuracy at significantly reduced computational cost (Batzner et al., 2022; Musaelian et al., 2022; Batatia et al., 2022; 2023; Xu et al., 2024). While equivariance is often deemed essential for improving generalization, strict symmetry preservation involves substantial tradeoffs (Xu al., 2024; Schreiner et al., 2023). Architectures that enforce exact equivariance at every layer often increase computational overhead, restrict model expressivity, and complicate optimization (Fuchs et al., 2020; Brehmer al., 2023; Elhag et al., 2025). It is unclear whether symmetry constraints can be relaxed without sacrificing accuracy for molecular dynamics.
"These authors contributed equally as senior authors. 1 "ooresponndecee andi han@sydney. edu. au.

---

Published as a conference paper at ICLR 2026
Furthermore, most existing methods for molecular dynamics autoregressive, predicting the next state based on the current one (Köhler et al., 2019: Fuchs et al., 2020; Thiemann et al., 2025). Autoregressive approaches often struggle to capture fong-horizon temporal dependencies and accu- mulate error as the prediction horizon grows (Bengio et al., 2015; Bergsma et al., 2023; Taieb & Atiya, 2016). Inference speeds are also constrained by the need for sequential integration, failing to exploit modern, highly parallel compute architectures. One exception is Equivariant Graph Neural Operator (EGNO) (Xu et al., 2024), which models the entire trajectory with neural operator learning. Nevertheless, EGNO enforces strict equivariance and is single-task in nature, i.e., it is trained and evaluated on trajectories of each molecule separately with fixed time horizon, which limits zero-shot generalization to unseen molecules or timeframes.
Our Main Contributions. this work, we address the above issues regarding equivariance, autoregression, and zero-shot generalization within unified framework, which we call Atomistic Transformer Operator for Molecules (ATOM). To this end, we propose a pre-trained neural operator with transformer backbone for molecular dynamics and introduce new MD dataset, TG80, which is both chemically diverse and numerically stable for multitask pretraining and benchmarking.
Design innovations. ATOM quasi-equivariant by employing an equivariant lifting layer that produces symmetry-aware features, while allowing subsequent transformer blocks to be uncon- strained for flexibility and expressiveness. Unlike autoregressive models, ATOM allows parallel decoding of molecule states across multiple timesteps, directly learning the trajectory operator. By encoding time lags via a novel temporal rotary position embedding, ATOM enhances temporal interpolation and extrapolation, enabling robust predictions across multiple time horizons. Finally, ATOM requires no predefined molecular graph and operates directly on point clouds, naturally accommodating long-range spatial interactions without the need for hand-crafted connectivity.
Performance highlights. ATOM sets state-of-ther woffthh-arttto single-task MD benchmarks. For larger, sparsely connected molecules in MD22, ATOM significantly outperforms existing graph- based baselines by capturing the long-range atomic interactions. In the multitask regime, we pretrain ATOM on TG80 trajectories from multiple molecules and varying timeframes, demonstrating significant zero-shot transfer to both unseen molecules and timesteps, improving existing baselines by 39.75% on average. This achieves performance on par with existing specialized baselines tailored for such molecules and timeframes. To the best of our knowledge, this the first method that demonstrates such generalization capability in molecular dynamics.
We believe our work represents shift in molecular dynamics modeling, where we demonstrate the potential of quasi-equivariance designs and zero-shot generalization to out-of-domain systems, which is enabled by the comprehensive TG80 MD dataset. ATOM and TG80 are available at this repository.
2 RELATED WORK
Equivariant Neural Networks. Equivariance (to transformations such rotation, reflection, and translation) has emerged as an essential physics- -nnfics-informed prior for deep learning models on molecular data (Bronstein al., 2021; Duval et al., 2023). Early works employed convolutional approaches to achieve translation equivariance in E(3) (Weiler et al., 2018; Wu et al., 2020) or tensor product attention and spherical harmonies enforce roto-translational equivariance in SE(3) (Fuchs et al., 2020; Thomas et al., 2018). In contrast, message passing neural network (MPNN) frameworks, such as Equivariant Graph Neural Network (EGNN) and others (Garcia Satorras et al., 2021: Gasteiger et al., 2021: Huang et al.. 2022), achieve equivariance by operating on strictly equivariant features, such as inter-node distances and directions. While effective, MPNNs typically assume a fixed molecular graph. This is problematic when the underlying structure contains non-local interactions and dynamic bonding effects (e.g., resonances, transient interactions). which render predefined graphs inaccurate over time (Knutson et al., 2022; Luo et al., 2021). To address this issue, model molecules as point clouds, with our attention represented as a fully connected graph that allows unrestricted information propagation across the molecule.
Time-coarsened Molecular Dynamics Time coarsening is a coarse moase-nraining method which preserves molecular structure, but compresses many short integration steps into few large-stride updates to reduce the cost of long-time simulation (Kmiecik et al., 2016). Stochastic coarse- coarse-grainig approaches often learn transition kernels on configuration space, bypassing explicit integration of the equations
2

---

Published as a conference paper at ICLR 2026
of motion. Klein et al. (2023) learns such a kernel with normalizing flow and uses it as an MCMC proposal targeting the Boltzmann distribution, Hsu et al. (2024) uses a conditional diffusion model to learn transition probability matrix, and Yu et al. (2025) uses fow-matching to learn vector field transporting current states to future states. Closer to our framework, deterministic methods such MDNeet (Zheng al., 2021) and TrajCast (Thiemann et al., 2025) learn GNN and EGNN, respectively, which autoregressively predict fixed strides 10-100 times larger than those of MD integrators. Bigi et al. (2025) incorporates Hamiltonian structure and explicit energy-conservation. Most of the methods require direct force learning and are sequential in nature, while ATOM may be interpreted as a force-free deterministic coarse- graineiggginggggggn approach, wherein temporal pushforward is approximated by a learned propagation operator which is decoded in parallel.
Neural Operators. Neural operators are deep learning methods for learning operators between function spaces (Kovachki et al., 2021). A wide variety of architectures have been proposed for such operator learning. Notably, Fourier Neural Operator (FNO) (Li et al., 2021) learns an operator in the Fourier domain, while its derivatives G-FNO (Helwig et al., 2023) and PINO (Li et al., 2023b), respectively, add group equivariance and physics-informed properties. Xu et al. (2024) bridges this framework with molecular dynamics by recasting the task as learning propagation operator that evolves historical atomic positions into their future configurations. Specifically, EGNO (Xu et al., 2024) integrates EGNN and FNO layers learn dynamic trajectories, capturing spatial and temporal correlations. Recently, transformer neural operators (Bryutkin et al., 2024; Hao et al., 2023; Li et al., 2023a) have surpassed the performance of FNO in most partial differential equation (PDE) tasks. Notably, OFormer (Li et al., 2023a) uses a linear Galerkin-type attention mechanism, which omits the softmax and instead interprets the latent column vectors as basis functions. General Neural Operator Transformer (GNOT) (Hao et al., 2023) employs a novel subquadratic cross- cross-attention methodology to integrate multiple feature types (e.g.. shape and point relationships) into their transformer blocks. With ATOM, we unify the MD problem formulation and temporal discretization approach introduced by EGNO with the increased representational power of transformers operator settings.
MD Benchmarks. Research on graph machine learning for molecular dynamics suffers from poor benchmarking (Bechler-Speichee et al., 2025). For example, despite the fact that MD17 Benzene exhibits non-physical noise approximately 1000 times higher compared to other compounds (Christensen & von Lilienfeld, 2020), is still regularly employed to benchmark new models (Bihani et al., 2023; Huang et al., 2022; Liao & Smidt, 2023; Xu et al., 2024). The practical relevance of single-task learning on these datasets is also dubious, as predicting trajectories for molecules with existing numerical solutions offers minimal benefit. We believe the strengths of neural approaches emerge in transfer learning, where models generalize to unseen compounds, thereby circunventing the computational costs associated with explicit numerical simulations. This motivates our development of TG80 to facilitate multitask dynamics learning across molecular systems.
3 ATOMISTIC TRANSFORMER OPERATOR FOR MOLECULES (ATOM)
In this section, we first introduce the problem formulation (Section 3.1) and then propose the framework of ATOM by introducing the key model and training designs (Section 3.2). We then discuss the multitask pretraining for ATOM and introduce TG80 MD dataset (Section 3.3).
3.1 PROBLEM FORMULATION
We follow (Xu et al., 2024) to cast molecular dynamics prediction as operator learning. We model molecule of N atoms as point cloud in 13,, which we denote as (t)) for a given system state time t. In partieular, we write (())_ = (xx"), (t) v ( that represent molecules in terms of the atom positions x and velocities v. Our objective is to predict future trajectory gif+a0), where At E [0, AT).
Similar to (Xu et al., 2024), we focus on predicting the position states only. Let [0, AT] + R- be the trajectory function mapping At U(Di) E NR3 representing molecule positions At in the future. We assume a solution operator Ft: ((3) - 24 exists which provides the underlying future trajectory given system states at t. Thus, the goal of molecular dynamics prediction be- comes training a neural operator (g(()) to approximate the target trajectory function Ft (g())): ming Eg0, ((Fo( (()))))))))))(((()))))))))))) for some loss function L: 24 x l - R. Here, expectation is with respect the different initial states, By discretizing over the temporal domain and considering
(t)))))(tt
3

---

Published as a conference paper at ICLR 2026
L2 loss, optimize the neural operator with discretized temporal sampling of the states:
min Σ gie |ra (g
(Atp) (a())
(F) (Alp)
-F (g())
where (At1,-. Atp) are discrete timesteps. We replace the true future state (T( with the known future ground truth node positions x(++A(p) for Alp E [0, AT).
Quasi-cquivariance. We formally define quasi-equivariance, motivated by (Elhag al., 2025).
Definition 3.1 (s-quasi-equivariancee) We call a function f : X + V e-quasi-equivarine with respect to group if it satisfies Exexll Jo (((())())) Sa (())(()))))()) S E. where μ denotes the normalized Haar measure.
In practice, we approximate the group integration with Monte Carlo samples from G.
Single- and multitask. Unlike prior works (Schreiner et al., 2023; Xu et al., 2024). we consider both single-task and multitask settings. Single-task refers the case where separate model is independently trained and evaluated on each molecule and fixed timeframes. This corresponds to the conventional practice in molecular dynamics benchmarks. Multitask instead pretrains one unified model on several molecules across varying time lags and evaluates out-of-domain trajectories on unseen molecules, thereby directly testing zero-shot cross-molecule generalization. Under multitask setting, the objective (1) computes the expectation over trajectories of multiple molecules.
3.2 ATOM MODEL AND TRAINING DESIGN
Here we outline the pipeline of ATOM. At its core is an equivariant lifting layer (Section 3.2. 1), which maps atomic positions, velocities and their phase features into richer embedding space while preserving symmetry under the Euclidean group E(3). The lifted embeddings are then processed by the ATOM attention block, which applies heterogeneous atlention over positions, velocities, and phase features with chemical augmentation (Section 3.2.2). To capture temporal dynamics, we incorporate temporal rotary position embedding (T-RoPE) (Section 3.2.2) that depends only on time lags and is shared across atoms, ensuring translation invariance in time and permutation invariance within each molecule. The parameterized ATOM can be written by
Fg = P (KL) c.... (KK,) Q
where Q, P denotes the equivariant lifting and projection operators respectively, =1, are the data-dependent kernels induced by cross attention (See Appendix G.I), and a denotes some nonlinear activation function.
e
Finally, to counter numerical noise in training trajectories, we inject randomly sampled position and velocity perturbations during training (Section 3,2.3), which improves robustness and acts as a regularizer against overfitting. The overall pipeline of ATOM is in Figure 1,
3.2.1 E(3) EQUIVARIANT LIFTING
To model atomic states in a symmetry-rsppecing way, each atom is encoded with its 3D position and velocity, augmented with their norms:
(x, /x2 + y2 + x2), v (Var, Ugy U2, 22 uz + v2). To construct higher-dimensional features that remain consistent with E(3) symmetry, we apply equivariant lifting that maps the inputs through learnable functions that preserve group actions. Specifically, we use E(3) equiivaria linear layers (Geiger & Smidt, 2022) that lift the position and velocity vectors to feature space wherein they satisfy equivariance con- straints by construction. We further construct phase space featrues for each atom by augmenting the position and velocity vectors with atomic number, which is subsequently processed by learnable equivariant layer to obtain lifted representation. The final lifted embedding for molecule given by (X, V, Z) 3x corresponding position, velocity and phase features. The second dimension aggregates nodes and time for attention and d, is the embedding space dimension.
We highlight that after the equivariant lifting layer, do not enforee equivariance in the subsequent transformer blocks. This relaxation improves performance over fully equivariant designs while still showing robustness random trajectory rotations versus non-equivariant baselines (see Section 4.4).
4

---

Published as a conference paper at ICLR 2026
Multitask Pretraining
Inference
ATOM
Decoded positions
TG80
Query molecule
Equivariant
*
lifting
Heterogenoms Temporal Aftention
Plane specs
T-Rope
to 251
Padiies
f2
Query times
Makes
750
LogUnet(s))
Nalcitin
tp
Figure I: ATOM Pipeline. We pretrain ATOM on the TG80 dataset across multiple molccules with stochastic time lags. At inference, ATOM takes query molecule and timestamps and directly outputs corresponding molecular states.
3.2.2 ATOM HETEROGENEOUS TEMPORAL ATTENTION
We employ a heterogeneous temporal attention mechanism to enable mixing between multiple features (X, V, Z) 3 NPxd, across spatial and temporal dimensions. We use the phase space embedding Z as the query and attend to the key-value pairs formed from all features X, V, Z E RxP. In Figure 5, we show that this improves performance by 6.36% over standard self-attention for single- task prediction. To encode temporal information, introduce Temporal RoPE (T-RoPE), adapting Rotary Position Embedding (RoPE) (Su et al., 2023) to irregular time lags by scaling the phase rotations according cumulative timestamps constructed from per-step increments At.
Let the hidden dimension per head be d, (even). We define frequencies W/ = b 2k2//d, for k 0,..., d,./2-,,2111 Given per-step time increments { we build timestamps L, = 1+ Di= Aly, and assign a single rotation to all N atoms at timestep p: Rp diag (R(O,.o)..... R R(0ppd//-1)) E Rdd xd., where 6p.
(lp fo) and R(0) E R22x2 is the rotation matrix with angle 0 and 7 > is a timescale hyperparameter. Suppose the query molecule state at time p is given as Op E Nxda and key molecule state at time p' is Kp E IR We apply Rp. Ry to Op. Kp respectively so that the rotary dot product Q,R,KK RR,) depends only the time interval by 1p. This makes attention trunslation invariant in time, which allows for interpolation and extrapolation across irregular increments {Atp}. In addition, sharing the same R, across all N atoms in a molecule ensures iottuttaationinvaiannn within a timestep. For aggregated query and key matrices Q K NPxd. denote the application of temporal RoPE across P timesteps and N atoms as T-RoPE(QQ, T-RoPE(K) E TR N Pxd,.
Specifically, single-head attention layer of ATOM computes
Σ OF softmax (-RoPF(Q(Z)) T-RoPE(K(F)) V(F), 2
ddh
where (-), V-) represent the query, key and value projections. We introduce learnable weights OF to modulate the relative importance of each feature. In Appendix G.I, we show that heterogeneous attention is equivalent to a kernel integral operator and discuss its properties.
3.2.3 TRAINING WITH LABEL NOISE REGULARIZATION
Many DFT datasets are inherently noisy (Christensen & von Lilienfeld, 2020), and MD models can overfit this noise. Motivated by the regularization effect of label noise (Damian al., 2021; HaoChen et al., 2020), we augment the observed node positions x and velocities v by random Gaussian noise Ex Ev o2) during training. Let g(()
(x!")
(()
Endt (() + E,,)) be the
5

---

Published as a conference paper at ICLR 2026
noised initial state at time t. We minimize the following regularized loss
P
min Σ
Egw E.E
Fe (ge)) (At,) (x (x))) + -0
similar strategy has also appeared in graph neural network (GNN)-based MD models and neural operator pretraining (Dauparas et al., 2022; Zhou et al., 2024a; Hao et al., 2024). We only apply noise augmentation during training and evaluate on the unperturbed ground-trum truuh trajectories.
Comparison to EGNO. We highlight that ATOM adopts different design choices compared to EGNO. First, EGNO is an EGNN operating on fixed bond connectivity, whereas ATOM uses an E())-qquivariant lifting layer followed by globally connected point-cloud attention, which better handles long-range and sparsely bonded interactions. Second, EGNO is strictly equivariant end-to- end. while ATOM is quasi-equivariant, enforcing equivariance only the lifting stage and relaxing in deeper transformer layers, which our ablations show improves accuracy. Third, EGNO models time via Fourier temporal convolution, whereas ATOM uses Temporal RoPE, allowing translation- invariant handling of irregular time gaps and stronger temporal extrapolation. T-RoPE also uniquely allows modifying the time-horizon AT al inference by modulating the rotary phases (Section 3.2.2). Consequently, pretrained ATOM can be evaluated at arbitrary AT values without retraining.
3.3 MULTITASK ATOM PRETRAINING AND TG80 DATASET
This section adapts ATOM for the multitask setting, where the aim is to predict future trajectories for unseen molecules. In order to more effectively distinguish molecules, we construct radius graph of 1.6 Å based on atomic positions, and apply random walk positional encoding (Ma et al., 2023; Lobato et al., 2021) to augment the phase vector z. We describe the process in detail in Appendix D.2 and highlight that such graph depends only on atomic positions, not chemical bonds.
During multitask training, each mini-batch contains trajectories from multiple molecules. In addition, we perform random sampling for the time lags At from a log-uniform distribution between Atmin and AT' namely At ~ LogUnif Atmin AT'). This aims enhance the robustness of interpolation and extrapolation in the temporal domain, a consideration that has been similarly explored in (Schreiner et al., 2023). Let M denote the set of training molecules and let ((I) represent the state of molecule m E M at timestamp t. We can write the pretraining multitask objective as
min
Σ MEM
E 9th Ar~LogUnif(MMmann AT)
Fo (g.h At) (At) -x
(1+At)
M
2
where we take expectation with respect initial states of multiple molecules in the training set, as well as the time lags. Here, ATOM also takes time lag, Al, as input, to modulate T-RoPE phase.
TG80 Dataset. To facilitate pretraining of our neural operator, we introduce TG80, superset the MD17 dataset. The initial seed set comprises 40 molecules: 8 MD17 compounds and 32 additional drug-like molecules selected through expert review. We then augment the seed molecules with structurally similar molecules from the PubChem dataset of 173 million compounds (Bolton et al., 2011). Accepted candidates had an BCFP-4 Tanimoto similarity between 0.875 and 0.925 to at least one seed molecule, and more than 0.80 similarity to previously accepted molecules, alongside other criteria detailed in Appendix C.4 (Landrum et al., 2025; Rogers & Hahn, 2010; Rogers & Tanimoto, 1960). These thresholds follow common practice in the literature, balancing diversity while avoiding collapse into overly narrow chemical subspaces (Matter, 1997; Menke et al., 2021; Eastman al., 2023; Harper et al., 2004: Zhang et al., 2023).
We generate all trajectories using ORCA V6.01 (Neese, 2022) with the PBE functional (Perdew et al., 1996), deP basis set (Weigend & Ahlrichs, 2005), A4 dispersion corrections (Caldeweyher et al., 2019; 2020; Wittmann et al., 2024) at one femtosecond resolution, 300K temperature, in vacuum. This resembles an enhanced RMD17, with more modern dispersion corrections to improve stability and allow for a larger step size (Christensen & von Lilienfeld, 2020). As a result. TG80 exhibits more diverse dynamics and improved numerical stability, with no compound exceeding 50 Á center-of-mass drift in Figure 8!
"Siuulations ran on 32 AMD EPYC 7543 cores with 256 GB RAM per molecule, totalling 806,400 CPU- hours (quoted market cost USD 150 000).

---

Published as a conference paper at ICLR 2026
4 EXPERIMENT RESULTS
Metrics. We use State-to-trajectory (S2T) and state-to-state (S2S) error to evaluate ATOM (Xu al., 2024). Specifically, S2T Ep=n | xyl12. measures the average discrepancy between
the predicted & and ground-truth positions x across entire trajectories, while S2S = || - xp quantifies the error at the final predicted timestep.
Baselines. For comparison, we include a range of classic to state-of-the-art baselines, including Radial Field (RF) (Köhler et al., 2019), Tensor Field Networks (TFN) (Thomas et al., 2018), SE(3) Transformer (SE(3)TTr.) (Fuchs et al., 2020), E(n) equivariant graph neural networks (EGNN) (Garcia Satorras et al., 2021), MACE (Batatia al., 2023), and EGNO (Xu et al., 2024). We note that MACE is pretrained on the authors' own dataset, so this is not strictly like-for-like comparison. Our EGNN baselines are EGNN-Rollout (EGNN-R), which predicts timesteps autoregressively, and EGNN-
(EGNN-S), which uses the output of each GNN as the prediction of given frame. We set all baseline hyperparameters following previous works (Xu et al., 2024; 2022; Shi et al., 2021) and tune ATOM and EGNO hyperparameters as in Table 19 and Table 20.
Training setups. For training of ATOM and EGNO, we consider two temporal discretization strategies in selecting the timestamps ip = + > Atr: (1) Uniform discretization selects tp = t P/PAT and (2) Tail discretization selects t,, = Á + P/P(AT A) for lag A [0, AT]. In the main paper, we present experiment results with uniform discretization and include the results with tail discretization in Appendix E. We perform early stopping on the lowest S2S validation loss checkpoint and report results as mean 22 over three training runs. All experiments are run on an NVIDIA @ RTX 5080 with wall-clock time and FLOP utilization detailed in Table 15.
4.1 SINGLE- -AALE-TASKAAA LEARNING
We benchmark on the MD17, RMD17, and MD22 DFT MD trajectory datasets (Chmiela et al., 2017: Christensen & von Lilienfeld, 2020; Chmiela et al., 2023). We partition the trajectories into cain/validation/test splits of sizes 500/2000/2000, set AT = 3300 fs and P = 8, and train for 2500 epochs following (Xu et al., 2024). For the performance MD17 (Table 1), we directly quote the results from (Xu et al., 2024) except for EGNO. We design ATOM to have six transformer blocks with hidden size of 256.
MD17 and RMD17. As shown in Table 1, ATOM compares favorably with ofate-of-the-art (SOTA) baselines MD17 dataset, yielding average reductions of 14.96% (S2S mean squared error (MSE) and 8.3% (S2T MSE) on average?. In Table 9 (Appendix E. we benchmark ATOM on RMD17, and observe similarly competitive performance against EGNO.
MD22. To evaluate performance on larger molecules, we consider Ac- Ala3-NHMe (20 heavy atoms), docosahexaenoic acid (DHA with 24 heavy atoms). and stachyose (45 heavy atoms) from the MD22 dataset (Chmiela et al., 2023). ATOM remains competitive on these systems; whereas EGNO fails to converge (Table 2). We attribute this discrepancy to differing inductive biases: GNNs such as EGNO restrict message passing to predefined bond or radius graph and can therefore under- represent long-range, non-bonded steric and electrostatic interactions that dominate the behavior of large, sparsely connected molecules (Alon & Yahav, 2021; Kosmala et al., 2023). This explains the poor performance of EGNO on MD22, which contains prototypically sparse molccules such as DHA. shown in Figure (Nv et al., 2003). We further disentangle the role of connectivity from the use of attention by training a variant, ATOM-GATv2. in which our heterogeneous temporal attention is replaced by graph attention network v2 (GATv2) layers (Brody et al., 2022) operating on the same bond/radius graph as EGNO. ATOM-GATV2 still substantially underperforms the full ATOM model, indicating that the performance gains stem from the fully connected point-clous interaction pattern rather than from attention alone.
OH
Docosahex- aenoic acid (DHA)
Figure 2:
We exclude henzene from the table due to the previously discussed high numerical noise.
7

---

Published as a conference paper at ICLR 2026
Table Single-task MSE (x10-3) on MD17. Upper part: S2S MSE. Lower part: S2T MSE
Aspirin
Ethanol
Malonaldehyde
Naphthalene
Salicylic
Toluene
Uracil
RF
10.9410.02
4.042
13.93_0.06
0.50
1.23.10.04
10.931..08
0.64
TFN
12.37.0.38
=.00
13.62_0.16
0.49
1.03.1.00
10.89
0.841..04
-H
11.12+0.12
4.74=0.02
13.89=0.04
0.52=0.02
1.13+00.4
10.88=0.12
0.79:0.04
EGNN
14.41=0.30
4.64=0.04
13.64=0.02
0.47=00.04
1.0220.04
11.78=0.14
0.64:0.02
EGNN-R R
9.96-20.14
4.61=0.01
13.04=0.03
0.44:0.05
0.96:20.01
10. 00:
1.11:0.04
EGNN-S
10.25:10.09
4.61=0.01
13.06-0.01
0.53
1.06...05
10. 0..09
0.62110.11
EGNO
9.64
=0.01
12.92
0.30010.00
1..01
10.00
0.58
MACE
6.95 +-0.00
2.06 +0.00
17.99 +0.26
0.72 +00.0
1.05 +-0.00
6.44 +00.00
0.75 +-0.00
ATOM
6.82+0.06
3.52+0.04
14.72=0.01
0.50=0000
0.88
4.66
0.63++0.00
EGNN-R
7.35+a.19
3.21=0.00
10.75#0.04
0.34=
1.09
4.53=2.00
0.89+0002
EGNN-S
9.01+0.34
3.21=0.00
11.20=0.03
0.42=0.01
1.41#0.00
4.86=2004
0.65++0.01
EGNO
9.64+0.15
4.57=0.01
12.92=0.00
0.39=0000
0.90=a.01
10.99=0.00
0.58=0.02
MACE
5,06 10.00
2.84 :0.00
16.09 10.03
0.57 20.00
0.55 1:0000
3.26 10.00
1.08 10000
ATOM
5.62
2.62:10.04
12.49=0.01
0.43=0000
0.86
2.27
0.61
Table 2: Single-task MSE (x10-3) on MD222. Upper: S2S. Lower: S2T
Table 3: Multitask S2T MSE (x10-2) on TG80 across five UMAP cluster assignments.
Ac-Ala-NHHME
DHA
Stachyose
Chuder
Chuder
Cluster
Chuder
Cluster
BGNO
357. 80
178.50g4.
42.11a0
44
141
150.020011
107 47.34
XTOM-GATV2
223.57am1
46.72
41 41.00
ATOM
9.651177
10.60
21.25
78.05%
80.805
88.005
89.78%
83 10%%
Gap
497. 37%
+94.00%
+-30.44%
MACE
184.25
221.12
826.97
316.26
229
EGNO
ATOM-GATV2
232 75
115. 115.3620
16.502033
116.4448.34
20.70
30.84
OOD
RONN-R
RONN-S
EGNO
45 15.8
11.85 en
d5.00am at
7S86 15a as
115
151 74
152 152.72ea so
188.55e
461 20..22aaom
161. S4ann
110 10
11.3 05
114 314.30m
ATOM
7.5G
9.66
18.13
ATOM
35. J5.0G.a.5
106 990
60.95
66 08 .8.c.
47 17.4911.50
Gap
+96.75%
+91.70%
Gap
21.98%
2.407
58.81%
68.7196
4.2 MULTITASK LEARNING ON TG80
We pretrain ATOM on TG80, scaling to six attention blocks with hidden size of 256. We se- lect stochastic horizons AT ~ LogUnn ( fs, 000 fs) and use a five-fold, cluster-based cross- validation. Specifically,, we compute ECFP-4 fingerprints (Rogers & Hahn, 2010), embed them using UMAP (Melnnes et al., 2018), and apply agglomerative clustering (Ward, 1963) to partition compounds into ten disjoint clusters. The folds are then formed by holding out clusters, ensuring that the train/validation/ees sets occupy distinct regions of chemical space. This cluster-wise protocol minimizes leakage and more closely reflects the prospective scientific setting in which models must generalize to unseen molecules. Cluster-based approaches present more challenging generaliza- tion problems than random splits or common ceanical-scafflle -aaeed splits (Guo et al., 2024). In Appendix E.2, also consider pretraining
standard random split of molecules.
Table 3 benchmarks ATOM by assessing both in-distribution (ID) and out-of out-of-domain (OOD) S2T performance. For the in-distribution setting, we train, validate, and test on molecules from the same cluster. We observe that ATOM outperforms existing baselines by an average of 83.96% in terms of S2T MSE. We then assess out-of-domain (OOD) generalization performance by predicting the dynamics of unseen compounds drawn from disjoint clusters. Under OOD settings, ATOM nearly halves the S2T MSE of EGNO, with an average improvement of 39.74% across five cluster splits. Notably, OOD ATOM beats ID EGNO performance in four of five folds. This striking zero-shot generalization, realized without any exposure to the test molecules. confirms that ATOM uniquely learns robust, transferable knowledge of molecular dynamics. In Appendix F.2, we show similar outperformance in S2S prediction. In Appendix F.2, we show that the significantly improved multitask performance comes with modest overhead in training time and in inference latency.
4.3 TEMPORAL GAP AND TIMESTEP INVARIANCE PROPERTIES
AT Invariance. We evaluate the performance of pretrained ATOM (at fixed AT = 3000) with varying AT at inference. We compare ATOM. EGNO, and EGNN on S2T MSE by fixing P = 8 and sweeping AT logarithmically from 10 to 10000 fs on an in-distribution (Cluster 1) multitask model. In Figure 3, we show that ATOM maintains its extrapolation advantage across the range compared to
8

---

Published as a conference paper at ICLR 2026
EGNO, particularly at larger AT. Ablating T-RoPE (NoPE) removes this advantage by exhibiting an EGNO-like error trend with substantially higher MSE. This underscores T-RoPE's s role in stable time-gap extrapolation.
P Invariance. Following the discrctization invariance in neural operators, we expeet ATOM and EGNO models to show consistent MSE as P varies under uniform discretization (Kovachki et al., 2021). Figure 4 confirms such a conjecture by showing that multitask ATOM pretrained at P = maintain constant S2T MSE as ranges from 24 inference.
46
14
42-
40-
36
20
34-
12-
107
12
16
20
Figure 3: Pretrained (AT = 3000, ID) multi-
Figure 4: Pretrained (P = 8) ATOM and
task S2T MSE across varying AT values.
EGNO are discretization invariant, showing stable S2T MSE.
4.4 ABLATION STUDIES
We perform extensive ablations to assess each design choice in ATOM. For single-task performance (Fig. 5) and multitask performance (Fig. 6), we independently toggle components and measure their contributions. Our analysis focuses on equivariant lifting, T-RoPE, label-noise regularization, heterogeneous attention, and random-walk positional encoding (under multitask pretraining).
Equivariant lifting. We assess the quasi- equivariant design against non-equivariant ATOM. As shown in Figure 5, replacing the equiv- ariant lifting introduced in Section 3.2.1 with stan- dard linear layers (no equivariant lifting) markedly degrades the performance of ATOM, increasing S2T MSE by 22.48. Table 4 further quantifies sensitivity SO(3) rotations: ATOM's S2T MSE increases by 10.80x under rotation, compared to 19.77x without equivariant lifting. Notably, the fully-equivariant variant of ATOM, described in appendix Appendix D.I, also underperforms ATOM in both single-task (Figure 5) and multitask (Figure 6) settings, with the gap exaggerated in the multitask setting. This aligns with recent findings on relaxed equivariance, suggesting that strict equivariance can limit model capacity and complicate the optimization process (Elhag et al., 2025). We present estimates of the quasi-equivariance E in Appendix E.3.
Table 4: S2T MSE (x10-2) of a fixed input frame rotated and unrotated by an SO(3) matrix.
ATOM
No cquivariant Lift
Unrotated
6. 7....0.0
33.44: 664.97.945.80
Rotated
73.04 .11
Increase
66.28=26.az
627.53+
Rotation penalty (x)
10.80
19.77
Heterogencous attention. We find that substituting heterogeneous temporal attention with standard self-attention the phase space features increases S2S MSE by 0.47, suggesting that cross-attention enables access to non-trivial feature interactions.
Temporal Rotary Position Embedding (T-RoPE). In the single-task regime with fixed AT (Fig- ure 5), T-RoPE contributes little to the performance of ATOM, as it effectively reduces to a constant rotational shift. By contrast, with stochastic AT, disabling T-RoPE (NoPE) increases MSE by 1.07, consistent with ATOM leveraging the 7 parameter to encode variable time gaps (Figure 6). An EGNO-style sinusoidal positional encoding produces similar performance degradation.
Label noise regularization. We also test the utility label noise regularization as in Section 3.2.3. From Figure 5, we observe that removing augmented noise from the position and velocity features in-

---

Published as a conference paper at ICLR 2026
equivarianz lift
Pully Equivariat-
Fully equivariant-
Sinusaidal PE-
No noise
NoPE
No oquávartias lif
NoFE
ATOM
ATOM
8
10
40.0
42.5
45.0
47.5
50.0
52.5
55.0
Mcan S25 MSE (xc10 P)
Mean S25 MSB (x10
"
Figure ATOM ablation on MD17 Aspirin.
Figure 6: ATOM ablation on TG80 Cluster
creased S2T MSE by 1.21. For the multitask ablation on TG80, suppress label noise regularization, as the dataset is designed to be numerically stable with small noise.
RWPE. We assess aadom--walk positional encoding (RWPE) in the multitask pretraining. Figure indicates that RWPE facilitates molecule identification, yielding improved multitask performance.
5 CONCLUSIONS
In this work, we demonstrate that carefully designed transformer neural operators enable zero-shot generalization to unseen chemical dynamics. Our experiments on MD17 demonstrate continued good single-task performance, and we present the first molecular neural operator that can successfully learn large molecule dynamics using MD22. Our multitask experiments show that our method learns transferable dynamics knowledge, even without explicit graph representations. In combination with our TG80 dataset, we provide a large-scale open-source benchmark and baselines to evaluate future models and spur further operator research with concrete scientific applicability.
Limitations We remark that TG80 does not contain trajectories for large molecules with more than 15 heavy atoms, despite their obvious chemical and pharmacological relevance. In follow-up work, we intend to enrich TG80 with such molecules, calculated with a higher-resolution DFT basis set, wB97X-3c (Müller al., 2023). Regarding ATOM, lacks an explicit energy- -ased inductive bias, which may permit long- horizoorizon drift. A natural extension is therefore framewise energy head (x.,,)) with force supervision F,, = -xx,, Ex,,,,. This energy term also defines the drift in the Langevin dynamies, where an additional stochastic term accounts for thermal perturbations of atomic positions. Incorporating such physics- stochastic dynamics into our operator learning framework a natural next step, and we view this as a promising direction for future MD research.
REPRODUCIBILITY STATEMENT
provide experiment details, such as choice of hy perrparrmeetess other training configurations in Appendix F. In addition, we will release the TG80 dataset upon acceptance under MIT license for reproducibility.
10

---

Published as a conference paper at ICLR 2026
REFERENCES
Uri Alon and Eran Yahav. On the Bottleneck of Graph Neural Networks and its Practical Implications, March 2021. URL httttp//arxiv.orree
052C5. arXiv:2006.0522 [cs].
Ilyes Batatia, Simon Batzner, Dávid Péter Kovács, Albert Musaelian, Gregor C. Simm, Ralf Drautz, Christoph Ortner, Boris Kozinsky, and Gábor Csányi The Design Space of E(3) E(3)-Equivariant Atom- Centered Interatomic Potentials, 2022. URL https: httppss//aaxxxvccoo org abs 220//22052220 06643. eprint: 2205.06644
Ilyes Batatia, Dávid Péter Kovács, Gregor N. C. Simm, Christoph Ortner, and Gábor Csányi. MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields, January 2023. URL http: arriiv orgg abs/ 2206. 077.7. arXiv:2206.0769) [stat].
Simon Batzner, Albert Musaelian, Lixin Sun, Mario Geiger, Jonathan P. Mailoa, Mordechai Kornbluth, Nicola Molinari, Tess E. Smidt, and Boris Kozinsky. E(3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials. Nature Communications, 13(1), May 2022. ISSN doi: 10. 1038//4 0088/444477-0000 011-442-22999--22 URL http: dx doi.. wwrg/0....... 541467-022-2992
2041-1723,
Publisher: Springer Science and Business Media LLC.
Maya Bechler-Speicher, Ben Finkelshtein, Fabrizio Frasca, Luis Müller, Jan Tönshoff, Antoine Siraudin, Viktor Zaverkin, Michael M. Bronstein, Mathias Niepert, Bryan Perozzi. Mikhail Galkin, and Christopher Morris. Position: Graph Learning Will Lose Relevance Due To Poor Benchmarks, 2025. URL https: htt//ariiiv////////v///./////aaa ht/p/aaxvvv.rg///. abs / 2502. 14546. -eprint: 1502.14546
Samy Bengio, Oriol Vinyals, Navdeep Jaitly, and Noam Shazeer. Scheduled Sampling for Sequence Prediction with Recurrent Neural Networks, 2015. URL https: /twpivrivvv..oo/aa/aaaa www. 03099. eprint: 1506.03099.
www//ba.aaa..oooooo..e/...e.e.ae abs /1506.
Shane Bergsma, Timothy Zeyl, and Lei Guo. SutraNets: Sub-series Autoregressive Networks for /2312.
Long- Song-Sequence, Probabilistic Forecasting, 2023. URL https:///atttte 14880. _eprint: 14880.
Filippo Bigi, Sanggyu Chong, Agustinus Kristiadi, and Michele Ceriotti. FlashMD: long-stride, universal prediction of molecular dynamics, October 2025. URL http: www.aawwwww....o.aaaooooa aab/ 2505.19355 arXiv:2505. 19350 [physics].
Vaibhav Bihani, Utkarsh Pratiush, Sajid Mannan, Tao Du, Zhimin Chen, Santiago Miret, Matthieu Micoulaut, Morten M. Smedskjaer, Sayan Ranu, and N. M. Anoop Krishnan. EGraFFBench: Evaluation of Equivariant Graph Neural Network Force Fields for Atomistic Simulations, 2023.
URL https:://// www.i.avvoo..... abs/ abs/a/////// /2110. 02 428. .eprint: 2310.02428.
Evan E, Bolton, Jie Chen, Sunghwan Kim, Lianyi Han, Siqian He, Wenyao Shi, Vahan Simonyan, Yan Sun, Paul A. Thiessen, Jiyao Wang, Bo Yu, Jian Zhang, and Stephen H. Bryant. PubChem3D: new resource for scientists. Journal of cheminformatics, 3(1):32, September 2011. ISSN 1758-2946. Place: England.
doi: 00.1186/1758-2946-1
Johann Brehmer, Pim de Haan, Sönke Behrends, and Taco Cohen. Geometric Algebra Transformer, 2023. URL https: / /awpiv.ivvv//////i///.../a/aaavv org/ abs http:/aab///3.111/ 2205
eprint: 2305.18415.
Shaked Brody, Uri Alon, and Eran Yahav. How Attentive are Graph Attention Networks?, 2022. URL eprint: 2105.14491
https: /tte:///xxxiie arxiv. / abs 2105.14499
Michael M. Bronstein, Joan Bruna, Taco Cohen, and Petar Veličković. Geometric Deep Learning: Grids, Groups, Graphs, Geodesics, and Gauges, 2021. URL https://arxiw.org/ 2104. 13478. _eprint: 2104.13478.
Andrey Bryutkin, Jiahao Huang, Zhongying Deng, Guang Yang, Carola-Bibiane Schönlieb, and Angelica Sviles-Rivers HAMLET: Graph Transformer Neural Operator for Partial Differential Equations, 2024. UtLhttps: //tx///a/vv//////////////// wwwp://ooo///b///////oo/e 03541. eprint: 2402.03541.

---

Published as a conference paper at ICLR 2026
Eike Caldeweyher, Sebastian Ehlert, Andreas Hansen, Hagen Neugebauer, Sebastian Spicher, Christoph Bannwarth, and Stefan Grimme, A generally applicable atomic-charge dependent London dispersion correction. The Journal of Chemical Physics, 150(15), April 2019. ISSN 63 /1 1 5090222.
1089-7690. doi: 10. 11636/155000020 URL http:// /xxxxxxoooo wex.doi..ogge 0
Publisher: AIP Publishing.
Eike Caldeweyher, Jan-Michael Mewes, Sebastian Ehlert, and Stefan Grimme. Extension and evaluation of the D4 London- -ispeersonnn model for periodic systems. Physical Chemistry Chemical Physics, 22(16))8499-851 8512, 2020. ISSN 1463-9084. doi: 10. 10.1039/00c0005022. URL http:
wwa.xdoi.orgo.. 1039
www.0039/DDD0P000 /DDD//DOOPPOO20.00 Publisher: Royal Society of Chemistry (RSC).
David Case, Hasan Metin Aktulga, Kellon Belfon, David S. Cerutti, G. Andrés Cisneros, Vinícius Wilian D. Cruzeiro, Negin Forouzesh, Timothy J. Giese, Andreas W. Götz, Holger Gohlke, Saeed Izadi, Koushik Kasavajhala, Mehmet C. Kaymak, Edward King, Tom Kurtzman, Tai- Sung Lee, Pengfei Li, Jian Liu, Tyler Luchko, Ray Luo, Madushanka Manathunga, Matias R. Machado, Hai Minh Nguyen, Kurt A. O'Hearn, Alcxey V. Onufriev, Feng Pan, Sergio Pantano, Ruxi Qi, Ali Rahnamoun, Ali Risheh, Stephan Schott-Verdugee Verdugo. Akhil Shajan, Jason Swails, Junmei Wang, Haixin Wei, Xiongwu Wu, Yongxian Wu, Shi Zhang, Shiji Zhao, Qiang Zhu, Thomas E. III Cheatham, Daniel R. Roe, Adrian Roitberg, Carlos Simmerling, Darrin M. York, Maria C. Nagan, and Kenneth M. Jr. Merz. AmberTools. Journal of Chemical Information and Modeling, 63(22)):6833611) October 2023. ISSN 1549-9596. doi: 10100///c..cimm... URL https: /oii/oi..oooo or org 112.0.1.11111 2000//21sccccco jeim. 3c01153 Publisher: American Chemical Society.
153.
Stefan Chmicla, Alexandre Tkatchenko, Huziel E. Sauceda, Igor Poltavsky, Kristof T. Schütt, and Klaus-Robert Müller. Machine learning of accurate energy-conserving molecular force fields. Science Advances, 3(5):e1633015 2017. doi: 10.11266sciadv 16030156 _eprint:
URL https: //t:://www www. www.:wscceene.. science, org/ /ooi/aabtt/aoo 110.00..00..
Stefan Chmiela, Valentin Vassilev-Galindo, Oliver T. Unke, Adil Kabylda, Huziel E. Sauceda, Alexandre Tkatchenko, and Klaus-Robert Müller. Accurate global machine learning force fields for molecules with hundreds of atoms. Science Advances, 9(2, January 2023. ISSN 2375-2548. doi: 10. 1 a.dfsciadv..dff0773. URL http://p/// //dx doi org/11 0 1 6/ s/ssciadk adddv73...877. Publisher: American Association for the Advancement of Science (AAAS).
Anders S Christensen and O Anatole von Lilienfeld. On the role of gradients for machine learning of molecular energies and forces. Mach. Learn. Sci. Technol., 1(4):045018, October 2020. Publisher: IOP Publishing.
Alex Damian, Tengyu Ma, and Jason D. Lee. Label Noise SGD Provably Prefers Flat Global Mini- mizers, December 2021. URL http: /ttp:///aaiiii/ wwwpxiiivorrg/aaae hbtp://bb//20000 06530. arXiv:2106.06530 [es].
J Dauparas, Anishchenko, N Bennett, H Bai, R J Ragotte, L F Milles, B I M Wicky, A Courbet, R de Haas, N Bethel, P J Y Leung, T F Huddy, S Pellock, D Tischer, Chan, B Koepnick, H Nguyen, A Kang, B Sankaran, A K Bera, N P King, and D Baker. Robust deep learning- based protein sequence design using ProteinMPNN. Science (American Association for the Advancement of Science), 378(661 5):49-56,, 2022. ISSN 0036-8075. Backup Publisher: Lawrence Berkeley National Laboratory (LBNLL), Berkeley, CA (United States) Place: United States Publisher: The American Association for the Advancement of Science.
Marco De Vivo, Matteo Masetti, Giovanni Bottegoni, and Andrea Cavalli. Role of Molecular Dynamics and Related Methods in Drug Discovery. Journal of Medicinal Chemistry, 59(9): 4035-4061, May 2016. ISSN 0022-2623. doi: 10.1021/aa 10.1021/acs.jmede httaa/s.smcccemm....t. 5b0teeem..b0...... URL https: ///doi.org/// 100 1021/acs. http://ccsjjmdcceem.... jmedchem. 5b016 68 Publisher: American Chemical Society.
Blaise Delattre, Quentin Barthélemy, Alexandre Araujo, and Alexandre Allauzen. Efficient Bound of Lipschitz Constant for Convolutional Layers by Gram Iteration, June 2023. URL http:
//waxxi..rg//as/////ooaooooooasa 16173. 73. arXiv:2305.16173 [es].
12

---

Published as a conference paper at ICLR 2026
Ron O. Dror, Robert M. Dirks, J. P. Grossman, Huafeng Xu, and David E. Shaw. Biomolecular simulation: a computational microscope for molecular biology. Annual Review of Biophysics, 41: 429-452, 2012. ISSN 1936-1238. doi: 10. Setpuuuree-booohs-4299110
Alexandre Duval, Simon V. Mathis, Chaitanya K. Joshi, Victor Schmidt, Santiago Miret, Fragkiskos D. Malliaros, Taco Cohen, Pietro Lio, Yoshua Bengio, and Michael Bronstein. A Hitchhiker's Guide to Geometric GNNs for 3D Atomic Systems. pp. arXiv:2312.07511 December 2023. doi: 10.485500arxxv22311.77111
Vijay Prakash Dwivedi, Anh Tuan Luu, Thomas Laurent, Yoshua Bengio, and Xavier Bresson. Graph Neural Networks with Learnable Structural and Positional Representations, 2022. URL
attps://azzee /awpvv.aivvvvvv org/ /2100, 07875. .print: 2110.07875.
Peter Eastman, Pavan Kumar Behara, David L. Dotson, Raimondas Galvelis, John E. Herr, Josh T. Horton, Yuezhi Mao, John D. Chodera, Benjamin Pritchard, Yuanqing Wang, Gianni De Fabritiis, and Thomas E. Markland. SPICE, A Dataset of Drug-like Molecules and Peptides for Training Machine Learning Potentials. Scientifiic Data, 10(1), January 2023. ISSN 2062-4463. doi: 10.1038/ URL http: / /xx doi. 1tt..//o///...13311131----0-2000
-41597-022-01889
Publisher: Springer Science and Business Media L.C..
Ahmed A. Elhag, T. Konstantin Rusch, Francesco Di Giovanni, and Michael Bronstein. Relaxed Equivariance via Multitask Learning, 2025. URL httpps///aaxxxv..///aaaaaa/////
_eprint: 2410.17878.
Fabian B. Fuchs, Daniel E. Worrall, Volker Fischer, and Max Welling. SE(3)-Transfommes:: 3D Transsrannltiion Equivariant Attention Networks. arXiv e-prints, pp. arXiv:2006. 10503, June 2020. doi: 10448550/arXiv.2000..000 10503. .eprint: 2006.10503.
Nicholas Gao, Eike Eberhard, and Stephan Günnemann. Learning Equivariant Non-Local Elec- tron Density Functionals, October 2024. URL http:h/t///cc htt////iww.orggb//////eaaaaasa// 07977777.
arXiv:2410.07972 [es] version: 1.
Victor Gareia Satorras, Emiel Hoogeboom, and Max Welling. E(n) Equivariant Graph Neural Networks. arXiv e-prints, pp. arXiv:2102.09844, February 2021. doi: 10.48550/axii22200009444
_eprint: 2102.09844.
Johannes Gasteiger, Florian Becker, and Stephan Günnemann. GemNet: Universal Directional Graph Neural Networks for Molecules. pp. arXiv:2106.08900: June 2021. doi: 10.48550/arXiv.21166 08903.
Mario Geiger and Tess Smidt. e3nn: Euclidean Neural Networks, 2022. URL https://a/xx org/abs/ 2207. 09453. eeprint: 2207.09453.
Qianrong Guo, Saiveth Hernandez-Hernandezz and Pedro J. Ballester. Scaffold Splits Overestimate Virtual Screening Performance, June 2024. URL arXiv:2406.00873 [q-bio].
wwwpa/////4000000 00873.
Zhongkai Hao, Zhengyi Wang, Hang Su, Chengyang Ying. Yinpeng Dong, Songming Liu, Ze Cheng, Jian Song, and Jun Zhu. GNOT: A General Neural Operator Transformer for Operator Learning, June 2023. URL http::tt//// //tr://aiiii// org/ 2otp//220...2 14376. arXiv:2302.1 14376.
Zhongkai Hao, Chang Su, Songming Liu, Julius Berner, Chengyang Ying. Hang Su, Anima Anand- kumar, Jian Song, and Jun Zhu. DPOT: Auto-Regressive Denoising Operator Transformer for Large-Scale PDE Pre- Trainiinn May 2024. URL http:: org htt.r//abb/2403.00.3 03542..
arXiv:2403.0348 [cs].
Z. HaoChen, Colin Wei, Jason Lee, and Tengyu Ma. Shape Matters: Understanding the Implicit Bias of the Noise Covariance, June 2020. URL arXiv:2006.08e8 [es].
08600
G Harper, S Pickett, and D. Green. (Research Papers) Design of Compound Screening Collection for use in High Throughput Screening. Combinatorial Chemistry &amp: High Throughput Screening, 7 (1):63-70, February 2004. ISSN 1386-2073. doi: 202111/113222044772888833 URL http:// dx doi. Ltd.
10. / /tt.:/0..227/11127777772222222
Publisher: Bentham Science Publishers
13

---

Published as a conference paper at ICLR 2026
Charles R. Harris, K. Jarrod Millman, Stéfan J. van der Walt, Ralf Gommers, Pauli Virtanen, David Cournapeau, Eric Wieser, Julian Taylor, Sebastian Berg, Nathaniel J. Smith, Robert Kern, Matti Picus, Stephan Hoyer. Marten H. van Kerkwijk, Matthew Brett, Allan Haldane, Jaime Fernández del Río, Mark Wiebe, Pearu Peterson, Pierre Gérard-Marchant, Kevin Sheppard, Tyler Reddy, Warren Weckesser, Hameer Abbasi, Christoph Gohlke, and Travis E. Oliphant. Array programming with NumPy. Nature, 585(7825)3357-362 September 2020. doi: 10. (088///4158-0022-26666 https: /top:://di... www.ii..gg/aii www.ooggg///oooo 1038 s11586 020 64 2. Publisher: Springer Science and
URL
Business Media LLC.
Jacob Helwig, Xuan Zhang, Cong Fu, Jerry Kurtin, Stephan Wojtowytsch, and Shuiwang Ji. Group Equivariant Fourier Neural Operators for Partial Differential Equations, 2023. URL https:
//aaxxiv.org//aaeee...aaaa 2306. 05697. -eprint: 2306.05697.
Tim Hsu, Babak Sadigh, Vasily Bulatov, and Fei Zhou. Score Dynamics: Scaling Molecular Dynamics with Picoseconds Time Steps via Conditional Diffusion Model. Journal of Chemical Theory and Computation, 2036):2335-23444 March 2024. ISSN 1549-9618, 1549-9626. doi: 10.1s2//accsjoo //ccc.jctc.. 3c01361. URL https: /pbbs.ub...uu acs. org/ doid/00000000011 /10. 021 acs. jetc. 3c01361.
Wenbing Huang, Jiaqi Han, Yu Rong, Tingyang Xu, Fuchun Sun, and Junzhou Huang. Equivariant Graph Mechanics Networks with Constraints, 2022. URL / 06442. -eprint: 2203.06442.
Keller Jordan. NanoGPT Speedrun 11/06/24, June 2024. URL
com/ http:////mmtttteer////oo/oo/ 2u1/255_Subbin//// 810e0-888442333 01818884229866888 txt.
Lalleeuurrdaa/mmoddddd-aaooooo 1d3bd933-
Keller Jordan, Yuchen Jin, Vlado Boza, Jiacheng You, Franz Cesista, Laker Newhouse, and Jeremy
Bernstein. Muon: An optimizer for hidden layers neural networks, 2024. URL https:
//kellerjordan batborddanniihuub..... wttghhi.ioopp..oo.eeeeeooeo wtcpi//psst.. s/mun//
Keller Jordan, Jiacheng You, Cesista Franz Louis, and Brendon Hogan. URL
NanoGPT Speedrun 01/16/25,
January 2025.
com/
Kel delllerdordanm/mddee htwp:///wwddddeedddnnndooooo/e/e haaadddnnnnoo//ooooooo http://pt..bbbb////e//eeeae/ee/e
httpe//ereoorrdd11//......1111 1d3bd93bba -90e--aa9e----00 4118 4118- aeb8- 9184239d7566.
Sékou-Oumar Kaba, Arnb Kumar Mondal, Yan Zhang, Yoshua Bengio, and Siamak Ravanbakhsh. Equivariance with Learned Canonicalization Functions, July 2023. URL http:://txix// org/ abs/2211. 06489, trXiv:2211.06433 [cs].
Hyeongchan Kim. vitorch..ottmiizg optimizer & Ir scheduler & loss function collections in Py PyTorch, September 2021. URL https: wttps//w...oooooeeeeeeeee/o www.kozzstt.. hccm///oossst/yytttcoooooooe
Leon Klein, Andrew K. Foong, Tor Erlend Fjelde, Bruno Mlodozeniee, Mare Brockschmidt, Sebastian Nowozin, Frank Noé, and Ryota Tomioka. Timewarp: Transferable Acceleration of Molecular Dynamics by Learning Time-Coarsened Dynamics, December 2023. URL http: 01170. arXiv:2302.01 [stat].
Sebastian Kmiecik, Dominik Gront, Michal Kolinski, Lukasz Wieteska, Aleksandra Elzbieta Dawid, and Andrzej Kolinski. Coarse-Grrinee Protein Models and Their Applications. Chemical Reviews, 116(14)77988-9966 July 2016. ISSN 1520-6890. doi: 10..002/acs..chemme...bul
Carter Knutson, Gihan Panapitiya, Rohith Varikoti, and Neeraj Kumar. Dynamic Molecular Graph- based Implementation for Biophysical Properties Prediction, 2022. URL atttss://axxe eprint: 2212.09991.
org/abs// 2212. 09991.
Arthur Kosmala, Johannes Gasteiger, Nicholas Gao, and Stephan Günnemann. Ewald-based Long- Range Message Passing for Molecular Graphs, June 2023. URL http:://axxxiw.tt///// org. org//bs/////aaaoo///////// 2303. 04791. arXiv:2303.08 .44999 [cs].
Nikola Kovachki, Zongyi Li, Burigede Liu, Kamyar Azizzadenesheli, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Neural Operator: Learning Maps Between Function Spaces. arXiv e-prints, pp. arXiv:2108.08493 August 2021. doi: 10.48550/a2x..21088.8822 2108.08481.
eprint:
14

---

Published as a conference paper at ICLR 2026
Dávid Péter Kovács, J. Harry Moore, Nicholas J. Browning, Ilyes Batatia, Joshua T. Horton, Yixuan Pu, Venkat Kapil, William C. Witt, Ban--oodaan Magdāu, Daniel J. Cole, and Gábor Csányi MACE-OFF: Short-Range Transferable Machine Learning Force Fields for Organic Molecules. Journal of the American Chemical Society, (47(21):1759--1177 17611, May 2025. ISSN 0002-7863. URL https: /doi. org 10. jacs 4c07099. Pub- lisher: American Chemical Society.
doi: 10.
G. Kresse and J. Furthmüllerr Efficiency of ab-initio total energy calculations for metals and semiconductors using plane-wave basis set. Computational Materials Science, 6(1):15-50, 1996. ISSN 0927-0256. doi: http:////c.....ee..a.ea/////////
URL https:
//ww.sscienciirees: com/
Jonas Köhler, Leon Klein, and Frank Noé. Equivariant Flows: sampling configurations for multi- body systems with symmetric energies. arXiv e-prints, pp. 1rX1v:19100007553, October 2019. doi:
10.48550/arXiv. http://rxxi..11..0000.0 1910.00753.
eprint: 1910.00753.
Greg Landrum, Paolo Tosco, Brian Kelley, Ricardo Rodriguez, David Cosgrove, Riccardo Vianello, sriniker, Peter Gedeck, Gareth Jones, Eisuke Kawashima, NadineSchneider, Dan Nealschneider, Andrew Dalke, Matt Swain, Brian Cole, Samo Turk, Aleksandr Savelev, tadhurst cdd, Alain Vaucher, Maciej Wójeikowski, Ichiru Take, Rachel Walker, Vincent F. Scalfani, Hussein Faara, Kazuya Ujihara, Daniel Probst, Juuso Lehtivarjo, guillaume godin, Axel Pahl, and Niels Maeder. rdkit/rdkit: 4025_03_2 (QI 2025) Release, 2025. URL https: /zoooo.. org/ /rgg/d3// do 1100. 5281/zenodo. 591637.
Ask Hjorth Larsen, Jens Jørgen Mortensen, Jakob Blomqvist, Ivano E. Castelli, Rune Christensen, Marcin Dulak, Jesper Friis, Michael N. Groves, Bjørk Hammer, Cory Hargus, Eric D. Hermes, Paul C. Jennings, Peter Bjerre Jensen, James Kermode, John R. Kitchin, Esben Leonhard Kolsbjerg, Joseph Kubal, Kristen Kaasbjerg, Steen Lysgaard, Jón Bergmann Maronsson. Tristan Maxson, Thomas Olsen, Lars Pastewka, Andrew Peterson, Carsten Rostgaard, Jakob Schiøtz, Ole Schütt, Mikkel Strange, Kristian S. Thygesen, Tejs Vegge, Lasse Vilhelmsen, Michael Walter, Zhenhua Zeng, and Karsten W. Jacobsen. The atomic simulation environment a Python library for working with atoms. Journal of Physics: Condensed Maiter, 29327),73302,, 2017. URL http: //stacks iop. org/ 898 4
Vu-Anh Le and Mehmet Dik A Mathematical Analysis of Neural Operator Behaviors, October 2024.
URL http: //rxivv org/ abs/ .11.111.1111.10 214 arXiv:2410.21481 [math].
Jianxiong Li, Boyang Li, Zhuoqiang Guo, Mingzhen Li, Enji Li, Lijun Liu, Guojun Yuan, Zhan Wang, Guangming Tan, and Weile Jia. Scaling Molecular Dynamics with ab initio Accuracy to 149 Nanoseconds per Day. In SC24: International Conference for High Performance Computing, Networking. Storage and Analysis, pp. 1-15. IEEE, November 2024. doi: 10.1109/sc41406.2024 00036. URL http: //tx//xx////// d orgg / 10. 1 1011119/11111 4066 2024 00036.
Zijie Li, Kazem Meidani, and Amir Barati Farimani. Transformer for Partial Differential Equa- tions' Operator Learning, 2023a. URL htttts:///aatttt 2205.13671.
/ttp:://arxiiii www.iiooorg//bbaa wttp:////aa//...1..222 -eprint:
Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Fourier Neural Operator for Parametric Partial Differential Equations, May 2021. URL htth: arxiv. org.or/g/aoo..... //as/////// aabs/2211 2019.08895. arXiv:2010.08895.
Zongyi Li, Hongkai Zheng, Nikola Kovachki, David Jin, Haoxuan Chen, Burigede Liu, Kamyar Azizzadenesheli, and Anima Anandkumar. Physics-Informed Neural Operator for Learning Partial Differential Equations, 2023b. URL https: /tapss///axxiv wwtpa/wwwwwwrg.s///oaa/aaaaaaas/ 2111.03794.
abs 03794. eeprint:
Yi-Lun Liao and Tess Smidt. Equiformer: Equivariant Graph Attention Transformer for 3D Atomistic abs 2200 11 1990. _eprint: 2206. 19900.
Graphs,
Alvaro Lobato, Miguel A. Salvad6, J. Mamuel Recio, Mercedes Taravillo, and Valentín G. Baonza. Highs and Lows of Bond Lengths: Is There Any Limit? Angewandte Chemie International Edition, 60(31) http: //l:://xxx.oo doi org 1ttp://...1000//1/1 /a./0////iiia/emm/ 2021 02967. Publisher: Wiley.
June 2021. ISSN 1521-3773. doi: a..1000/miie.2222229667 URL
15

---

Published as a conference paper at ICLR 2026
Iyy Loshchilov and Frank Hutter. Decoupled Weight Decay Regularization. arXiv e-prints, pp- arXiv:1711.0500., November 2017. doi: 00448500aaxii.111.05100.. .eprint: 1711.05101.
Shitong Luo, Chence Shi, Minkai Xu, and Jian Tang. Predicting Molecular Conformation via Dynamic Graph Score Matching. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P. S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 19784-19795.
Curran Associates, Inc., 2021. URL
cttpppsscccpp//e cc/paper. pdff.
Liheng Ma, Chen Lin, Derek Lim, Adriana Romero-Soriano, Puneet K. Dokania, Mark Coates, Philip Torr, and Ser-Nam Lim. Graph Inductive Biases in Transformers without Message Passing, May 2023. URL /tttp://arxive org/ abss/////////// 2305. 17589, arXiv:2305. [cs].
Hans Matter. Selecting Optimally Diverse Compounds from Structure Databases: A Validation Molecular Descriptors. Journal of Medicinal Chemistry, 40(8))1219 April 1997. ISSN 1520-4804. doi: 10. 1mm00052m00 URL http:
Study of Two-Dimensional and
//dx. doi /tt 1021/ jm9 60352. Publisher: American Chemical Society (ACS).
Leland McInnes, John Healy, and James Melville. UMAP: Uniform Manifold Approximation and Projection for Dimension Reduction. arXiv e-prints, pp. arXiv:1802.03422 February 2018. doi: 10.48550/arXiv. 1802.03426. _eprint: 1802.03426.
Janosch Menke, Joana Massa, and Oliver Koch. Natural product scores and fingerprints extracted from artificial neural networks. Computational and Structural Biotechnology Journal, 19:4593- 4602, 2021. ISSN 2001-0370. doi: (tt0///...ssj.200000 (at.j.2021.07..00 07.032. URL http: /xx doi org/ 10.1016 / csbjssssj... 2021 07. 032. Publisher: Elsevier BV.
Albert Musaelian, Simon Batzner, Anders Johansson, Lixin Sun, Cameron J. Owen, Mordechai Korn- bluth, and Boris Kozinsky. Learning Local Equivariant Representations for Large-Scale Atomistic Dynamics, 2022. URL https:/ /ttpp::////xvvvv.//// org/ aaa//bbs//22000 05249. .eprint: 2204.05249.
Marcel Müller, Andreas Hansen, and Stefan Grimme. omegaB97X-3c: composite range-separated hybrid DFT method with a molecule-optimized polarized valence double-zeta basis set. The Journal of Chemical Physics, 158(1), January 2023. ISSN 1089-7690. doi: 10. 1063//.00133066. URL http:/tp//cc
org. 1063/5 0133026.
Publisher: AIP Publishing.
Frank Neese. Software update: The ORCA program system-Versiinn 5.0. WIREs Computational
Molecular Science, 12(5)xe600,, 2022. doi: https://oo.orr///.100//vwwww 1606. URL https:
htwwrres. www.:eswwnnnnnnnnenne........eee witaaaa..wiwlll..... com com doi abs 100.10000000000000 /wems 1606. _eprint:
Eldho Nv, Feller Se, Tristram-Nagle S, Polozov Iv, and Gawrisch K. Polyunsaturated docosahexaenoic vs docosapentaenoic acid-differences in lipid matrix properties from the loss of one double bond. Journal of the American Chemical Society, 125(21), May 2003. ISSN 0002-7863. doi: 10.1021/ ja029029o. URL https: /pwbmed..ddd/ ncbi nlm. nih gov/12785780. Publisher: J Am Chem Soc.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Edward Yang, Zach DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Py Torch: An Imperative Style, High- Perfooreere Deep Learning Library. arXiv e-prints, pp. arXiv:1912.01103 December 2019. doi: 104485500arxx..1112217700 URL https: / /ui adsabs harvard. edu/ abs 221222111222 9arl1iaiiiiiii w1t.vv11111111111
John Perdew, Kieron Burke, and Matthias Ernzerhof. Generalized Gradient Approximation Made Simple. Physical Review Letters, 77(18):385-38668 October 1996. ISSN 1079-7114. doi: 10. 10.113/pyssreeeet.7....3. 77.3886. URL http:///// dx www.oooorg/000 wtt.://og//.....//a.e.ee.... w1t...11/////yeeeett.t..o 3865. Publisher: American Physical Society (APS).
David Rogers and Mathew Hahn. Extended-CComectiiit Fingerprints. Journal of chemical informa- fion and modeling, 50(5):742-754, May 2010. ISSN 1549-9596. doi: 10.1021/ct 1000000..
16

---

Published as a conference paper at ICLR 2026
David J. Rogers and Taffee T. Tanimoto. A Computer Program for Classifying Plants. Science, 132: 1115-1118, October 1960. ISSN 0036-8075. doi: 10. 1226sscience.32.3333... 132.3434.1115.
Mathias Schreiner, Ole Winther, and Simon Olsson. Implicit Transfer Operator Learning: Multiple Time-Resolution Surrogates for Molecular Dynamics, October 2023. URL http://arxiv. org/abs/ 2305. 18046 6, arXiv:2305. 18046 [physics].
Peter Schwerdtfeger and David J. Wales. 100 Years of the Lennard-Jones Potential. Journal of Chemical Theory and Computation, 20(9):3379 3405, May 2024. ISSN 1549-9618. doi: 10.1021/ Acs.jetc.4c00135. URL hww.ttttt///t...t. //oii wrg/orrg.//./. 11011111111111 1021 1020001csssssc /acs /tt///ccssccc.....oeeaaeaaaaaoet jctc. 4c00135. American Chemical Society.
Publisher:
Chence Shi, Shitong Luo, Minkai Xu, and Jian Tang. Learning Gradient Fields for Molecular Conformation Generation, 2021. URL https: /ttps://arxiv.cro aaaaiiiiivv..... www.vvorr//aaa abs http://a./222....5.1 2105.03902.
.eprint:
Frederick Stein, Jürg Hutter, and Vladimir V. Rybkin. Double-Hybrid DFT Functionals for the Condensed Phase: Gaussian and Plane Waves Implementation and Evaluation. Molecules, 25(21), 2020. ISSN 1420-3049. 1420- 3049 00/044//2////////
doi: 10.330//olecules52115174. URL https: : //www mdpi com/
Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu. RoFormer: Enhanced Transformer with Rotary Position Embedding, November 2023. URL hwwpthtpt////tt/ http://wxxvvvrro./// abs/2108 04.004..00 098 64, arXiv:2104.0968 [cs].
Souhaib Ben Taieb and Amir F. Atiya. A Bias and Variance Analysis for Multistep-Aheet Time Scries Forecasting. IEEE Transactions on Neural Networks and Learning Systems, 27(1):62-76, 2016. doi:
Fabian I. Thiemann, Thiago Reschüitzegeer, Massimiliano Esposito, Tseden Taddese, Juan D Olarte- Plata, and Fausto Martelli. Force-free molecular dynamics through autoregressive equivariant networks. arXiv preprint arXiv:2503.23794, 2025.
Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang. Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation- and translation-eque -equivarnee neural networks for 3D point clouds. arXiv e-prints, pp.
February 2018. doi: 11.550//rXiv..802.082111 eprint:
1802.08219.
Joe H. Ward. Hierarchical Grouping to Optimize an Objective Function. Journal of the American Statistical Association, 58(301):236-244, 1963.
doi: 10.10
016214559.9963.1050088 1963. 10500845. Publisher: Taylor & Francis .eprint:
URL https: //w::///wwww tandfonl wiiiieccom. ////doiaal
10. 10.1080001622110
1055000085
Florian Weigend and Reinhart Ahlrichs. Balanced basis sets of split valence, triple zeta valence and quadruple zeta valence quality for H to Rn: Design and assessment of accuracy. Physical Chemistry Chemical Physics, 7(18):3297, 2005. ISSN 1463-9084. doi: 10. 10.1039/b508541a
URL http:// doi. orrii.org/100 1039//50881// Publisher: Royal Society of Chemistry (RSC).
Maurice Weiler, Mario Geiger, Max Welling, Wouter Boomsma, and Taco Cohen. 3D Steerable CNNs: Learning Rotationally Equivariant Features in Volumetric Data, 2018. URL https: //
02547. -eprint: 1807.02547.
Lukas Wittmann, Igor Gordiy, Marvin Friede, Benjamin Helmich- Helmich-Paris, Stefan Grimme, Andreas Hansen, and Markus Bursch. Extension of the D3 and D4 London dispersion corrections to the full actinides series. Physical Chemistry Chemical Physics, 26(32):21379 2024. ISSN 1463- 9084. doi: 10. 1r.103994cp01151555 URL http:// /xx www.doo..ooooooooooooooooooo..o. ort///rgggo/ 10 10100//00000.1130/0/000000002220 Publisher: Royal Society of Chemistry (RSC).
Wenxuan Wu, Zhongang Qi, and Li Fuxin. PointConv: Deep Convolutional Networks on 3D Point Clouds, 2020. URL https: /ttpx//wvvvvo/////// http://wwwww///////////////////e 07246. eprint: 1811 1811.07246
17

---

Published as a conference paper at ICLR 2026
Minkai Xu, Lantao Yu, Yang Song, Chence Shi, Stefano Ermon, and Jian Tang. GeoDiff: Geometric Diffusion Model for Molecular Conformation Generation, 2022. URL eprint: 2203.02923.
abs/2203. 02923.
Minkai Xu, Jiaqi Han, Aaron Lou, Jean Kossaifi, Arvind Ramanathan, Kamyar Azizzadenesheli, Jure Leskovec, Stefano Ermon, and Anima Anandkumar. Equivariant Graph Neural Operator for Modeling 3D Dynamics, June 2024. URL attp://aarxvv..orr//a////// arXiv:2401.11077.
Ziyang Yu, Wenbing Huang, and Yang Liu. UniSim: A Unified Simulator for Time-Coarsened Dynamies of Biomolecules, October 2025. URL
03157.
arXiv:2506.03157 [q-bio].
Ziqiao Zhang, Bangyi Zhao, Ailin Xie, Yatao Bian, and Shuigeng Zhou. Activity Cliff Prediction: Dataset and Benchmark, February 2023. URL httte://arai..orrr//////222...
arXiv:2302.07541 [q-bio].
Tianze Zheng, Weihao Gao, and Chong Wang. Learning Large-Time-Step Molecular Dynamics with Graph Neural Networks, December 2021. URL htt.ttp://axxx..orr///// abs/ httpa/bb//11111111111
arX1v:2111.15176 [physics].
Anthony Zhou, Cooper Lorsung, AmirPouya Hemmasian, and Amir Barati Farimani. Strategies for Pretraining Neural Operators, October 2024a. URL 08473. arXiv:2406.0847) [cs].
Zhanchao Zhou, Tianyi Wu, Zhiyun Jiang, and Zhenzhong Lan. Value Residual Learning For Alleviating Attention Concentration In Transformers, 2024b. URL https: /arxiv.org.//////......./../.... abs/2410 17897. epprin: 2410.17897.
18

---

Published as a conference paper at ICLR 2026
APPENDICES
LLM Acknowledgment
20
Background
20
B.1 Groups
20
B.2 Group Representations
20
C Datasets
21
Licences
C.2 Model Inputs and the Dataloader
21
C,3 Numerical Stability
22
C.4 TG80 Generation Algorithm
23
C.5 Molecular Dynamics Simulations
23
D Architectural Details
24
D.1 Fully Equivariant ATOM
24
D.2 Random-War Walk Positional Encodings
24
D.3 Value-residual Learning...
25
D.4 Delta-prediction
25
Further Experiments
25
E.1
Further Single-task Learning Experiments
25
E,2 Further Multitask Learning Experiments
26
E.3 Monte Carlo Estimation of Quasi-equivariance
27
F
Experimental Details
27
F.1
Software and Hardware Details
................................
27
F2
Training Times and Compute Requirements
27
................................
F3
Inference Times
28
F4
ATOM Hyperparameters
29
F.5
EGNO Hyperparameters and Experimental Details
30
G Propositions and Proofs
31
G.1 Kernel Integral form of Cross-attention.
Trajectory Samples
32
19

---

Published as a conference paper at ICLR 2026
A LLM ACKNOWLEDGMENT
The authors acknowledge the use LLM for grammar corrections and for improving the clarity of the manuscript. The LLM not used for generating original ideas, content, experimental results. All conceptual contributions, analyses, and conclusions presented in this work are entirely from the authors.
B BACKGROUND
This section provides an introduction to the preliminaries of group theory.
B.1 GROUPS
A group (G, o) consists of a non-empty set G and a binary operation o: G x - G satisfying the following axioms:
1. Closure: For all a, G, the result of the operation o b is also in G: aobeC.
2. Identity Element: There exists an element e G such that, for all a E G, ace eoa a.
3. Associativity: For all a, b, E G, (a o b) c a (b c).
4. Inverses: For each G, there exists element
G such that aoa a - oa
In general, not all groups are abelian. That is, the binary operation does not necessarily commute: g h hog. Vg, h E
B.2 GROUP REPRESENTATIONS
group representation homomorphism p: G GV. that assigns matrix each group element E G, realizing as a linear transformation. Representations must preserve the binary operation for all members of the group such that:
pga h) plg)e(h), Vg,
representation p(g) is reducible if be represented the direct sum of other representations:
p(g) p1(g) @ p2(g), Vg E G.
For example, reducible 4 x representation of (U) can be decomposed into two 2 sub- representations:
(g())
p(g)
Ag SU(2),
0
p2(9)
where pi (g) and p2(g) the following irreducible representations of SU(2):
40 0
pi ()
(2(g)
e 10
By contrast, irreducible representations or irreps cannot represented as such direct sum. Formally, they have no non-trivial invariant subspaces W C V such that P(g) c C,W. Vg E G.
Representing inputs as irreps ensures equivariance by constraining each feature to transform pre- dictably under group actions. Given V = θ V, with irreps V., the transformation of an input x e V under g E is:
p(g) x + p:(g)z,
Each component x; transforms independently according to Ps. preserving symmetry. Scalars remain invariant, while vectors rotate according to standard representations. This decomposition prevents the mixing of differently transforming features, ensuring that all subsequent operations, linear or non-linear, respect the group's symmetry, thereby maintaining equivariance throughout the network.
20

---

Published as a conference paper at ICLR 2026
Intuitively, the tensor products capture interactions between features in manner akin multiplication, producing a nigher-dimensionll representation. Crucially, this new representation is reducible, so we may decompose it into irreps:
V
Vi-
is this decomposition that allows the network to project onto individual irreps, achieving feature mixing whilst preserving symmetry constraints.
C DATASETS
We present a visualization of a sample trajectory of uracil from three datasets in Figure 7.
0.2
6/50
0.00
0:3
2.4
1D Uncil
RMD 1? Umeil
TG30 Umeil
Nig--r
Figure 3000 timesteps of uracil trajectory from MD17, RMD17, and TG80.
C.1
LICENCES
Table 5: Dataset sources and licenses. We release TG80 under the MIT license.
Dataset
Source
License
MD17
httpa : //www//www/.o// sgdml. .org/
CC BY 4.0
RMD17
satps://archs wtre.materiaa Lsc wtw.::/sssslad..o.gg..o..o btr.://oddddrr/////////2222222
CC Zero V1.0 Universal
MD22
https: //www wwwwl .org/
BY
TG80
released
MII
C.2 MODEL INPUTS AND THE DATALOADER
Our compound representations follow (Shi al., 2021; Xu et al., 2024). We model hydrogen atoms implicitly and concatenate the position and velocity norms for each node with their respective vectors. Unlike their implementations, avoid explicit graph construction and do not include edge labels describing atomic bond geometries.
We duplicate all frames (4) + { during dataset initialization, producing a five-fold improve- ment in throughput compared to previous dataloaders in Table 6.
Table 6: Mean time (seconds) to produce 10000 batches 100 benchmark runs. Batch size 100, 500 samples, Al 3000, 500 warmup batches.
Aspirin
Ethanol
Naphthalene
Toluene
EGNO
0. 0.060+0.024 0.024+0.015 016
8.056+0.024 0.008
0.039=0.024
ATOM
0.005 000002
0.007 =0.002
0.006.2.002
21

---

Published as a conference paper at ICLR 2026
C.3 NUMERICAL STABILITY
We evaluate the numerical stability of MD17, RMD17, and TG80. MD17 benzene exhibits substantial center-of-mass drift Figure 8a, which is also partially visible in the consistent motion trails shown in Figure 10a. RMD17 exhibits improved stability, with no center-of-mas ma-s drift exceeding 1 x 10 TG80 shows the lowest drift of all datasets, and expectedly includes more molecules with high per-step drift (due to more complex sterically hindered geometries).
ofderr Di MV
Mas Thll
(a) MD17 molecules are largely consistent, except benzene, which exhibits substantial drift.
(b) RMD17 molecules are more numerically stable, supporting their use in future benchmarks.
x10
1.0
High Internal Motion
High Drift Internal Motion
Heplanol
0.8
Per-Stop Internal Motion (A)
0.6
Bipbenyl acmmecene
0.4
0.2
Static
Drifting from Origin
10
10
10
10²
10
10*
Center of Mass Drift (A)
(c) TG80 molecules exhibit lowest centre-of-mass drift among the evaluated MD datasets.
Figure 8: Comparison of numerical stability across MD17, RMD17, and TG80 datasets. Dashed lines denote the mean centre-of-mass drift and per-step motion; datapoints exceeding two standard deviations annotated.
22

---

Published as a conference paper at ICLR 2026
C.4 TG80 GENERATION ALGORITHM
We first recall the definition of Tanimoto T similarity between two bit vectors X, Y as
XXNY]
T(X,)) XUYI
which is identical to the definition of the Jaccard similarity in this case (Rogers & Tanimoto, 1960).
To generate TG80, we randomly shuffled the PubChem dataset, then iterated through all compounds until 40 were found that matched the following criteria:
1. Simplified Molecular-ing input Line-entry System (SMILES) encode valid molecular structure
Randomity sample from seed
2. No more heavy atoms than the corresponding seed molecule
3. Only contain {C, H. N} atoms
4. No more than five oxygen atoms
5. No more than three nitrogen atoms
Pub C hem
6. No disconnected molecular fragments (e.g., salts)
Bcreen compound against seed
7. Tanimoto similarity to at least one seed molecule greater than 0.875, less than 0.925
Figure Construction of TG80 from an initial seed using the PubChem database.
8. Tanimoto similarity to previously selected molecule no more than 0.2
This controlled selection procedure generates structurally analogous subsets around each seed molecule whilst preventing convergence to highly similar molecules across different seed groups.
Only 2,488 of the 173 million in the PubChem library satisfied the filtration criteria above. This low yield largely reflects the cumulative effect of criterion 8: as more molecules are added, it becomes harder find candidates sufficiently dissimilar to prior selections. Given that the average Tanimoto similarity to our seed set was just 0.14 the 0.875 threshold was highly selective. Dataset generation code available at ANONYMIZED.
C.5 MOLECULAR DYNAMICS SIMULATIONS
We present complete overview of the DFT parameters used generate MD17 (Chmiela al., 2017), RMDI7 (Christensen & von Lilienfeld, 2020), MD22 (Chmiela et al., 2023), and TG80,
Table 7: An overview of the methodologies used to generate the MD datasets featured.
DFT Functional
Dispersion Corrections
Basis set
Timestep
Temperature
MD17
PBE
TS
NAO
0.5 fs
500K
RMD17
PBE
None MBD
def2-SVP
0.5 fs
500K
MD22
PBE
NAO
1.0 fs
500K
TG80
PBE
A4
def2-SVP
1.0 fs
300K
23

---

Published as a conference paper at ICLR 2026
D ARCHITECTURAL DETAILS
In this section, we discuss the finer architectural details of ATOM and the architectural adjustments made for our ablations.
D.1 FULLY EQUIVARIANT ATOM
To achieve the full equivariance discussed in Figure 5, we employ a canonicalization network approach, which removes Euclidean gauge before learning and then reinstates it afterwards (Kaba et al., 2023). This preserves equivariance of the whole network, even with the use of non-equivariant architectures the trunk. This provides a controlled comparison between equivariant and non- equivariant ATOM: it enforces equivariance by construction while otherwise keeping the core architecture unchanged.
We first make data translation equivalent by centering
T2
24
μ..
(2)
We then remove rotations by aligning to the second moment
S
E AecReh
(A 22 Xz),
(3)
and choose ej as the principal axis and orthonormalise
e1)e1
63 X 62.
(4)
[2 (e, e1)eıll
We can then form Q = [e1, e2 ea] E (3) and canonicalise
f; = (x: Q,
8. == Q.
(5)
fix the eigenvector sign ambiguity using the chirality pseudoscalar cp = Σ TG X 8/ at the reference time (flip E1 to satisfy e co 2 0, then adjust e2, E3 jointly to keep right- hegtthanndddnes).. Let F be an arbitrary trunk acting in the canonical frame; with per-atom canonical outputs y F({Ø;, 0g N )), we decanonicalise by
y1 5J Q +
(6)
This results in exact
(Kaba et al., 2023) and permits non-equivariant trunks.
D.2 RANDOM- WALK POSITIONAL ENCODINGS
In the multitask case, we add row-normalized random walk positional encoding (RWPE) to equip ATOM and EGNO with multiscale connectivity features, enhancing their ability to distinguish non- isomorphic graphs (Dwivedi et al., 2022; Ma et al., 2023). We first form e-neighborhood graph from our pointclouds as:
G E),
V= E ((i.j) [1 v. 2), (.,,)), e)-
We set = = 1.6, as covalent bonds typically range from 1 Á to 2.0 Á in length (Lobato et al., 2021) and highlight that this construction does not necessitate prior knowledge of the graph structure.
Let A E p denote the adjacency matrix of this graph, and let D = diag (A1) represent its degree matrix. We construct the random walk transition matrix as M D 1 AA then compute matrix powers of M up to maximum walk-length K, defining the self-return probabilities for each node as
P
(M*)
-= = K.
(8)
These probabilities are collected into vectors p; E R and concatenated with the phase space to form (v || Z p) Wn. Here, the input feature space is redefined as Win Vnn @ n @ (pgron d0 RR), and the subsequent equivariant maps modified in kind.
24

---

Published as a conference paper at ICLR 2026
D.3 BALUE-RESIDUAL LEARNING
We employ value-residual learning wherein each transformer block receives the output of the first block via residual connection stabilize training and information flow through the network (Zhou et al., 2024b). Inspired by (Jordan, 2024), we add learned coefficient to weight this residual. Here, denotes the current block's value output, and Ej represents the initial block's value. A learnable parameter a is passed through a sigmoid to obtain the weighting coefficient:
a(a).
(9)
The combined output is then given by:
Ao (1 A) v1. In practice, we lock the first block's À value to 0.5.
(10)
D.4 PREIAAPRED
When delta-prediction is enabled, as in Figure 5, we incorporate the initial positions as residual term, reformulating the model as an operator that learns displacement field rather than predicting absolute positions. We express this as:
xt Project + x.
(11)
Although this approach implemented in both EGNN and EGNO, we found was disabled by default in the codebase of the latter (Garcia Satorras et al., 2021; Xu et al., 2024). Based on empirical results from our ablations, Figure we argue there is sufficient evidence discourage the use of delta- peteapreeiition in neural operator- ppeset molecular dynamics simulations.
E FURTHER EXPERIMENTS
We conduct further experiments on single-task and multitask learning. We consider performance under tail discretization and report results the RMD17 dataset. For multitask learning, report performance under random cluster assignment and S2S metrics for the experiments in Section 4.2.
E. FURTHER KINGLE-TASK LEARNING EXPERIMENTS
MD17 with tail discretization We find the performance of both EGNO and ATOM on MD17 with tail discretization remains similar to the performance under uniform discretization discussed in table 1.
Table 8: EGNO and ATOM with final frame sampling. Upper part: S2S MSE. Lower part: S2T MSE.
Aspirin
Benzene
Ethanol
Malonaldehyde
Naphthalene
Salicylic
Toluene
Uracil
EGNO
9.6610.12
39 09:
4.57
12.9210.00
0.30
0.88
10.99
0.60
ATOM
2.38
33.033
3.62
15.26
0.39
0.83
5.26
a.500
Gap +33.97%
+0.15%
+20885%
-18.06%
+1.62%
+4.75%
+52.13%
+9.28%
EGNO
9.66
39.15g2.0
4.5720.0
12.9220001
0.3920.00
0.88=0.01
10.99
0.60
ATOM
6.38
39.03313.5
2.63
15.21
0.38
0.83
5.2720.79
0.55 10.00
Gap
+33.91%
+0.30%
+ +20.66%
-17.71%
+1 +1.82%
+5.02%
+52.08%
+9.44%
Revised MD17 Dataset We reach performance parity with EGNO on RMD17, shown in Table 9.
Table 9: EGNO and ATOM with final frame sampling. Upper part: S2S MSE. Lower part: S2T MSE.
Azobeuzene
Ethanol
Malonaldehyde
Naphthalene
Paracetamol
Salicylic
Toluene
Uracil
BGNO ATOM
8.9622..0
23.26
40.1 20.05
1.42.0.00
28.0820.0r
1.0629.0
28. 20:0
0.8820. 00
8.88+0.05
23.49
10.29
1.36=0.00
30.12+0.a7
1.03#0.00
28.56+0.04
0.86=0.00
Gap
+0.90%
-0.99%
-0.15%
+3.93%
-7.26%
+3. 1..00%
-0.99%
+ +1.90%
BGNO
8.5110.08
23.6 23.61200.8
40.32:0.08
1.4220.00
10000000
1.07490.01
0.8749..0
ATOM
8.38+0.05
23.90=0.15
40.67=0.17
1.36=0.00
30.03:a.78
1.04#0.00
28.5840.06
0.83=0.00
Gap
+1.47%
-1.27%
0.88%
+-1.9%%
-7.21%
-2.8%
-1.23%
+2.00%
25

---

Published as a conference paper at ICLR 2026
E.2 FURTHER MULTITASK LEARNING EXPERIMENTS
Random-split cross-validation on TG80. For completeness, we report multitask results under compound-level random cross-validation, where compounds are randomly assigned the train, validation, and test sets. Relative to the more challenging out-of-domain (UMAP-based) split in Table EGNO comparatively stronger, nevertheless, ATOM maintains consistent lead across folds, with mean improvements of 24.43% on S2S and 23.93% on S2T.
Table 10: S2S MSE (x10 ?) on TG80 across five random cluster assignments.
Cluster
Cluster 2
Cluster
Cluster
Cluster
EGNO
71.83=0.00
76.92+0.00
68.99
101.27=0.00
83.20#0.00
OOD
ATOM
53.93
62.40
49.37
70.75
66.75+10.00
Gap
+24.92%
+ +8888%
+28.45%
+33..4%
+19.77%
Table 11: S2T MSE (x10-2) on TG80 across five random cluster assignments.
Cluster 1
Cluster
Cluster
Cluster 4
Cluster
EGNO
63.23=0000
64.09=0.00
59.18:00.0
85.87+00.00
69.46:0000
OOD
ATOM
46.0940.00
54.4471..00
42.90.1..0
55.6440.00
59.55
Gap
+27.10%
+ +15.54%
+27.51%
+35.21%
+14.28%
Multitask S2S results TG80 under UMAP cluster cross-validation. The S2S side of the multitask learning results follow closely from their S2T counterparts presented in Section 4.2.
Table 12: S2S MSE (x10-2) on TG80 across five UMAP cluster assignments.
Cluster 1
Cluster 2
Cluster
Cluster 4
Cluster 5
EGNO
51.9840.81
95.86+0.53
142.51=0.5s
155.25+0.67
109.25=0.21
ID
ATOM
15.4941.04
26.55-2213
28.7442..0
29.8142.77
2..33
Gap (%)
70.20%
72.30%
79.83%
80.80%
75.90%
EGNO
52.90++072
114.14=13.21
149.99=0..
163.47+11.0
112.36
EGNN-S
52.3940.44
12512.07
149.41
663.541865.22
11 111.08
OOD
EGNN-R
52.08:+0.79
108.89 1.66
148.67==077
163.27
109.9 199.94+0.3.
ATOM
41.97
127.95
74.53
80.95
58. 1.
Gap (%)
19.41%
-17.50%
49.87%
50.42%
47.01%
26

---

Published as a conference paper at ICLR 2026
Multitask S2S Versus Single Task We evaluate mul- titask ATOM model trained on all available trajectories, pooling both ID and OOD clusters, against single-task ATOM trained separately on each compound. The mul- titasking model achieves losses comparable or lower than those of the single-task baselines, despite being trained with the same compute resources.
Table 13: MSE (x10-2) on full-dataset TG80 ATOM and single-task ATOM. S2S upper, S2T lower.
Formic Acid
Single-task ATOM All-data ATOM
26.40
22.39
E.3 MONTE CARLO ESTIMATION OF EQASS-EQUIVARIANE
Gap (%)
15. 19%
Single-task ATOM
18.10
All-data ATOM
18.72
For both the pretrained ATOM and the non-equivariant lifting variant, we estimate the quantity Section 3.1 by Monte Carlo, drawing N random timesteps 2n and, for each, R random rotations In.r E G. We approximate
Gap (%)
33.43%
N
R
(f(o(on)))xn)) - (f(n))) Σ
E
R
n=1
We report our estimates of for various ATOM models trained on MD17 single task learning in Table 18. In all cases, ATOM shows a substantially lower equivariance defect, supporting our claim that quasi-equivariant design achieves a kind of middle middle-ground ground in the trade-off between expressiveness and strict equivariance.
Table 14: Estimates of Quasi-equivariince via Monte Carlo over rotations and 10 timesteps with 2SD intervals.
Aspirin
Ethanol
Malonaldehydee
Naphthalene
Salicylic
Tolucne
Uracil
Non equivariant
120 a118.7
102.50:a4..0
37.97
37.
57822,9 11222.95=4a77222
32
ATOM
28. 3 31=25.. 10
11. 19..94
9.99=5.82
17.09
25. 16-14421
27 0733.07733
28. 18.722=10.20
F EXPERIMENTAL DETAILS
F.1 SOFTWARE AND HARDWARE DETAILS
All experiments were conducted using Python 32, NumPy 2.2.1 (Harris al., 2020), PyTorch 2.5.1 (Paszke et al., 2019), e3nn 0.5.6 (Geiger & Smidt, 2022) and PyTorch Optimizer 3.5.0 (Kim, 2021). We use RDKit 2024.9.6 (Landrum al., 2025) and PubChemPy 1.0.4 construct TG80. All single-task training was performed on an NVIDIA RTX 5080 (16 GB) with CUDA 12.4, running on Ubuntu 24.04. We use Ase 3.26.0 (Larsen al., 2017) and MACE-Torde 0.3. 14 (Kováes al., 2025) in the experiments of Appendix F.3.
F.2 TRAINING TIMES AND COMPUTE REQUIREMENTS
Single-task training time We roughly wall-clock normalised our ATOM and EGNO parameter counts, resulting in respective learnable parameter counts of 754 468 and 335 770.
Table 15: Compute cost of single-task training on all MD17 molecules over 1000 epochs. Both ATOM (335 770 params) and EGNO (754 468 params) under torch. comp ill on Titan V.
Model
Metric
Azobenzenc
Eihanol
Moloealdehyde
Naphibalene
Paracetamol
Salicylic
Toluene
Uracil
=0.02
4.35
EGNO
Total PLOPSS (x110 By
3681 .24
1257. 70
3282.00 00
4513.37
8114.44
4641.50)
3539.92 92
3915.98
Epochs/mmin
244.48
276.27
274.22
199.41
110.91
193.90
254.24
229 83
5.83
ATOM
Total OPS (3
82286 49
5212.50 53
52133 50
5271.19
5207. 84
52661 76
52244.9 91
5247.33
Epochs/min
172.20
172.66
172.63
170.74
169.88
170.98
172.25
171.52
FLOPS Reduction
85%
16.79%
+31.71%
60.60
-3100%
27

---

Published as a conference paper at ICLR 2026
Multitask training time In multitask training TG80, our upsized ATOM model contained 3.57 840 parameters, compared to 335 770 for EGNO. Despite this, ATOM only trained between 5% and 30% slower than EGNO. This is perhaps unsurprising given the much higher of the transformer architecture upon which ATOM is based.
Table 16: Compute cost of single task training on five TG80 molecules over 1000 epochs. Both ATOM (335 770 params) and EGNO (754 468 params) are under torch. compice on a Titan V.
Model
Fold1
Fold2
Fold3
Fold4
Fold5
Time (mins)
11.21
8.5641.06
9.04
9.02
8.98
EGNO
Total FLOPS (x 10
8645.66
7703.33
8136.98
8378.06
8084.98
Epochs/min
104.10
116.83
110.61
107.42
111.32
Time (mins)
10. 1620.49
10.55:0.02
I 1..62
12.38
10.39
ATOM
Total FLOPS (x (x))
9140.57
9497.79
10455.34
11141.73
9355.37
Epochs/min
98.46
94.76
86.08
80.78
96.20
Total FLOPS Reduction (%)
5.72%
23.29%
28.49%
32.99%
15.71%
F.3 INFERENCE TIMES
We compare inference times on MD17 and MD22 across ATOM, the pretrained machine learning interaction potential MACE-OFF24 (Medium) (Kovács et al., 2025), and the classical Lennard-Jones potential (Larsen et al., 2017; Schwerdtfeger & Wales, 2024) and the molecular forcefield. We report inference times in seconds with 2SD) intervals. We exclude AMBER results on Uracil as we were unable to run simulations for this molccule (Case et al., 2023).
Table 17: Seconds to produce timestep AT 3000 of each MD17 trajectory at float 32 precision.
Aspirin
Ethunol
Makonaldehydee
Nophtbalene
Salicylic
Toluene
Unaiil
MACE-OFF
42.803 42.805+a..0 +2
39.676 =0 560
39.782 too
39.916 mo 125
39.620 m30
40.271 +bs
40. 10B thusce
AT' 3000
AMBER
Lemmand-Jnnes
37,745 1ussa
2.400-tm30
38.512
1.604
-n
-:1 204
36.806 Eu.
1.520 trom
2.174 m0
38.512 13:
37.621 1114
1.981
-he
36. 061
1.906
40e
1ic.ee
1,804 th.
ATOM
0.849 Ln.so
0.259
0.714
0.367
0.450
0,341
0.373 La.ons
MACE-OFF
142.500
140.372 +..24
011.508
121 th
121.607
Lemmand-Jones
7.797 11:30
5.533 22 425
5.167 2u.z
7.453 200370
6.721
6.477
5.921 La
ATOM
2.554 70.05
1.002 -m 184
1.634 200000
1.099 -10
1.116
03.003
0.064
40 155
0.900 1m.
Table 18: Seconds to produce timestep AT 3000 of each MD22 trajectory at float 32 precision.
Ac-Ala3-NHME
DHA
Stachyose
MACE-OF
44.737 14.61
42.042 70.
47.913 E0.546
Lennard-Jones
4.700 1.2212
3.790
0.246
6.613 1.0.177
ATOM
1.302 +
0.914 0.055
3.075 t0.080
28

---

Published as a conference paper at ICLR 2026
F.4 ATOM HYPERPARAMETERS
We employ the same dataset splitting and discretization parameters reported in Xu et al. (2024) for the MD17. We set the batch size 192, the AdamW-AMSGrad optimizer (Loshchilov & Hutter, 2017) with an of x 10-10 to avoid instability associated with the small gradients produced by zero-initialised weight matrices in early training (Jordan al., 2025). During multitask training, we reduce the number of epochs to 250 and employ the Muon optimizer (Jordan et al., 2024; Kim, 2021). We present a complete overview of our hyperparameters in Table 19.
Table 19: Hyperparameters for ATOM. MD17 hyperparameters are shared across all molecules unless otherwise noted.
Modulc
MD17, RMD17, TG80
TG80 Multitask
Training
Batch size Epochs
192
192
1000
250
Max grad norm
1.0
1.0
Label noise
0.1
0.1
At Timesteps P Train/Val/Test
3000
10 000
8
8
(500, 000, 3 300)
56500, 13 000, 13 000) 8
RWPE length
Optimiser
Optimiser type Learning rale $1, S8 Weight decay
AdamW-An W-AMSGrddd 10-3
Muon
1 x 10-3 (0.9, 0.999)
0.9, 0.999
1 x 10-6
1 X 10-5
1 x 10-10 10
1 X 10
Model
Embedding dim
128
256
No. layers
6
No. attention heads
No. output heads
8
A Attention dropout
0.2
0.2
RoPE frequency
1000
1000
MLP layers
2 SwiGLU
MLP activation
SwiGLU
MLP dropout
0.0
0.0
Norm type
RMS norm True
RMS norm
Learnable value residuals
True
29

---

Published as a conference paper at ICLR 2026
F.5 EGNO HYPERPARAMETERS AND EXPERIMENTAL DETAILS
We generated the EGNO results reported in Table with the same discretization parameters and hyperparameters used in their experiments. We reduce the number of epochs from 10 000 to 2 500, use a batch size of 192 with the Adam W-AMSGrad optimizer (Loshchilov & Hutter, 2017), and select the best validation loss epoch for testing. In the multitask case, we further reduce the number of epochs to 250 and employ the Muon optimizer (Jordan et al., 2024; Kim, 2021). Complete hyperparameters are displayed in Table 20.
Table 20: Hyperparameter values for EGNO across each benchmark dataset.
Module
MD17, RMD17, TG80
TG80 Multitask
Training
Batch size Epochs
192
192
2500
250 Uncapped 0.1 10 000 8
Max grad norm Label noise a At Timesteps P Train/Val/Test
Uncapped 0.1
3000
8
(500, 3 300030000))) 8
(6 500, 13 000, 13 000) 8
RWPE length
Optimiser
Optimiser type Learning rate
Adam W-AMSGrdd 1 x 10 (0.9, 0.999)
Muon
1 x 10 (0.9, 0.999) x 10-5 10-5
$1, 32 Weight decay
1 x 10-5
1 10-10
Scheduler
Scheduler type Step size 7
StepLR 2500
StepL 2500 0.5
0.5
Model
Embedding dim
64
64
No. EGNO layers
5
5
Temporal convolution activation
LeakyRELU
LeakyRELU
MLP layers
2
2
MLP activation
SiLU
SiLUU
MILP dropout
0
0
Time embedding dim
32
32
Fourier modes
2
2
30

---

Published as a conference paper at ICLR 2026
G PROPOSITIONS AND PROOFS
G.1 KERNEL INTEGRAL FORM OF CROSS-ATTENTION
Proposition G.1. The cross-attention is equivalent a kernel integral operator, i.e.. (K)
ff 6z x));(z)ddp (x), where Ki denotes the kernel induced by softmax function, U/ (;)x) denotes the values as a function of x, and JuN denotes the empirical measure supported {x, {x
Proof of Proposition G.I. Following (Gao et al., 2024) we may view our attention as kernel integral transform by considering x; as being sampled from the continuum domain C. R33 for which we define the empirical measure with support on {x,} Q::
N
N Σ ((x)) 1=1
| (xx)duN(x)= 8xzy
MN (N(x)
= N
(12)
N
where 5 is the Dirac delta function "selecting" the values X,. Given T-RoPE-rotated query and key maps go (z) (o,)). kz(x,) ko, (x,) we form the data-dependent kernel for feature F:
exp ((()), (x,)) vdn)
Kayi (z, xj)
(13)
exp( (((Z)), (x))) / / dd)) dxx((x)
Thus, for any F E FF we may represent our cross-attention as the kernel integral operator:
(Ko,v,)(2z= = / (e, (x) (x),
1,
(14)
which is row-stochastic under the measure in Equation (12).
We remark that the kernel fails to satisfy global Lipschitz continuity (Delattre et al., 2023), unlike FNO (Li et al., 2021), 1), and certain generalization theorems fail result (Le & Dik, 2024).
31

---

Published as a conference paper at ICLR 2026
H TRAJECTORY SAMPLES
Aspida
Arehoraxar
0.2
Estooo
Alaalbetydd
Napltokke
Naphthaleza
Sallettic
Saligyte
Tokess
Uneil
Tolussc
NGP-O
(0-9
CS24
O.2-9
(a) MD17 trajectories
(b) RMD17 trajectories
Figure 10: 3000 steps MD trajectories from the MD17 and RMD17 datasets.
32

---

Published as a conference paper at ICLR 2026
-02
1.2Diibieeetllm:
1.3.undiam:
13-Cycloanatiime
6.0
Acctalishyde
10
6.0
80
-22
Beszaloetdde
40
Boncylmiie
Hiptaayl
80
0.0
-19
19
180
-22
-22
-011
Conaaic
Cyclalssnol
(12-4
N12-5
33

---

Published as a conference paper at ICLR 2026
14
-1
122
-13
Cydlppepan:
Bbylamin:
Edylen:
90
022
0.0
-11
-01
Fonnalietyde
19
422
-22
413
48
Reptaol
Iopropanoo
Makedulddtyde
1.0
-2.
.01
-10
14
Osdicacié
pCral
Poppllane
SalcsficaoM2
Sweinkead
(12-4)
12--9
34

---

Published as a conference paper at ICLR 2026
108
15
16
100
05
-1.57
233
Repea22
Repons
822
1.2
Unell
Figure 11: 3000-step MD trajectories from TG80. Molecules generated by our dataset expansion algorithm are named according to their seed molecule and the order of their selection.
35


Reviews:
[{"review_id": "Zox96iTtD2", "reviewer": "Reviewer_h8zv", "summary": "This paper introduces ATOM, a neural operator for molecular dynamics that introduces a quasi-equivariant architecture for better flexibility and simulation efficiency. Molecular systems are modeled as fully connected point clouds. It employs a equivariant lifting layer to embed atomic positions and velocities into a symmetry-aware latent space and a temporal attention mechanism to enable parallel decoding of multiple future states. ATOM demonstrates better performance than state-of-the-art approaches on standard MD benchmarks and shows that it can transfer to unseen molecules better than prior approaches. The authors also present TG80, a large-scale benchmark dataset with simulation data for a diverse set of molecules.", "questions": "Please address comments in weakness.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 1, "contribution": 2, "strengths": "* This is an important and timely problem\n\n* For a multitask set up, the authors demonstrate better performance that state-of-art baselines\n\n* Parallel decoding can be more efficient\n\n* An MD dataset, TG80, is produced, which intended to be chemically diverse and numerically stable to help in multitask pretraining and benchmarking tasks", "weaknesses": "* The choice and implications of modeling atoms as a point cloud is not discussed \n\n* As a non-expert, I found few insightful discussions or learnings in the paper. It would be good if the authors discussed:\n(1) the explicit novelty of each part of their architecture in comparison to prior works; (2) the importance of generality as opposed to specialized models; (3) the insights behind structure of the attention-based mechanism.\n\n* It is not clear what \"heterogeneous attention\" means and its importance.\n\n* There is no evaluation of the efficiency benefits of the proposed mechanism \n\n* Figure 2 does not add much value to the paper", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "6dd4vq0K8K", "reviewer": "Reviewer_L5rw", "summary": "The paper introduces ATOM, a pretrained neural operator designed to improve molecular dynamics (MD) modeling by incorporating approximate geometric equivariance. The method proposes a heterogeneous temporal attention mechanism along with a temporal RoPE (T-RoPE) adapation, and further enables large-scale pretraining and generalization across different MD tasks. Experimental results demonstrate that ATOM achieves competitive performance on multiple benchmarks compared to strong baselines such as EGNO.", "questions": "*  In the single-task setting, what is the key difference between ATOM and the main baseline EGNO? Can ATOM degenerate into EGNO if some critical components are removed? If so, such a description would clarify the key improvements and help highlight the method’s contribution.\n\n* Does multi-task training improve the performance of each individual task? The paper does not include an ablation study that isolates the effect of multi-task learning itself.", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "* The paper addresses a critical problem, that is, improving neural operators for molecular dynamics, and further enables pretraining on large-scale datasets, which is an important step toward generalizable molecular representation learning.\n\n* This paper constructs a new large-scale molecular dynamics dataset, TG80, which not only facilitates the pretraining of ATOM itself but also provides valuable resources for future research in this area.\n\n* The empirical results are generally strong and show that ATOM performs well across multiple benchmarks, indicating the effectiveness of the proposed approach.", "weaknesses": "* **Unclear Architecture Design.** The paper repeatedly employs the term **quasi-equivariant** without providing a clear mathematical definition. Moreover, while the ablation studies compare various architectural variants, the main text and appendix lack a detailed, mathematically grounded description of the original model design. The absence of released code further limits reproducibility. The authors should clearly define quasi-equivariance, explain its advantages, and elaborate on how it is concretely maintained in the model.\n\n* **Limited Discussion of Related Works.** The problem studied in the paper can also be framed as a **Time-Coarsened Dynamics** problem, yet many relevant works in this area are not sufficiently discussed [A, B, C]. The authors are encouraged to analyze the connections and differences between ATOM and prior works, including distinctions in problem formulation, modeling strategy, and evaluation metrics.\n\n[A] Klein, Leon, et al. \"Timewarp: Transferable acceleration of molecular dynamics by learning time-coarsened dynamics.\" Advances in Neural Information Processing Systems 36 (2023): 52863-52883.\n[B] Hsu, Tim, et al. \"Score dynamics: Scaling molecular dynamics with picoseconds time steps via conditional diffusion model.\" Journal of Chemical Theory and Computation 20.6 (2024): 2335-2348.\n[C] Yu, Ziyang, Wenbing Huang, and Yang Liu. \"UniSim: A Unified Simulator for Time-Coarsened Dynamics of Biomolecules.\" Forty-second International Conference on Machine Learning.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "sZ02l89JVt", "reviewer": "Reviewer_hg7v", "summary": "This work presents the atomistic transformer operator for molecules (ATOM), which preserves quasi-equivariance and enables parallel decoding of molecule states across multiple timesteps. The model leverages E(3)-equivariant linear layers to produce symmetry-aware features, and proposes temporal RoPE (T-RoPE) to build the heterogeneous temporal attention mechanism. For both single-task experiments on MD17, MD22 and multitask experiments on the curated TG80 dataset, ATOM consistently outperforms the selected baselines, showing an out-of-distribution generalization across unseen molecules.", "questions": "1. Whether ATOM offers efficiency improvements over sequential generative models remains to be verified. The authors are encouraged to provide a comparison of inference time between ATOM and conventional sequential generation models when generating trajectories of equal wall-clock duration.\n2. Although ATOM is theoretically capable of predicting future states at different time intervals through T-RoPE, the maximum predictable time interval is pre-defined as $\\Delta T$. Therefore, when generating a long trajectory, using ATOM appears equivalent to employing a sequential generative model with a fixed time step of $\\Delta T$. Please provide a reasonable explanation for this or correct me if I understand wrong.\n3. In line 316, the experimental setup reports $\\Delta T=3000$. Is this value expressed relative to the frame spacing of trajectories in the training datasets? What wall-clock time does it correspond to?\n4. To my knowledge, for molecular systems with fewer than 20 heavy atoms, MD simulations using empirical force fields can already achieve high efficiency. Please clarify the advantages of ATOM over classical MD methods, preferably supported by additional experimental results.\n5. In Table 3, EGNO appears to exhibit a much smaller performance drop from in-distribution to out-of-distribution settings compared to ATOM. Please provide an explanation for this observation.\n6. ATOM predicts future states at different time intervals given the current state, which typically relies on the assumption that the trajectory is uniquely determined by the initial state. However, in broader application scenarios, biomolecules generally interact in aqueous solution, where thermal fluctuations must be considered. Thus, trajectories should be represented as distributions rather than being uniquely determined. ATOM appears not directly applicable to such settings. How do the authors address this limitation?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- The paper is written with a clear logical flow, facilitating reader comprehension. Each component of the model architecture is well justified and supported by comprehensive ablation studies.\n- The proposed T-RoPE method ingeniously enables attention to achieve translation equivariance along the temporal dimension, thereby supporting parallel prediction of future molecular states at various timesteps.\n- The newly introduced TG80 dataset represents a valuable extension over datasets such as MD17 and MD22 in terms of both scale and trajectory accuracy, contributing to the advancement of the field.\n- ATOM demonstrates out-of-distribution generalization on the TG80 dataset, which is crucial for practical applications.", "weaknesses": "- The authors claim that ATOM enables parallel computation through the T-RoPE technique, thereby overcoming the limitation of sequential sampling. However, this claim lacks experimental evidence. For instance, it would be helpful to compare the inference time of ATOM with that of sequential generative models when generating trajectories of equal wall-clock duration.\n- Given the molecular state at time $t$, the model outputs future states at different time intervals based on the assumption that the trajectory is uniquely determined by the initial state. However, in realistic scenarios like protein–ligand interactions in aqueous solution, thermal fluctuations must be taken into account, which are typically modeled using Langevin dynamics. In such cases, the trajectory is not uniquely determined but rather represented as a probabilistic distribution. Therefore, the training and evaluation of ATOM are not fully aligned with real physical settings.\n- As the authors mentioned in the Limitations section, ATOM has not yet been validated on larger molecular systems. However, for small-molecule systems with fewer than 20 heavy atoms, MD simulations using empirical force fields are already highly efficient. The authors should therefore further clarify the advantages of ATOM over classical MD simulations based on empirical force fields.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)