Improving Molecular Force Fields with Minimal Temporal Information

Ali Mollahosseini; Mohammed Haroon Dupty; Wee Sun Lee

physics.chem-phApr 14, 2026

Improving Molecular Force Fields with Minimal Temporal Information

Authors: Ali Mollahosseini, Mohammed Haroon Dupty, Wee Sun Lee

Organizations: School of Electrical and Computer Engineering, University of Tehran, Tehran, Iran · School of Computing, National University of Singapore, Singapore

Abstract

Accurate prediction of energy and forces for 3D molecular systems is one of fundamental challenges at the core of AI for Science applications. Many powerful and data-efficient neural networks predict molecular energies and forces from single atomic configurations. However, one crucial aspect of the data generation process is rarely considered while learning these models i.e. Molecular Dynamics (MD) simulation. MD simulations generate time-ordered trajectories of atomic positions that fluctuate in energy and explore regions of the potential energy surface (e.g., under standard NVE/NVT ensembles), rather than being constructed to steadily lower the potential energy toward a minimum as in geometry relaxations. This work explores a novel way to leverage MD data, when available, to improve the performance of such predictors. We introduce a novel training strategy called FRAMES, that use an auxiliary loss function for exploiting the temporal relationships within MD trajectories. Counter-intuitively, on two atomistic benchmarks and a synthetic system we observe that minimal temporal information, captured by pairs of just two consecutive frames, is often sufficient to obtain the best performance, while adding longer trajectory sequences can introduce redundancy and degrade performance. On the widely used MD17 and ISO17 benchmarks, FRAMES significantly outperforms its Equiformer baseline, achieving highly competitive results in both energy and force accuracy. Our work not only presents a novel training strategy which improves the accuracy of the model, but also provides evidence that for distilling physical priors of atomic systems, more temporal data is not always better.

Figures & tables

Figure 1: An overview of our proposed training and inference framework. On the left, the shared GNN body processes a history of frames ( $S_{t},S_{t-1}$ , etc.) to produce latent embeddings. The primary Output Head uses these embeddings to predict energies and forces for the entire window, supervised by a primary loss. Concurrently, the embeddings are concatenated $\bigoplus$ and fed to an Auxiliary Head, which is trained with an auxiliary loss to predict the displacement to the next frame $\Delta\mathbf{r}_{t}$ . On the right, at test time, the auxiliary head is detached. The model operates as a simple, static predictor, taking a single frame $S_{t}$ as input to efficiently predict its corresponding energy and forces.

Figure 2: Mean squared error (MSE) of the linear predictor on the spring–mass toy system as a function of the history length $T$ (x-axis). The plot highlights the large performance improvement when moving from $T{=}1$ to $T{=}2$ , followed by degradation once redundant temporal information ( $T{\geq}3$ ) is included. A 95 $\%$ confidence interval is shown.

	Aspirin		Benzene		Ethanol		Malonaldehyde		Naphthalene		Salicylic acid		Toluene		Uracil
Model	energy	forces	energy	forces	energy	forces	energy	forces	energy	forces	energy	forces	energy	forces	energy	forces
SchNet ( Schütt et al., 2017 )	16.0	58.5	3.5	13.4	3.5	16.9	5.6	28.6	6.9	25.2	8.7	36.9	5.2	24.7	6.1	24.3
DimeNet ( Gasteiger et al., 2020 )	8.8	21.6	3.4	8.1	2.8	10.0	4.5	16.6	5.3	9.3	5.8	16.2	4.4	9.4	5.0	13.1
PaiNN ( Schütt et al., 2021 )	6.9	14.7	-	-	2.7	9.7	3.9	13.8	5.0	3.3	4.9	8.5	4.1	4.1	4.5	6.0
TorchMD-NET ( Thölke and Fabritiis, 2022 )	5.3	11.0	2.5	8.5	2.3	4.7	3.3	7.3	3.7	2.6	4.0	5.6	3.2	2.9	4.1	4.1
NequIP ( $L_{max}=3$ ) ( Batzner et al., 2022 )	5.7	8.0	-	-	2.2	3.1	3.3	5.6	4.9	1.7	4.6	3.9	4.0	2.0	4.5	3.3

