Teaching Molecular Dynamics to a Non-Autoregressive Ionic Transport Predictor

This paper proposes a non-autoregressive learning framework that utilizes atomic trajectories as an auxiliary modality during training to enable fast, accurate, and dynamic ionic transport prediction from static structures without requiring sequential inference or trajectory data at inference time.

The Gist

Imagine you are trying to predict how fast a crowd of people (ions) can move through a crowded room (a solid material) to get from one side to the other. This speed is crucial for things like how fast your phone battery charges.

Traditionally, scientists have tried to figure this out in two ways, both of which have big problems:

1. The "Slow Motion" Method (Molecular Dynamics): They simulate every single step the people take, second by second. It's incredibly accurate, but it takes so much computer power and time that it's like trying to watch a movie in slow motion just to see if the actors can run. It's too slow for testing thousands of materials.
2. The "Snapshot" Method (Non-Autoregressive Models): They look at a single photo of the room (the static atomic structure) and guess the speed. It's instant, but because they can't see how the people move, their guesses are often wrong. They miss the "dynamics" of the crowd.

The Problem:
There is a third option: a method that generates a movie of the movement step-by-step (autoregressive). But this is still slow and prone to errors piling up (like a game of "telephone" where the message gets garbled). Also, most of the data scientists have is either just the "snapshot" (no movement data) or the full "movie" (movement data), but rarely both.

The Solution: "Teaching" the Predictor
The authors of this paper created a new framework that acts like a smart teacher. They want a student (the predictor) that can look at just a "snapshot" and instantly guess the crowd's speed, but they want that student to be as smart as if they had watched the whole "movie."

Here is how they do it, using a creative analogy:

1. The "Dual-Modal" Teacher (Training with the Movie)

First, they build a "Teacher" model. This teacher gets to see both the static photo of the room and the full movie of the people moving. Because it sees the movement, it learns the deep, complex rules of how the crowd flows. It becomes an expert.

2. The "Student" (The Fast Predictor)

Next, they build a "Student" model. This student is designed to be super fast. It can only look at the static photo (no movie allowed during the test). The goal is to make the student so good that it can guess the speed without ever seeing the movie.

3. The "Secret Transfer" (Model-Level Learning)

How do they teach the student without showing it the movie?

• They don't just ask the student to copy the teacher's final answer.
• Instead, they force the student to mimic the internal thoughts (hidden representations) of the teacher.
• The Magic Trick: They use a mathematical shortcut (called "closed-form initialization," which is like solving a puzzle with a direct formula rather than guessing and checking) to instantly align the student's brain with the teacher's brain. The student learns, "Oh, when the teacher sees this specific room layout, it thinks this about the movement." The student memorizes the logic of the movement without needing the actual video.

4. The "Chain Reaction" (Data-Level Learning)

Here is the really clever part. Most real-world data only has the "snapshot" (no movie).

• The authors realized that even if a new dataset has no movies at all, they can still use the knowledge from the dataset that did have movies.
• They take the "Teacher" and the "Student" (who learned from the movie) and use them to initialize a new student for the "snapshot-only" data.
• It's like taking a master chef who learned to cook with fresh ingredients (the movie data) and teaching them to cook with canned ingredients (the snapshot-only data). The chef still knows the flavor profile and techniques, so they can make a great dish even without the fresh ingredients.

The Results

• Speed: Their method is 200 times faster than the slow "step-by-step" simulation methods. It's like switching from watching a movie in slow motion to snapping a photo.
• Accuracy: It is much more accurate than other fast methods that just look at the photo. By "learning" the dynamics from the teacher, the fast predictor makes fewer mistakes.
• Versatility: It works even when the data is messy, comes from experiments (not just simulations), or involves different types of ions (like swapping Lithium for Sodium).

In Summary:
The paper presents a way to train a fast AI to predict how ions move through materials. It does this by using a "teacher" that watches the movement to train a "student" that only sees the static structure. The student learns the essence of the movement so it can make lightning-fast, accurate predictions without needing to run expensive, slow simulations. This helps scientists screen new battery materials much faster than before.

Technical Summary

Technical Summary: Teaching Molecular Dynamics to a Non-Autoregressive Ionic Transport Predictor

Problem Statement
Predicting ionic transport properties (e.g., diffusivity, conductivity) from static equilibrium atomic structures is a fundamental challenge in materials science, particularly for rechargeable batteries. Unlike static properties, ionic transport is inherently dynamic, requiring the inference of long-time atomic motion from static inputs. The current gold standard, Molecular Dynamics (MD) simulations, is computationally prohibitive for large-scale screening due to the need for extremely small time steps and long simulation times to capture rare diffusion events.

Existing machine learning approaches face a trade-off between speed and accuracy:

