Teaching Molecular Dynamics to a Non-Autoregressive Ionic Transport Predictor

TL;DR

This paper introduces a non-autoregressive learning framework for predicting ionic transport properties that learns dynamics efficiently without sequential inference, achieving significant speedup and accuracy improvements.

Contribution

It proposes a novel auxiliary modality learning approach that leverages atomic trajectories during training but not inference, enhancing prediction speed and accuracy.

Findings

Over 200 times speedup compared to autoregressive models.

Substantially reduced prediction error relative to non-autoregressive benchmarks.

Effective learning of dynamics without sequential inference.

Abstract

Unlike most static material properties widely studied in the machine learning literature, ionic transport properties are inherently dynamic, making their fast and accurate prediction from static atomic structures challenging. The current standard approach, molecular dynamics (MD) simulations, suffers from prohibitively high computational cost. Recent autoregressive learning-based MD acceleration methods requiring sequential inference remain slow and prone to error accumulation; in contrast, existing non-autoregressive material property prediction models are less accurate because they fail to exploit dynamics. Moreover, existing methods typically benefit from datasets either with or without atomic trajectories, but not both. To overcome these limitations, we propose a non-autoregressive learning framework based on auxiliary modality learning, which treats atomic trajectories as an auxiliary modality during training but does not require them at inference. This enables the predictor to learn dynamics without sequential inference while benefiting from both types of datasets. As a result, our framework achieves over 200 times speedup compared to autoregressive models on the dataset with atomic trajectories while substantially reducing prediction error relative to non-autoregressive benchmarks across both types of datasets. Our code is available at https://github.com/jykim-git/MD.