ATOM: A Pretrained Neural Operator for Multitask Molecular Dynamics
Luke Thompson, Davy Guan, Slade Matthews, Dai Shi, Junbin Gao, Andi Han
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.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Cited by top-tier papers1
Ask how each one uses itBuilds on26
- Fourier Neural Operator for Parametric Partial Differential EquationsZongyi Li, Nikola Borislavov Kovachki, Kamyar Azizzadenesheli, Burigede Liu et al.ICLR 2021 · 3,911 citations
- How Attentive are Graph Attention Networks?Shaked Brody, Uri Alon, Eran YahavICLR 2022 · 1,717 citations
- MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force FieldsIlyes Batatia, Dávid Péter Kovács, Gregor N. C. Simm, Christoph Ortner et al.NeurIPS 2022 · 1,448 citations
- GeoDiff: A Geometric Diffusion Model for Molecular Conformation GenerationMinkai Xu, Lantao Yu, Yang Song, Chence Shi et al.ICLR 2022 · 695 citations
- GemNet: Universal Directional Graph Neural Networks for MoleculesJohannes Gasteiger, Florian Becker, Stephan GünnemannNeurIPS 2021 · 665 citations
Related papers
- Equivariant Transformers for Neural Network based Molecular PotentialsPhilipp Thölke, Gianni De FabritiisICLR 2022 · 277 citations
- May the Force be with You: Unified Force-Centric Pre-Training for 3D Molecular ConformationsRui Feng, Qi Zhu, Huan Tran, Binghong Chen et al.NeurIPS 2023 · 16 citations
- UniSim: A Unified Simulator for Time-Coarsened Dynamics of BiomoleculesZiyang Yu, Wenbing Huang, Yang LiuICML 2025
- Conformal Prediction Sets for Graph Neural NetworksSoroush H. Zargarbashi, Simone Antonelli, Aleksandar BojchevskiICML 2023 · 49 citations
- DiffMD: A Geometric Diffusion Model for Molecular Dynamics SimulationsFang Wu, Stan Z. LiAAAI 2023 · 48 citations