Table 1: Mean absolute error results on the MD17 testing set. Energy and force are in units of meV and meV/Å, respectively. This table compares several baseline models against our Equiformer-based approach, which is tested using both two and three frames of temporal context to investigate the effects of data redundancy.

	Aspirin		Benzene		Ethanol		Malonaldehyde		Naphthalene		Salicylic acid		Toluene		Uracil
Model (T=2)	energy	forces	energy	forces	energy	forces	energy	forces	energy	forces	energy	forces	energy	forces	energy	forces
Equiformer (T=1, Baseline)	5.3	7.2	2.2	6.6	2.2	3.1	3.3	5.8	3.7	2.1	4.5	4.1	3.8	2.1	4.3	3.8
FRAMES (T=2) with Aux: $E_{t+1},F_{t+1}$	5.3	7.2	2.3	6.1	2.2	3.2	3.2	5.6	3.6	2.2	4.4	4.2	3.8	1.9	4.1	3.6
FRAMES (T=2) with Aux: $\Delta\mathbf{r}_{t}$	5.2	7.0	2.4	6.3	2.2	3.2	3.3	5.6	3.6	2.2	4.3	4.1	3.6	2.0	4.2	3.5

Table 2: Ablation study on the choice of auxiliary loss for the

T=2

model. We compare our proposed method, which uses a loss on atomic displacements (Aux:

E_{t+1},F_{t+1}

) against our proposed method of predicting atomic displacements (Aux:

\Delta\mathbf{r}_{t}

). Results are mean absolute error (MAE) in meV for energy and meV/Å for forces.

	Within Distribution		Outside Distribution
	Energy	Forces	Energy	Forces
Equiformer (Baseline)	0.13228	0.0093	0.1460	0.0174
FRAMES $(T=2)$	0.00569	0.0053	0.0248	0.0154
FRAMES $(T=3)$	0.07009	0.0101	0.0639	0.0187

Table 3: Mean Absolute Error on the ISO17 test sets. Our FRAMES (T=2) model is compared against the baseline Equiformer and a T=3 model. Within Distribution tests generalization to new conformations of known molecules, while Outside Distribution tests generalization to entirely new isomers.

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Hyper-parameter	Aspirin	Benzene	Ethanol	Malonaldehyde	Naphthalene	Salicylic acid	Toluene	Uracil
Energy coefficient $\lambda_{E}$	1	1	1	1	2	1	1	1
Force coefficient $\lambda_{F}$	80	80	80	100	20	80	80	20
FRAMES coefficient $\lambda_{aux}$	1	0.25	0.25	0.25	1	1	1	0.25

Appendix

