Online Test-time Adaptation for Interatomic Potentials

Three stages of TAIP include training, test-time adaptation, and inference, as shown in Figure 1. In the training phase, the network parameters are updated according to the supervised learning loss of the main task, namely energy and force prediction, as well as three self-supervised learning losses. In the test-time adaptation phase, we first input the test sample into the encoder and update its parameters once per test sample to reduce the self-supervised learning losses. This unique process facilitates the model’s adaptation to the test data, thereby raising its generalization capabilities on test samples. Since the parameter update during adaptation is performed only once, the computational time and cost will not increase substantially. At the inference stage, the test sample passes through the encoder to predict the energy and forces, which is updated during the adaptation phase.

The main task shown in block 1 of Figure 1 is a typical MLIP task that predicts energy and force. We use a graph to represent a molecular conformation and feed it into the encoder to generate the graph feature. The graph feature is then utilized to predict the potential energy of the molecule, and the force exerted on each atom is calculated by taking the negative gradient of the potential energy.

The noise intensity prediction task is shown in block 2 of Figure 1. For each sample, we perturb the coordinates of atoms with random noises drawn from a Gaussian distribution, and the variance is defined as the noise intensity. The encoder takes the perturbed structure as input and produces the noisy graph feature. Subsequently, the noisy graph feature is combined with the original clean graph feature to jointly estimate the noise intensity, which reflects the dissimilarity between noisy and clean data. Instead of directly predicting the precise value of noise, we categorize the intensity into several bins and determine the bin to which the feature corresponds. This approach naturally groups similar configurations while separating dissimilar ones and thus reduces abrupt changes or outlines, leading to a smoother configuration representation space. By focusing on the key features that distinguish between different levels of noise, the model can learn more robust features and enhance its generalization ability to new scenarios.

The atom feature recovery task is shown in block 3 of Figure 1. We randomly mask some atoms in each sample with a certain ratio and then obtain the atomic representations by feeding the masked atomic structure to the encoder. The representations are subsequently used to rebuild the original atomic features. We can infer the most probable types of atoms in the vicinity of each atom by restoring the masked atoms, and thereby improving the extraction of local information.

The pseudo force recovery task is shown in block 4 of Figure 1. We mask the pseudo forces, which are defined as the negative gradient of the pseudo energy obtained from the masked node features. Then, we feed the masked pseudo forces into the decoder to reconstruct pseudo forces. By restoring the pseudo forces, the model can capture the relationship among the forces on neighboring atoms, which can better understand the local environment of the atoms and is helpful for the prediction of the atomic forces.

We conducted a series of experiments across multiple datasets to evaluate the performance of TAIP on the widely used MLIPs of SchNet [38] and PaiNN [39], which are invariant and equivariant, respectively. First, we tested the enhanced accuracy in predicting energies and forces on the MD-17 dataset [40]. Next, we extended our evaluation to the ISO17 [38] dataset and complex systems featuring periodic boundary conditions, including water and electrolyte solutions. The test data from these datasets were designed to have a larger distribution shift from the training data, as indicated in Supplementary Figure S1, and therefore are used to investigate the improvement in generalizability achieved by TAIP. We then performed MD simulations with MLIPs on the aforementioned complex systems to assess the practical utility of TAIP. Moreover, we explored the impact of TAIP on feature distributions through dimensionality reductions using the t-SNE method [41], illustrating the mechanism behind the improvement from the feature embedding perspective.

MD-17. We evaluated the performance of TAIP on the widely used MD-17 dataset, which contains small organic molecules with reference values of energies and forces generated by ab initio MD simulations. Table 1 compares the mean absolute errors (MAEs) of the predicted energies and forces for different molecules using different models. The results show that TAIP-based models remarkably outperform their corresponding baseline models for all eight molecules, demonstrating the generalization ability of our method.

ISO17. The ISO17 dataset contains trajectories of isomers of C₇O₂H₁₀. We consider two scenarios with the ISO17 dataset. In the first scenario (known molecules/unknown conformation), the isomers in the test set are also present in the training set. In the second scenario (unknown molecules/unknown conformation), the test set contains a different subset of isomers. This task is more challenging and is used to evaluate TAIP in the case where training and test data are drawn from different distributions. Table 2 shows that, for the energy and force MAEs on the ISO17 dataset, TAIP achieves an average reduction of 40% and 31% in scenario 1 and scenario 2, respectively.