1. Autoregressive MD acceleration methods generate atomic trajectories sequentially. While they capture dynamics, they suffer from slow inference and error accumulation, which can cause trajectory divergence.
2. Non-autoregressive material property predictors offer fast, single-pass inference but fail to exploit dynamical information, leading to lower accuracy because they cannot access atomic trajectories as input.
3. Data Scarcity: Ionic transport datasets are scarce. Some contain atomic trajectories (from MD), while others (often experimental or large-scale MD-derived) contain only static structures and target properties. Autoregressive models cannot train on structure-only data, while non-autoregressive models cannot utilize the dynamic information present in trajectory-based datasets.

Methodology
The authors propose a non-autoregressive learning framework based on auxiliary modality learning. The core idea is to treat atomic trajectories as a "privileged" modality available only during training to teach the model dynamics, while the final predictor operates solely on static structures during inference.

The framework consists of two main components:

1. Model-Level Auxiliary Modality Learning:

   • Dual-Modal Trainer ($g$): A model trained on trajectory-based datasets ($\mathcal{D}_{trj}$) using both equilibrium structures ($x$) and atomic trajectories ($p$) as inputs. It employs a trajectory encoder ($W_p$) and a structure-temperature encoder ($W_{x,T}$).
   • Regularization: To prevent the model from relying solely on the trajectory encoder, a regularization term forces the structure encoder to produce accurate predictions independently.
   • Closed-Form Initialization: The knowledge from the dual-modal trainer is transferred to a non-autoregressive predictor ($f_1$) via a closed-form ridge regression solution. This aligns the hidden representations of the predictor (using only structure inputs) with those of the dual-modal trainer (using both inputs). This avoids iterative gradient-based distillation, which is less effective in data-scarce regimes.
   • Embeddings: The framework leverages scientific foundation models: SevenNet (an MLIP foundation model) for extracting structural embeddings from equilibrium structures, and MOMENT (a time-series foundation model) for condensing atomic trajectories into embeddings via Fourier transforms.

2. Data-Level Auxiliary Modality Learning (Optional):

   • Designed for structure-based datasets ($\mathcal{D}_{str}$) that lack atomic trajectories.
   • It initializes a new predictor ($f_2$) by transferring the structure encoder from the dual-modal trainer and the decoder from the trajectory-trained predictor ($f_1$).
   • This enables models trained on structure-only data to benefit from the dynamical knowledge learned from trajectory-based datasets, even when the datasets differ in ion species, data sources (simulation vs. experiment), or target definitions.

Key Contributions

• Dynamics-Aware Non-Autoregressive Prediction: The first framework to formulate atomic trajectories as a privileged modality for ionic transport prediction, enabling accurate, trajectory-free inference.
• Efficient Knowledge Transfer: Introduction of a closed-form initialization based on ridge regression. This method is shown to be more effective than conventional gradient-based distillation in data-scarce settings, allowing the predictor to reproduce the hidden representations of a teacher model without iterative optimization.
• Cross-Dataset Generalization: The ability to transfer dynamical knowledge from trajectory-based datasets to structure-based datasets (and across different ion species and target properties) using data-level auxiliary modality learning.
• Integration of Foundation Models: Effective utilization of pre-trained scientific foundation models (SevenNet and MOMENT) to extract informative embeddings without task-specific fine-tuning of the backbone.

Experimental Results
The framework was evaluated on three datasets: a trajectory-based MD dataset (Dataset 1), a structure-based MD dataset (Dataset 2), and a real-world experimental dataset (Dataset 3).

• Speed: On the trajectory-based dataset, the proposed method achieves a 200× speedup in inference time compared to state-of-the-art autoregressive models (e.g., LiFlow), while maintaining comparable or better accuracy.
• Accuracy:
  • On trajectory-based data, the method significantly outperforms non-autoregressive benchmarks (MatFormer, ComFormer, DenseGNN) and even surpasses autoregressive baselines in Mean Absolute Error (MAE) for log-scaled targets.
  • On structure-based datasets (including experimental data), the framework substantially reduces prediction error compared to existing non-autoregressive benchmarks. For example, on the experimental dataset (Dataset 3), the MAE was reduced from ~2.0 to 1.388 (log scale), a level of error comparable to the natural variability of experimental measurements.
• Generalization: The model successfully generalizes to unseen ion species (Na) and different material classes (polymers), demonstrating the transferability of the learned dynamical knowledge.
• Ablation Studies: Confirm that both model-level and data-level auxiliary modality learning, the closed-form initialization, and the use of foundation models are critical to performance.

Significance and Claims
The paper claims that this framework offers a general pathway to accelerate MD-based material property prediction. By decoupling the need for atomic trajectories during inference from the training process, it enables fast, accurate, and stable inference without the error accumulation inherent in autoregressive methods.

The authors emphasize that while the method is designed for initial screening to filter candidate materials, the achieved error levels on experimental data are practically meaningful. They note that the framework is readily extensible to other material properties governed by atomic dynamics. However, they modestly acknowledge limitations, such as the need for further systematic analysis on how scientific foundation models affect the framework and the conditions under which the linear encoder assumption holds. The work aims to reduce the computational cost and energy footprint of large-scale materials screening, thereby accelerating the discovery of ion-conducting materials for energy technologies.