Table 4: Hyperparameters used for training our ‘FRAMES‘ models on the MD17 dataset. We report the coefficients for the primary loss (

\lambda_{E},\lambda_{F}

) and the initial value for the auxiliary loss (

\lambda_{aux}

Machine Learning Interatomic Potentials

Machine learning interatomic potentials predict atomic energies and forces for fast simulations of molecules and materials.

See all 91 papers

Jul 31, 2026

Implicit Machine Learning Force Fields Accelerate Molecular Dynamics Simulations

We introduce implicit machine learning force fields (I-MLFFs), which replace explicit stacks of neural network layers with self-consistent fixed-point equations. In molecular simulations, this formulation enables intermediate representations to be reused across successive timesteps, thereby warm-starting force evaluation. The resulting models effectively combine the computational footprint of a shallow, single-layer MLFF with the representational capacity and accuracy of a deep neural network. Our approach unlocks architecture-agnostic efficiency gains that are inaccessible when force prediction and trajectory integration are considered separately. We demonstrate this across three major classes of graph neural networks: invariant, equivariant Cartesian tensor, and SO(3)-equivariant spherical-tensor architectures. Each yields a two- to five-fold reduction in compute and memory footprint. Crucially, these gains are achieved while retaining full atomistic resolution and the original integration timestep, avoiding spatial or temporal coarse graining. Our contribution therefore advances the scaling frontier of quantum-mechanically faithful molecular simulation, enabling longer trajectories and larger atomistic systems within fixed GPU memory and compute budgets, and thereby opening access to new insights across biomolecular and material systems.

Johannes Maeß, Leon Werner, J. Thorben Frank +5

BIFOLD, Berlin, Germany. · Machine Learning Group, TU Berlin, Berlin, Germany. · Department of Artificial Intelligence, Korea University, Seoul, South Korea. +1

Implicit Neural Representations Machine Learning Interatomic Potentials Molecular Dynamics Simulation

Ask PDF

Jun 16, 2026

ASTEROID: A Spatiotemporal Information Transformer for Forecasting Multi-Step Time Series of Molecular Dynamics

Molecular dynamics (MD) simulation is computationally demanding, particularly for large-scale systems requiring long-term analysis. Accurate forecast of the outcomes of a MD simulation is not only an attractive scientific challenge but also has substantial practical value. In this work, we developed a data-driven framework, termed ASTEROID (Advanced Spatiotemporal TransformER fOr Inferring Dynamics), that can directly predict multi-step atomic coordinates, avoiding conventional iterative integration. For this purpose, our ASTEROID reformulates MD trajectories as high-dimensional spatiotemporal sequences and integrates the Spatiotemporal Information (STI) Transformation equation into a Transformer architecture. The core innovation of ASTEROID lies in its ability to model multiscale spatiotemporal dependencies. In particular, for spatial dependencies, a local-global self-attention mechanism captures both short- and long-range interactions. For temporal dependencies, an encoder-decoder structure integrates global context with autoregressive forecasting. ASTEROID was evaluated on several quantum-mechanics derived molecular datasets. Our results indicate that ASTEROID achieved not only a higher level of accuracy in multi-step prediction than existing methods on various benchmarks, but also significantly reduced computational cost of conventional MD simulation. Moreover, the model supports iterative multi-step forecasting over an extended time scale. This work establishes a robust and generalizable data-driven paradigm for accelerating MD simulations.

Kexin Wu, Luonan Chen, Renxiao Wang

Department of Medicinal Chemistry, School of Pharmaceutical Sciences, Fudan University, 826 Zhangheng Road, Shanghai 201203, People's Republic of China · School of Mathematical Sciences and School of AI, Shanghai Jiao Tong University, Shanghai 200240, People's Republic of China

Spatiotemporal Forecasting Molecular Dynamics Simulation Trajectory Prediction Time Series Transformers Multiscale Time-Series Modeling

Ask PDF

Apr 5, 2026

Align Your Structures: Generating Trajectories with Structure Pretraining for Molecular Dynamics

Generating molecular dynamics (MD) trajectories using deep generative models has attracted increasing attention, yet remains inherently challenging due to the limited availability of MD data and the complexities involved in modeling high-dimensional MD distributions. To overcome these challenges, we propose a novel framework that leverages structure pretraining for MD trajectory generation. Specifically, we first train a diffusion-based structure generation model on a large-scale conformer dataset, on top of which we introduce an interpolator module trained on MD trajectory data, designed to enforce temporal consistency among generated structures. Our approach effectively harnesses abundant structural data to mitigate the scarcity of MD trajectory data and effectively decomposes the intricate MD modeling task into two manageable subproblems: structural generation and temporal alignment. We comprehensively evaluate our method on the QM9 and DRUGS small-molecule datasets across unconditional generation, forward simulation, and interpolation tasks, and further extend our framework and analysis to tetrapeptide and protein monomer systems. Experimental results confirm that our approach excels in generating chemically realistic MD trajectories, as evidenced by remarkable improvements of accuracy in geometric, dynamical, and energetic measurements.

Aniketh Iyengar, Jiaqi Han, Pengwei Sun +3

Stanford University · Lambda, Inc

Molecular Dynamics Simulation 3D Molecular Generation Diffusion Models

Ask PDF

Abstract

Machine Learning Interatomic Potentials

Explore similar work

Implicit Machine Learning Force Fields Accelerate Molecular Dynamics Simulations

ASTEROID: A Spatiotemporal Information Transformer for Forecasting Multi-Step Time Series of Molecular Dynamics

Align Your Structures: Generating Trajectories with Structure Pretraining for Molecular Dynamics