Liquid Water and Ice. Water has complex phase behaviors that pose considerable challenges for computational studies. As shown in Figure 2 and Figure 2, there exists a substantial difference between the microscopic structures of liquid water and ice. In the liquid state, water molecules form a highly dynamic network through hydrogen bonding. In contrast, water molecules form a stable hexagonal lattice structure in an ice crystal. The structural differences between different phases can lead to a decrease in prediction accuracy. Thus, we train the models only with liquid water, using a training set of 1000 frames and a validation set of 100 frames, and report the test accuracy on randomly sampled 1,000 liquid water and ice structures from the remaining dataset, respectively. The results of MAE on liquid water (known phase/unknown configuration) and on ice (unknown phase/unknown configuration) are shown in Table 2. The TAIP method notably reduces the errors of force and energy predictions by an average of 40% for unknown configurations across known and unknown phases in the two baseline models.

Electrolyte solutions of different concentrations. We examine TAIP on the electrolyte solution dataset, which is developed in our previous work [23], to further analyze compatibility with complex systems. Consisting of charged ions and electrolyte solvents, the dataset includes more elements, stronger electrostatic forces, and more complicated interatomic interactions, thereby exhibiting a higher degree of complexity than water. Training MLIPs for such complex systems and generalizing the trained models on different concentrations are both challenging. Here, we train the models on the 1 M electrolyte solutions and test the generalization across configurational space at different concentrations. The atomic structures of 1 M and 4 M electrolyte solutions are illustrated in Figure 2 and Figure 2, respectively. The MAE results for the test set of 1 M solutions (known concentration/unknown configuration) and 4 M solutions (unknown concentration/unknown configuration) are presented in Table 2. The data reveals that TAIP substantially enhances the precision of energy and force predictions for unknown configurations, irrespective of whether the concentration is known or unknown.

MD simulations. The MAEs of predicted forces and potential energies on fixed datasets are not enough to characterize the performance of MLIPs on long-time MD simulations [15]. During MD simulations, new molecular structures are constantly produced, and the false prediction of the interatomic forces on a single atom may cause the collapse of the entire simulation system. Therefore, we perform MD simulations using SchNet with and without TAIP on periodic systems (water and electrolyte solutions) to further analyze how TAIP can affect the predicted forces and, thus, the quality of MD simulations.

All simulations are set up with a timestep of 0.5 fs, using the Berendsen thermostat [42] as the temperature coupling method with a coupling temperature of 300 K and a decaying time constant ${\tau}$ of 100 fs. Without TAIP, the simulations for electrolyte solutions with concentrations of 1 M and 4 M collapse due to losing several atoms at the step of 258 ($\sim$0.1 ps) and of 1843 ($\sim$0.9 ps), respectively. The simulation for water exhibits undesired force and energy instability after 250 ps, and that for ice collapses at 248 ps. In contrast, all the simulations with TAIP can run stably for more than 500 ps. The simulation trajectories are provided in Supplementary Figure S2 and Supplementary Video S1-S4. We show the potential energies, kinetic energies, and the maximum absolute force of these trajectories on Figure 3. It is evident that the failure of MD simulations is caused by the divergence of predicted potential energies and forces. Instability is a fundamental issue for MLIP-based MD simulations, as the accuracy of MLIPs without TAIP is crucially compromised when encountered with unknown molecular configurations. The inclusion of the TAIP method remarkably enhances the overall performance of MILP-based MD simulations when dealing with unknown molecular structures, and it can be suggested as an applicable solution to MD instability.

Feature visualization. To visually identify the effect of TAIP. We perform a series of dimensionality reductions using the t-SNE method [41] to map the feature embeddings of the training dataset and the testing dataset on benchmark models and TAIP models. In Figure 4, we can notice that the feature embeddings of the test dataset in the TAIP models are much closer to that of the training dataset than the feature embeddings in the baseline models. It indicates that the TAIP enables the MLIPs to process unknown molecular structures with less uncertainty, which helps the energies and forces predictions to harness more accuracy and generalizability.

MD-17. The MD17 dataset is composed of ab initio molecular dynamics trajectories of eight small molecules [40]. Obtained from http://www.sgdml.org/#datasets, we use the training set of 1,000, the validation set of 1,000, and the test set of 1,000 configurations for each small molecule.

ISO17. ISO17 dataset consists of ab initio MD trajectories of 129 isomers whose energies and forces are calculated by Density Functional Theory (DFT) [38]. We adopted the same splitting strategy reported in the original literature[38], using 80% trajectory steps of the 80% isomers for training and validation (totaling 404,000 molecular conformations for training and 4,000 for validation). We further set up two test datasets to evaluate the effectiveness of TAIP: (1) The other 20% unseen trajectory steps of the isomers included in the training set (totaling 101,000 conformations) and (2) all of the remaining 20% isomers not included in the training and validation set (totaling 130, 000 conformations).

Liquid Water and ice. The atomic structures are sampled from MD trajectories, which consist of a training set of 1,000 snapshots, a validation set of 500 snapshots, and a test set of 1,000 liquid water. In addition, 1,000 ice snapshots in the form of hexagonal ice crystals are sampled to be the other test set. There are 96 molecules or 288 atoms in each snapshot. The forces and energies for each sampled snapshot are calculated by DFT using the cp2k package[43]. The DFT calculations are conducted using PBE exchange-correlation functional[44] with the Projector Augmented-Wave (PAW) pseudo-potential[45]. The DFT-D3 method is used for the vdW-dispersion energy correction[46]. A single gamma k-point is used to sample the Brillouin zone, the cutoff energy for the plane-wave-basis set is set to be 400 eV, and the electronic self-consistency is considered to be achieved if the change of total energy between two steps is smaller than $10^{-6}$ eV.

Electrolyte solutions. The electrolyte solutions dataset is taken from our previous work [23]. Here, we used the combination of [$\rm Li^{+}$$\rm PF_{6}^{-}$DME], $[\rm Na^{+}$$\rm PF_{6}^{-}$DME], $[\rm Li^{+}$$\rm Tf_{2}N^{-}$DME], $[\rm Na^{+}$$\rm Tf_{2}N^{-}$DME], [$\rm Li^{+}$$\rm PF_{6}^{-}$EC+DMC], $[\rm Na^{+}$$\rm PF_{6}^{-}$EC+DMC], $[\rm Li^{+}$$\rm Tf_{2}N^{-}$EC+DMC], and $[\rm Na^{+}$$\rm Tf_{2}N^{-}$EC+DMC] as electrolyte solutions with ionic concentrations of 1 M and 4 M. The training and validation data are randomly selected from the 1 M solutions. The training set contains 1000 samples and the validation set contains 250 samples. We construct two different test sets to evaluate TAIP: one is randomly sampled from the remaining 1 M solutions and the other is sampled from the 4 M solutions.

The TAIP framework is implemented based on PyTorch 1.8.0. The experiments are conducted with NVIDIA GeForce RTX 3090 GPU. The models are trained using the Adam optimizer [47], employing single-GPU training for efficient processing. The hyperparameters are provided in Supplementary Table S1-S3.

In order to achieve optimal performance on both properties, we incorporate both energy $E$ and forces $\mathbf{F}_{i}$ into the criterion for training:

For the periodic systems including water and electrolyte solutions, we employ the atomization energy as the target energy by shifting the original potential energy according to the energy of each single atom:

Noise intensity prediction. For each sample, we randomly choose a noise intensity $l$ from a uniform distribution $U(1,L)$ and perturb the coordinates accordingly, which will facilitate the exploration of the configurational space. The GNNs are trained to predict the noise intensity instead of the specific noise value. The graph feature of the perturbed structure $\mathbf{\hat{u}}$ is concatenated with that of the original structure $\mathbf{u}$ and passes through a classification head $\mathbf{\Phi}_{\text{scale}}$ to produce the classification logit. The training objective can be formulated as follows:

Atomic feature recovery. The masked atomic structures are fed into the encoder to get the masked node features. We utilize the cosine error as the criterion for reconstructing the original node features. Additionally, we reduce the weight of easy samples in the training set by scaling the cosine error with a power of $\gamma$, as shown below:

Pseudo force recovery. The pseudo force is defined as the negative gradient the pseudo energy, which is obtained from the masked node features. We utilize the L1 loss as the criterion for reconstructing the pseudo force: