Title: Transfer learning for atomistic simulations using GNNs and kernel mean embeddings
Venue: NeurIPS 2023
Accepted: yes
Keywords: GNN

Mean Embedding

Kernels

Atomistic Simulations

OCP

Transfer Learning

Molecular Dynamics

Kernel Ridge Regression

Neural Networks
Abstract: Interatomic potentials learned using machine learning methods have been successfully applied to atomistic simulations. 
However, accurate models require large training datasets, while generating reference calculations is computationally demanding. To bypass this difficulty, we propose a transfer learning algorithm that leverages the ability of graph neural networks (GNNs) to represent chemical environments together with kernel mean embeddings. We extract a feature map from GNNs pre-trained on the OC20 dataset and use it to learn the potential energy surface from system-specific datasets of catalytic processes. Our method is further enhanced by incorporating into the kernel the chemical species information, resulting in improved performance and interpretability. We test our approach on a series of realistic datasets of increasing complexity, showing excellent generalization and transferability performance, and improving on methods that rely on GNNs or ridge regression alone, as well as similar fine-tuning approaches.

Transfer learning for atomistic simulations using GNNs and kernel mean embeddings

Abstract

Interatomic potentials learned using machine learning methods have been successfully applied to atomistic simulations. However, accurate models require large training datasets, while generating reference calculations is computationally demanding. To bypass this difficulty, we propose a transfer learning algorithm that leverages the ability of graph neural networks (GNNs) to represent chemical environments together with kernel mean embeddings. We extract a feature map from GNNs pre-trained on the OC20 dataset and use it to learn the potential energy surface from system-specific datasets of catalytic processes. Our method is further enhanced by incorporating into the kernel the chemical species information, resulting in improved performance and interpretability. We test our approach on a series of realistic datasets of increasing complexity, showing excellent generalization and transferability performance, and improving on methods that rely on GNNs or ridge regression alone, as well as similar fine-tuning approaches.

1 Introduction

Atomistic simulations have become a pillar of modern science and are pervasively used in many areas of physics, chemistry, and biology. Among these techniques, molecular dynamics plays a prominent role. This method simulates the time evolution of a system of atoms by integrating Newton's equation of motion [1]. The forces acting on the particles are determined from a model for the interactions, called the potential energy surface (PES), on whose accuracy the reliability of the simulation depends. For a long time, the interactions were modeled in a rather empirical, and of course, not very accurate way [1, 2]. A significant step forward was made with the introduction of ab initio molecular dynamics in which the interactions are computed on the fly from accurate electronic structure calculations [3, 4]. This implies solving at every step the Schroedinger equation, typically with the use of some approximation such as the popular Density Functional Theory (DFT) scheme [5]. This approach is much more accurate but at the same time more computationally expensive, limiting the system size (i.e. the number of atoms) and the time scale that can be simulated.

Starting with the work of Behler and Parrinello [6], machine learning potentials have emerged as promising candidates to alleviate the tension between accuracy and efficiency [7, 8]. They regress the

potential energy, as a function of the atomic positions and the chemical species, on a (large) set of expensive ab initio calculations. Once successful, this strategy results in an ab initio -quality model of the potential energy at a fraction of the cost, speeding up simulations by orders of magnitude.

For this procedure to be successful, however, a good representation of the physical system, including its symmetries, is necessary. The use of handcrafted physical descriptors [6, 9-11] is often a laborious procedure that limits applicability to systems with few chemical species. In recent years, graph neural networks (GNNs) have proven to be a viable alternative for directly representing the chemical environment, capable of scaling to large datasets with numerous chemical species and encoding symmetries directly in the architecture, e.g., through SE(3)-equivariant layers [12-16]. Nevertheless, obtaining an accurate model for real-life applications is as of today still a challenging task, requiring high-quality data samples which are scarce and/or very expensive to obtain.

As shown by the recent advancements in large language modeling [17, 18] and, before that, image classification [19, 20], fine-tuning a pre-trained large-scale representation of the data provides a highly effective paradigm to solve downstream tasks for which only a handful of data points are available. This approach paved the way for the concept of foundation models [21], at the core of the current generative AI revolution. Mirroring these developments, in this work we leverage the representation power of GNNs trained on large datasets of molecular configurations. In particular, we rely on the Open Catalyst dataset [22], which contains DFT relaxations for ∼ 1.2M catalytic systems, totaling over 260M data points. We show how, by exploiting the availability of this heterogeneous dataset, it is possible to learn interatomic potentials for specific systems taken from realistic chemical applications in a fast and data-efficient manner.

Contributions This paper makes the following contributions:

We propose a transfer learning algorithm, which we refer to as mean embedding kernel ridge regression (MEKRR), for modeling the potential energy surface of atomic systems. MEKRR combines GNN representations pre-trained on large datasets with fine-tuning via kernel mean embeddings. This combination allows to satisfy the physical symmetries inherent to atomistic systems. Specifically, GNN features take care of roto-translational invariance, while kernel mean embeddings are chosen to satisfy the permutational symmetry.

We introduce a new kernel function in the context of modeling potential energy surfaces, which exploits chemical species information. This shows superior performance and facilitates monitoring the chemical evolution of the system.

We demonstrate excellent transferability and generalization performance on increasingly complex datasets. Remarkably, they include configurations sampled out-of-distribution with respect to the GNN pre-trained representation.

Related work There is a long list of relevant works (see also [7] and references therein) on representing the potential energy surface by machine learning methods. In particular, the first models employed a set of physical descriptors in combination with either neural networks [6] or kernel methods [23-25]. Later, it was proposed to model with neural networks the descriptors as well [26]. Recently, graph networks have been used to directly represent physical systems and regress energy and forces [12, 15, 27, 28]. A combination of GNN features and kernel-based methods has been investigated in [29], but without doing transfer learning from a larger dataset.

In terms of transfer learning, there is a long line of work in transferring deep representations on images, see [19, 20] and references therein. Current progress in few-shot image classification, where algorithms can adapt quickly to new classification problems, can partially be attributed to adding a preprocessing step where each image is mapped through a meta-learned or pre-trained representation, see e.g. [30-36]. In particular, [36] employs kernel ridge regression (KRR) on top of a meta-learned feature map. More recently, the same ideas have been applied to language modeling (see e.g. [17, 18]), revolutionizing the field. Pre-training strategies have started very recently to emerge also in the context of interatomic potentials, as a way to interpolate across different levels of theory [37-39] or to exploit the release of large and heterogeneous datasets such as OC20 [40].

Organization In Section 2 we specify the machine learning problem we are aiming to solve. In Section 3 we introduce the kernel mean embedding framework together with KRR and in Section 4 we introduce our method MEKRR. In Section 5 we validate our method on a variety of realistic datasets of increasing complexity. Finally, in Section 6 we conclude and outline future directions.

2 Learning the potential energy surface from atomistic simulations

2.1 Setting

We consider systems composed of n atoms and S different chemical species, described via one-hotencoding over [ S ] , where [ S ] = { 1 , . . . , S } . Each atom is described by its Euclidean position r ∈ R 3 and its chemical species z ∈ { 0 , 1 } S . We denote a state of the system by x = ( r i , z i ) n i =1 = ( R,Z ) ∈ X using the design matrices R ∈ R n × 3 and Z ∈ { 0 , 1 } n × S .

The quantity we want to regress is the potential energy, which is a scalar function E : X → R . The target values E t are calculated by querying the ab initio method of choice, DFT in the case of the OC20 dataset. A training dataset is therefore composed by a sequence of T atomic configurations ( x t ) T t =1 together with the corresponding labels ( E t ) T t =1 .

The physics of the system posits that any estimator of the energy ˆ f : X → R should possess certain invariant properties reflecting the underlying physical symmetries. These invariances are roto-translational invariance with respect to the position vectors and permutation-invariance within each group of chemical species . In addition, the energy estimator should be curl-free and smooth [27]. A wealth of previous works introduced deep learning architectures incorporating these invariances by design [e.g. 16, 27, 41, 42]. Our transfer-learning scheme works by exploiting these invariant architectures so that the final fine-tuned estimator is invariant as well.

2.2 Objective function

In order to learn potential energies, one considers models f w from a set of functions (e.g., a reproducing kernel Hilbert space or neural network functions) parameterized by w and fitted to a given dataset ( x t , E t ) T t =1 . This is achieved by minimizing an objective function, which can be split into a data-fitting term and a regularizer which encourages functions of low complexity.

Typically, the data fitting term is a least squares loss between outputs and predictions. In this case the regularized empirical risk minimization reads

̸

for some regularizer Ω and regularization parameter λ > 0 and a loss weight γ ∈ [0 , 1] . A loss weight γ = 0 is used whenever the forces F ( x, r ) = -∇ r E ( x ) are regressed alongside the energy. Typically we will use Ω( w ) = ∥ w ∥ 2 as the regularizer, but in Section 4.3 we extend this to incorporate additional information. Since in the experiments discussed below we focus on energy prediction, we let γ = 0 .

3 Kernel ridge regression and kernel mean embeddings

Our method is designed around kernel methods, a well-established tool at the heart of most nonparametric machine learning algorithms [43, 44]. A kernel K is any positive definite function (see 44, Definition 4.15) K : X × X → R on the input space X . For any kernel it exists a feature space H K and a feature map ϕ : X → H K such that K ( x, x ′ ) = ⟨ ϕ ( x ) , ϕ ( x ′ ) ⟩ for all x , x ′ ∈ X .

Kernel ridge regression is a supervised learning algorithm that parametrize the dependence between inputs x ∈ X and scalar outputs y ∈ R as y = ⟨ w,ϕ ( x ) ⟩ for some vector w in the feature space. Given a dataset ( x t ) T t =1 ∈ X T of T input points, and a corresponding dataset of outputs ( y t ) T t =1 ∈ R T , KRRlearns a functional relation between inputs and outputs by solving the least squares minimization problem

A basic result known as the representer theorem (see e.g. 44, Theorem 5.5) prescribe that the solution of (2) is just a linear combination of the feature map evaluated at the training points, that is ˆ w = ∑ T t =1 c t ϕ ( x t ) for some ( c t ) T t =1 ∈ R T . The representer theorem and the fact that K ( x, x ′ ) = ⟨ ϕ ( x ) , ϕ ( x ′ ) ⟩ also imply that predicting any new point x with KRR is as simple as evaluating

⟨ ˆ w,ϕ ( x ) ⟩ = ∑ T t =1 c t ⟨ ϕ ( x t ) , ϕ ( x ) ⟩ = ∑ T t =1 c t K ( x t , x ) . KRR can also be readily extended [45] to regress the gradient of y with respect to x , a technique used e.g. in [29] to concurrently learn potential energies and forces.

In the standard case in which the input space is a subset of an Euclidean space X ⊆ R d , many kernels have been designed; see, for instance, 46, Section 4.2. This notwithstanding, kernel functions can be defined on arbitrary input spaces X . In the case of atomic configurations, the potential energy is invariant under the permutation of atoms with the same chemical species. It is therefore advantageous to design kernels able to preserve such symmetry. We now introduce the concept of kernel mean embeddings [47], which allow us to define permutationally-invariant kernels over atomic configurations.

Kernel mean embeddings consider the case in which inputs x are sets of points living in an Euclidean space x = { r i ∈ R d : i ≤ n } . For example, an atomic configuration is a set containing the position of each atom in the system ( d = 3 ). For any feature map ψ ( r ) and the corresponding kernel function k ( r, r ′ ) = ⟨ ψ ( r ) , ψ ( r ′ ) ⟩ on points r ∈ R d we can define a feature map acting on the whole set x as

with C x > 0 a positive normalization constant. We note in passing that the value of C x depends on the chemical property that we want to learn. Indeed, setting C x = 1 returns extensive properties such as the potential energy, while C x = 1 /n is appropriate to model intensive ones. This allows us to define kernels on sets x as

Using the definition of kernel mean embeddings (3) inside the KRR algorithm (2) enables us to learn scalar functions over sets of points. In the following section, we will combine these algorithms together with pre-trained feature maps to define a principled and efficient method to learn potential energy surfaces from data.

4 Method

Our method relies on pre-trained GNN representations of chemical environments and uses them as feature maps to define a kernel mean embedding (3) acting on atomic configurations. In practice, we do so by leveraging pre-trained feature maps on the OC20 dataset [22]. Once a kernel based on mean embeddings is defined, we can use KRR to regress the potential energy from sampled data (see Fig. 1 for a diagram of the method). In this respect, we note that our method shares similarities with the idea behind foundation models [21], where a representation trained on large datasets is transferred to novel settings by means of fine-tuning. While in this work we show the performance of pre-trained representations based on SCN [14] and SchNet [27], any other GNN architecture is a perfectly valid choice, and can be used in place of ours without the need of further adjustments.

4.1 GNN representations of chemical environments

SCN and SchNet are instances of graph neural networks (GNNs) [48-50], a class of architectures designed to learn mappings over graphs. A graph G = ( V , E ) is a collection of n nodes V and edges E ⊆ V × V between pairs of nodes. Representing the chemical environment with GNNs involves a preprocessing step that turns a configuration x into a graph by associating each atom to a node and constructing the edges according to the matrix of pairwise distances between atoms. Each node i is then initialized to a feature vector h i encoding the chemical species of the atom via a (possibly learnable) embedding layer. Each GNN layer then updates the node features via a nonlinear message-passing scheme

Here, N ( i ) is the neighborhood of i i.e. the set of nodes connected to i through an edge, g θ is a learnable vector-valued function such as an MLP and η θ is a message-passing function. A graph

H

K(H, H')

H'

2

hei

a.Atomic feature maps from a pretrained GNN

neural network is formed by concatenating multiple message-passing layers. The weights of the embedding and GNN layers are then learned in an end-to-end fashion typically using first-order optimization.

Among the many GNN architectures available [51], a growing number is being developed to specialize in computational chemistry settings, see for instance [16, 28, 41, 42, 52-56].

4.2 Our method: mean embeddings of GNN representations

We are now ready to combine together the techniques described so far and describe in full detail our method. Starting with an atomic configuration x , we pass it to a pre-trained GNN representation to get a collection of node features

In the case of equivariant GNNs we additionally need to extract invariant features from equivariant ones via a pooling method such as averaging. In case of SCN [14], for example, we average the features over the sphere (see Appendix C.2). We then evaluate the kernel mean embedding of the node features

Finally, at training time we can use the kernel mean embedding ϕ ( H ( x )) to solve the KRR problem as

while at inference time we can get an estimation for the potential energy at x by simply evaluating ⟨ ˆ w,ϕ ( H ( x )) ⟩ .

Some practical considerations are now in order. Given a pre-trained GNN representation, one is not bound to use the node features of the last layer. Intermediate representations have the benefit of being faster to evaluate, and we empirically observed that the performance of MEKRR is not hindered when truncating the GNNs pre-trained on OC20 at an intermediate layer (see Appendix C). Furthermore, there are many strategies to solve the KRR problem (4) based e.g. on the solution of a linear system or on gradient-descent. In Algorithm 1 we report the pseudo-code describing our implementation of the training and prediction steps of MEKRR.

4.3 Chemically-informed kernel mean embeddings

The chemical species of each atom is implicitly encoded in the GNN representation through its initial layer, embedding each atom of the same chemical species to the same dense feature vector. In MEKRR we strengthen this dependence by encoding the same information again at the level of mean embeddings. For a configuration x let now H s ( x ) := { h i ( x ) | species of i = s } be the subset

Algorithm 1 Mean Embedding Kernel Ridge Regression (MEKRR)

Parameters : Training dataset ( x t , E t ) T t =1 of configurations x t and potential energies E t . Kernel function k . Interpolation parameter of the multi-weight formulation α ∈ [0 , 1] . Tikhonov regularization λ > 0 . Pre-trained GNN feature map x ↦→ H ( x ) .

of node features corresponding to atoms of the s 'th chemical species. We then consider a composite kernel

The kernel K α is a direct generalization of the mean embedding kernel (3) where the parameter α allows to interpolate between emphasizing all atomic interactions equally ( α = 0) and only within-species atomic interactions ( α = 1 ).

Solving KRR (4) with the composite kernel (5) is also equivalent to considering an extension of (4) where to each chemical species is given its own weight while encouraging the weights to be close to each other and of small magnitude through a regularizer. Precisely, let ( w s ) S s =1 be the set of displacements from a center weight w 0 for each chemical species. The weight of a chemical species s is given by w s + w 0 and the potential energy function for a configuration x takes the form ∑ S s =1 ⟨ w s + w 0 , ϕ ( H s ( x ) ⟩ . In Appendix A we provide a full characterization of this extension, highlighting its connection to multi-task learning [57, 58], in which multiple tasks are learned jointly.

5 Experiments

In this section, we consider the realistic problem of modeling the potential energy surface of catalytic reactions occurring on metallic surfaces. We evaluate our approach against methods that are representative of the KRR and GNN approaches, testing them on datasets of increasing complexity taken from realistic applications. As is customary in the literature, we employ the root mean squared error (RMSE) normalized by the number of atoms as a metric, to facilitate the comparison between systems of different sizes. We make the code repository available at https://github.com/IsakFalk/atomistic_transfer_mekrr .

5.1 Baselines and MEKRR

We consider baselines spanning different categories. Firstly we examine supervised learning algorithms trained from scratch on the provided datasets, either through GNNs or kernel methods with hand-crafted physical features. In this category we have Schnet, SCN and GAP. SchNet [59] is one of the first GNNs to be applied successfully to chemistry, which uses a radial basis function representation of the interatomic distances. Spherical Channel Networks ( SCN ) [14] is another GNN and its atom embeddings are a set of spherical functions represented via spherical harmonics. SCN is one of the state-of-the-art models on the OC20 dataset. For both SchNet and SCN we use the codebase of [22]. Finally, Gaussian Approximation Potential ( GAP ) is a kernel-based method that builds a Gaussian Process using the Smooth Overlap of Atomic Positions (SOAP) descriptors [60], which we use through the QUIP/quippy code base [61, 62].

The second category of baselines concerns transfer learning methods based on the OC20 dataset [22]. In this case we consider the fine-tuning of Schnet ( Schnet-FT ) which is done by keeping the parameters fixed up to the representation used for MEKRR and then optimizing the subsequent layers on the new dataset.

These baselines are tested against our method ( MEKRR ), which uses a kernel mean embedding with Gaussian kernel based on different pre-trained GNN features. The length-scale of the Gaussian kernel is chosen according to the median heuristic [63]. We will denote MEKRR-(SchNet) and MEKRR-(SCN) the variants using Schnet and SCN node features as inputs, respectively.

5.2 Datasets

We first describe the dataset which has been used for the construction of the pre-trained GNN feature map ( OC20 ), and then present the system-specific MD datasets where our method is fine-tuned on ( Cu/formate , Fe/N 2 ).

OC20 The Open Catalyst (OC) 20 is a large dataset of ab initio calculations aimed at estimating adsorption energies on catalytic surfaces. It comprises ∼ 250 millions of DFT calculations, generated from over 1.2 million relaxations trajectories of different combinations of molecules and surfaces. In each relaxation, the positions of the molecule and of the surface upper layers are optimized via gradient descent in order to compute the adsorption energy. The adsorbate is selected out of 82 molecules relevant to environmental applications, while, for each of them, up to 55 3 surfaces are selected, including binary and ternary compounds. We underline that, for each adsorbate-surface pair, the configurational space sampled is very limited, and especially it does not cover out-of-equilibrium and reactive (e.g. bond forming or breaking) events.

We fine-tune the method and test it on two datasets that are representative of reactive catalytic events, obtained by means of molecular dynamics simulations coupled with enhanced sampling methods [64, 65] to avoid mode collapse into the local minima of the potential energy landscape. Indeed, whereas the OC20 dataset contains short, correlated relaxations toward the nearest equilibrium state, typical catalytic reaction datasets require sampling all local minima (adsorption states of the molecule) and especially reactive events, in which, due to interaction with the surface, bonds can be broken or formed. For this reason, these applications are challenging as they relate to realistic datasets containing mostly reactive events that are outside the distribution of the OC20 dataset. We split all the below datasets into a train, validation, and test set using random splitting of 60 / 20 / 20 .

Cu/formate The first dataset is a collection of molecular dynamics simulations of the dehydrogenation reaction of formate on a copper (Cu) <110> surface [15], initialized along the reaction path (obtained with the Nudged Elastic Band method [65]), in which the molecule loses its hydrogen atom upon interaction with the surface.

Fe/N 2 ( D i ) The second dataset consists of molecular dynamics simulations of a nitrogen molecule adsorbing on an iron (Fe) <111> surface at high temperature ( T = 700 K) and breaking in two nitrogen atoms [66]. A peculiarity of this dataset is that it contains data from different sources (e.g. standard and biased molecular dynamics) and system sizes, allowing us to also assess the transferability of the methods across different conditions. For this reason, we divide it into 4 subsets, denoted with D i :

D 1 : AI-MD Ab initio molecular dynamics simulations. The resulting configurations are highly correlated and cover a small portion of the configurational space related to the adsorption process, thus being the closest dataset to the OC20 one.

D 2 : AI-METAD Here the ab initio MDsimulation is accelerated with the help of the metadynamics [64] technique. This is an importance sampling method that allows rare events to be observed, and thus it has been employed for collecting reactive configurations in the training set [67]. Due to the metadynamics approach, a larger region of configurational space is sampled with respect to D 1 , allowing one to sample one bond-breaking event.

D 3 : AL-METAD Dataset built from an active learning procedure using an ensemble of NNs combined with metadynamics. In this simulation, multiple reaction events are observed, covering a wider region of the configurational space and providing a large number of

uncorrelated samples. Hence, these configurations are far from those used to pre-train the feature map.

D 4 : AL-METAD-72 Same as D 3 but the surface is composed of 72 atoms (8 layers) to test the transferability across systems of different sizes.

5.3 Interpolating between shared and independent weights

The α parameter in the K α kernel can vary in the range [0 , 1] , which are the limiting cases between a shared or independent set of weights for each chemical species. We use cross-validation to set this parameter in practice. To initially fit the regularization parameter λ we set α = 0 and cross-validate λ ∈ { 10 -3 , . . . , 10 -9 } using the same datasets. Despite this simple heuristic cross-validation scheme, as we will see, the scheme is effective, which we believe is a strength as it shows that the MEKRR method is simple to tune while still having the strongest performance among the competitors.

In Fig. 2 we show the cross-validation curves for MEKRR-(SchNet), related to the two datasets Cu/formate and Fe/ N 2 . In the latter case, we perform the cross-validation only on ( D 2 ), which is representative of the family of datasets, and then use the found parameters also for the other datasets. The two plots show different behavior with the optimal α for the Cu/formate dataset occurring around 10 -2 . This means that the potential energy can be well described with shared weights across chemical species together with a small perturbation. Instead, in the Fe/ N 2 dataset the optimal α occurs at the boundary leading to a kernel in which the weights for the two species are learned independently.

Cu/formate

Fe/N2 (D2)

2.00

atom

RMSE

1.75

MAE

a

e

1.50

1.25

err

1.00

0.75

0.50

10-4

10-2

100

10-4

10-2

100

5.4 Potential energy regression

In this section, we consider the setting of predicting the potential energy surface. We first evaluate the performance of the models in predicting the energy for each of the datasets and then assess their generalization performance through the transfer learning setting, where we train and test on similar but distinct datasets.

Same-dataset energy prediction From Table 1 we see that MEKRR achieves the best performance in all datasets, both when the input features are extracted from SchNet and SCN. We note that, in general, the transfer learning algorithms (SchNet-FT and the two MEKRR variants) outperform the ones trained from scratch, with MEKRR being significantly faster (see Appendix B). Furthermore, it is worth highlighting that our method performs better than the baselines even when it is applied to datasets that are out-of-distribution for the pre-trained feature map. This is particularly evident for D 3 and D 4 which contain multiple reactive events (bond-breaking) that were never seen in the relaxations composing the OC20 dataset. This demonstrates the ability of the GNN trained on large and heterogeneous datasets to effectively represent chemical environments.

Across-dataset energy prediction Here we evaluate the performance of the algorithms and MEKRR on transferring from different systems in the Fe/N 2 family of datasets. To do this we consider the task of zero-shot transfer learning (see e.g. [18] and references therein) where we evaluate a model trained on a source dataset D source on a target dataset D target . While the two datasets D source and D target may be sampled from arbitrary systems, we consider here systems that share some characteristics as we are evaluating the transfer capability of the models [68, 69]. Due to the ordering of the datasets D 1 , . . . , D 4 in increasing complexity on several axes (size, ab initio vs. active sampling, standard vs. biased dynamics, etc.) we consider a transfer from simpler to more complicated systems. Successfully transferring from simpler to more complicated systems has real-world impact as it can alleviate the high computational cost required for labeling via DFT calculations by reducing the number of points. From Table 2 we see that MEKRR-(SCN) has the lowest error in four out of five tasks, while MEKRR-(SchNet) has the lowest error in the remaining task. Furthermore, the relative transferability of MEKRR compared to the other methods even improves as the task becomes harder. To this respect, it is worth noting that D 1 , D 2 , D 3 are qualitatively similar, being all composed of 5 layers of Fe and differing for the sampling method used. Instead, the atomic environments contained in D 4 are different as they refer to a slab with a different number of layers. This explains the different order of magnitudes in the last two columns. Despite this, MEKRR still performs very well compared to the baselines.

5.5 Leveraging MEKRR beyond supervised learning

In the previous sections, we have shown how MEKRR performs very well on both supervised and transfer learning tasks. The effectiveness comes from the combination of a pre-trained feature map together with the K α kernel. However, this idea is not restricted to supervised learning. We can indeed leverage the similarity measure provided by the kernel for tasks beyond potential energy regression. As a simple example, in Fig. 3 we plot the kernel matrix, when using SchNet as the feature map, of a part of the trajectory of D 2 containing an N-N bond breaking event in the cases α = 0 and α = 1 . In both images we can see a clear structure that highlights at least two distinct states, but with the second heatmap having more signal. We can then use the kernel to perform spectral clustering with two classes, the result is visualized on the top margin of the heatmap along with the time evolution of a physical quantity that signals the N-N bond-breaking. This facilitates a

d(N, N) [A]

[V] (NND)P

0

0

1.6

F

0.62

150

1.4

Similarity

0.60

1.2

300

0.58

1.0

450

0.8

0.56

0

150

300

450

0

150

300

450

t

t

physical interpretation, as the two classes correspond to configurations containing the reactants (the N 2 molecule) and products (two N atoms) of the chemical reaction. Interestingly, the α = 1 case correlates more closely with the handpicked physical quantity, shown in the top panel. The reason for this is the fact that learning the weights independently gives more weight to chemical species that are under-represented, which typically correspond to adsorbed atoms in surfaces. This allows us to give more weight to the most important actors in catalytic applications.

6 Conclusion and future work

In this work, we introduced an approach to model the potential energy surface of atomistic systems. Our method employs GNN representations trained on the large OC20 dataset along with kernel ridge regression, which we tested on two catalytic processes that are outside the distribution of the pre-training dataset. We devised a kernel function incorporating GNN features, blending kernel mean embedding with information related to the atom's chemical species. Our approach outperforms standalone GNN or kernel methods, demonstrating impressive transferability. This suggests promising avenues for transfer learning application in computational chemistry. However, we recognize certain limitations. Firstly, our method is based upon KRR which scales poorly to large scale datasets, although potential ways around this such as random features [70] or Nyström approximations [71] can overcome this limitation. Secondly, although we tested MEKRR on out-of-distribution datasets for representations pre-trained on OC20, we still focused our analysis on catalytic systems similar to those in the dataset. In this regard, it would be interesting to understand the extent to which MEKRR can predict well the chemical properties of generic systems. In addition, it would be important to incorporate forces into the loss function in order to use it in molecular dynamics applications. We believe that addressing these aspects will further improve the impact of this framework in computational chemistry.

CRediT author statement J. I. Falk : Conceptualization, Methodology, Software, Investigation, Formal analysis, Writing - Original Draft; L. Bonati : Conceptualization, Methodology, Formal analysis, Visualization, Writing - Original Draft; P. Novelli : Conceptualization, Methodology, Writing - Original Draft; M. Parrinello : Conceptualization, Writing - Review & Editing, Supervision; M. Pontil : Conceptualization, Methodology, Writing - Review & Editing, Supervision.

Acknowledgements Weacknowledge the financial support from the PNRR MUR Project PE000013 CUP J53C22003010006 "Future Artificial Intelligence Research (FAIR)", funded by the European Union - NextGenerationEU, EU Project ELIAS under grant agreement No. 101120237, and the "Joint project TransHyDE_FP3: Reforming ammonia - transport of H 2 via derivatives", funded from the German Federal Ministry of Research (BMBF), funding code: 03HY203A-F.

References

Daan Frenkel and Berend Smit. Understanding Molecular Simulation: From Algorithms to Applications , volume 1. Elsevier, 2001.

J. E. Lennard-Jones. Cohesion. Proceedings of the Physical Society , 43(5):461, 1931.

Roberto Car and Michele Parrinello. Unified approach for molecular dynamics and densityfunctional theory. Physical Review Letters , 55(22):2471, 1985.

Dominik Marx and Jurg Hutter. Ab initio molecular dynamics: Theory and implementation. Modern Methods and Algorithms of Quantum Chemistry , 1(301-449):141, 2000.

Walter Kohn and Lu Jeu Sham. Self-consistent equations including exchange and correlation effects. Physical review , 140(4A):A1133, 1965.

Jörg Behler and Michele Parrinello. Generalized neural-network representation of highdimensional potential-energy surfaces. Physical Review Letters , 98(14):146401, 2007.

Oliver T. Unke, Stefan Chmiela, Huziel E. Sauceda, Michael Gastegger, Igor Poltavsky, Kristof T. Schütt, Alexandre Tkatchenko, and Klaus-Robert Müller. Machine Learning Force Fields. Chemical Reviews , 121(16):10142-10186, 2021.

Jörg Behler and Gábor Csányi. Machine learning potentials for extended systems: a perspective. The European Physical Journal B , 94:1-11, 2021.

Albert P. Bartók, Risi Kondor, and Gábor Csányi. On representing chemical environments. Physical Review B , 87(18):184115, 2013.

Michael J Willatt, Félix Musil, and Michele Ceriotti. Atom-density representations for machine learning. The Journal of Chemical Physics , 150(15):154110, 2019.

Felix Musil, Andrea Grisafi, Albert P Bartók, Christoph Ortner, Gábor Csányi, and Michele Ceriotti. Physics-inspired structural representations for molecules and materials. Chemical Reviews , 121(16):9759-9815, 2021.

Ilyes Batatia, David P Kovacs, Gregor Simm, Christoph Ortner, and Gábor Csányi. Mace: Higher order equivariant message passing neural networks for fast and accurate force fields. Advances in Neural Information Processing Systems , 35:11423-11436, 2022.

Albert Musaelian, Simon Batzner, Anders Johansson, Lixin Sun, Cameron J Owen, Mordechai Kornbluth, and Boris Kozinsky. Learning local equivariant representations for large-scale atomistic dynamics. Nature Communications , 14(1):579, 2023.

Larry Zitnick, Abhishek Das, Adeesh Kolluru, Janice Lan, Muhammed Shuaibi, Anuroop Sriram, Zachary Ulissi, and Brandon Wood. Spherical channels for modeling atomic interactions. Advances in Neural Information Processing Systems , 35:8054-8067, 2022.

Simon Batzner, Albert Musaelian, Lixin Sun, Mario Geiger, Jonathan P. Mailoa, Mordechai Kornbluth, Nicola Molinari, Tess E. Smidt, and Boris Kozinsky. E(3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials. Nature Communications , 13(1):2453, 2022.

Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E(n) Equivariant Graph Neural Networks, 2022.

Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. In Jill Burstein, Christy Doran, and Thamar Solorio, editors, Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACL-HLT 2019, Minneapolis, MN, USA, June 2-7, 2019, Volume 1 (Long and Short Papers) , pages 4171-4186. Association for Computational Linguistics, 2019.

Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language Models are Few-Shot Learners. In Advances in Neural Information Processing Systems , volume 33, pages 1877-1901, 2020.

Eleni Triantafillou, Tyler Zhu, Vincent Dumoulin, Pascal Lamblin, Utku Evci, Kelvin Xu, Ross Goroshin, Carles Gelada, Kevin Swersky, Pierre-Antoine Manzagol, and Hugo Larochelle. Meta-Dataset: A Dataset of Datasets for Learning to Learn from Few Examples. In International Conference on Learning Representations , 2020.

Xiaohua Zhai, Joan Puigcerver, Alexander Kolesnikov, Pierre Ruyssen, Carlos Riquelme, Mario Lucic, Josip Djolonga, Andre Susano Pinto, Maxim Neumann, Alexey Dosovitskiy, Lucas Beyer, Olivier Bachem, Michael Tschannen, Marcin Michalski, Olivier Bousquet, Sylvain Gelly, and Neil Houlsby. A Large-scale Study of Representation Learning with the Visual Task Adaptation Benchmark, 2020.

Rishi Bommasani, Drew A. Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney von Arx, Michael S. Bernstein, Jeannette Bohg, Antoine Bosselut, Emma Brunskill, Erik Brynjolfsson, Shyamal Buch, Dallas Card, Rodrigo Castellon, Niladri Chatterji, Annie Chen, Kathleen Creel, Jared Quincy Davis, Dora Demszky, Chris Donahue, Moussa Doumbouya, Esin Durmus, Stefano Ermon, John Etchemendy, Kawin Ethayarajh, Li Fei-Fei, Chelsea Finn, Trevor Gale, Lauren Gillespie, Karan Goel, Noah Goodman, Shelby Grossman, Neel Guha, Tatsunori Hashimoto, Peter Henderson, John Hewitt, Daniel E. Ho, Jenny Hong, Kyle Hsu, Jing Huang, Thomas Icard, Saahil Jain, Dan Jurafsky, Pratyusha Kalluri, Siddharth Karamcheti, Geoff Keeling, Fereshte Khani, Omar Khattab, Pang Wei Koh, Mark Krass, Ranjay Krishna, Rohith Kuditipudi, Ananya Kumar, Faisal Ladhak, Mina Lee, Tony Lee, Jure Leskovec, Isabelle Levent, Xiang Lisa Li, Xuechen Li, Tengyu Ma, Ali Malik, Christopher D. Manning, Suvir Mirchandani, Eric Mitchell, Zanele Munyikwa, Suraj Nair, Avanika Narayan, Deepak Narayanan, Ben Newman, Allen Nie, Juan Carlos Niebles, Hamed Nilforoshan, Julian Nyarko, Giray Ogut, Laurel Orr, Isabel Papadimitriou, Joon Sung Park, Chris Piech, Eva Portelance, Christopher Potts, Aditi Raghunathan, Rob Reich, Hongyu Ren, Frieda Rong, Yusuf Roohani, Camilo Ruiz, Jack Ryan, Christopher Ré, Dorsa Sadigh, Shiori Sagawa, Keshav Santhanam, Andy Shih, Krishnan Srinivasan, Alex Tamkin, Rohan Taori, Armin W. Thomas, Florian Tramèr, Rose E. Wang, William Wang, Bohan Wu, Jiajun Wu, Yuhuai Wu, Sang Michael Xie, Michihiro Yasunaga, Jiaxuan You, Matei Zaharia, Michael Zhang, Tianyi Zhang, Xikun Zhang, Yuhui Zhang, Lucia Zheng, Kaitlyn Zhou, and Percy Liang. On the Opportunities and Risks of Foundation Models, 2022.

Lowik Chanussot, Abhishek Das, Siddharth Goyal, Thibaut Lavril, Muhammed Shuaibi, Morgane Riviere, Kevin Tran, Javier Heras-Domingo, Caleb Ho, Weihua Hu, Aini Palizhati, Anuroop Sriram, Brandon Wood, Junwoong Yoon, Devi Parikh, C. Lawrence Zitnick, and Zachary Ulissi. The Open Catalyst 2020 (OC20) Dataset and Community Challenges. ACS Catalysis , 11(10):6059-6072, 2021.

Albert P. Bartók, Mike C. Payne, Risi Kondor, and Gábor Csányi. Gaussian Approximation Potentials: The Accuracy of Quantum Mechanics, without the Electrons. Physical Review Letters , 104(13):136403, 2010.

Stefan Chmiela, Alexandre Tkatchenko, Huziel E. Sauceda, Igor Poltavsky, Kristof T. Schütt, and Klaus-Robert Müller. Machine learning of accurate energy-conserving molecular force fields. Science Advances , 3(5):e1603015, 2017.

Stefan Chmiela, Huziel E. Sauceda, Igor Poltavsky, Klaus-Robert Müller, and Alexandre Tkatchenko. sGDML: Constructing accurate and data efficient molecular force fields using machine learning. Computer Physics Communications , 240:38-45, 2019.

Linfeng Zhang, Jiequn Han, Han Wang, Wissam Saidi, Roberto Car, et al. End-to-end symmetry preserving inter-atomic potential energy model for finite and extended systems. Advances in Neural Information Processing Systems , 31, 2018.

Kristof Schütt, Pieter-Jan Kindermans, Huziel Enoc Sauceda Felix, Stefan Chmiela, Alexandre Tkatchenko, and Klaus-Robert Müller. SchNet: A continuous-filter convolutional neural network for modeling quantum interactions. In Advances in Neural Information Processing Systems , volume 30. Curran Associates, Inc., 2017.

Johannes Klicpera, Janek Groß, and Stephan Günnemann. Directional Message Passing for Molecular Graphs. In Eighth International Conference on Learning Representations , 2020.

Niklas Frederik Schmitz, Klaus-Robert Müller, and Stefan Chmiela. Algorithmic Differentiation for Automated Modeling of Machine Learned Force Fields. The Journal of Physical Chemistry Letters , 2022.

Liam Collins, Aryan Mokhtari, Sewoong Oh, and Sanjay Shakkottai. MAML and ANIL Provably Learn Representations. In Proceedings of the 39th International Conference on Machine Learning , pages 4238-4310. PMLR, 2022.

John Isak Texas Falk, Carlo Ciliberto, and Massimiliano Pontil. Implicit kernel meta-learning using kernel integral forms. In Proceedings of the Thirty-Eighth Conference on Uncertainty in Artificial Intelligence , pages 652-662. PMLR, 2022.

Ruohan Wang, Massimiliano Pontil, and Carlo Ciliberto. The role of global labels in few-shot classification and how to infer them. Advances in Neural Information Processing Systems , 34:27160-27170, 2021.

Ruohan Wang, Isak Falk, Massimiliano Pontil, and Carlo Ciliberto. Robust meta-representation learning via global label inference and classification. IEEE Transactions on Pattern Analysis and Machine Intelligence , 2023.

Guneet Singh Dhillon, Pratik Chaudhari, Avinash Ravichandran, and Stefano Soatto. A Baseline for Few-Shot Image Classification. In International Conference on Learning Representations , 2023.

Jake Snell, Kevin Swersky, and Richard Zemel. Prototypical Networks for Few-shot Learning. In Advances in Neural Information Processing Systems , volume 30, 2017.

Luca Bertinetto, Joao F. Henriques, Philip Torr, and Andrea Vedaldi. Meta-learning with differentiable closed-form solvers. In International Conference on Learning Representations , 2023.

Noah Hoffmann, Jonathan Schmidt, Silvana Botti, and Miguel AL Marques. Transfer learning on large datasets for the accurate prediction of material properties. arXiv preprint arXiv:2303.03000 , 2023.

Viktor Zaverkin, David Holzmüller, Luca Bonfirraro, and Johannes Kästner. Transfer learning for chemically accurate interatomic neural network potentials. Physical Chemistry Chemical Physics , 25(7):5383-5396, 2023.

Alice EA Allen, Nicholas Lubbers, Sakib Matin, Justin Smith, Richard Messerly, Sergei Tretiak, and Kipton Barros. Learning together: Towards foundational models for machine learning interatomic potentials with meta-learning. arXiv preprint arXiv:2307.04012 , 2023.

Duo Zhang, Hangrui Bi, Fu-Zhi Dai, Wanrun Jiang, Linfeng Zhang, and Han Wang. Dpa-1: Pretraining of attention-based deep potential model for molecular simulation. arXiv preprint arXiv:2208.08236 , 2022.

Tian Xie and Jeffrey C. Grossman. Crystal Graph Convolutional Neural Networks for an Accurate and Interpretable Prediction of Material Properties. Physical Review Letters , 120(14):145301, 2018.

Kristof T. Schütt, Oliver T. Unke, and Michael Gastegger. Equivariant message passing for the prediction of tensorial properties and molecular spectra, 2021.

Bernhard Schölkopf and Alexander Johannes Smola. Learning with Kernels: support vector machines, regularization, optimization, and beyond . Adaptive computation and machine learning series. MIT Press, 2002.

Ingo Steinwart. Support Vector Machines . Information Science and Statistics. Springer, New York, NY, 2008.

Lei Shi, Xin Guo, and Ding-Xuan Zhou. Hermite learning with gradient data. Journal of Computational and Applied Mathematics , 233(11):3046-3059, 2010.

Carl Edward Rasmussen. Gaussian processes in machine learning. In Advanced Lectures on Machine Learning , pages 63-71. Springer Berlin Heidelberg, 2004.

Krikamol Muandet, Kenji Fukumizu, Bharath Sriperumbudur, and Bernhard Schölkopf. Kernel Mean Embedding of Distributions: A Review and Beyond. Foundations and Trends® in Machine Learning , 10(1-2):1-141, 2017.

Michael M. Bronstein, Joan Bruna, Taco Cohen, and Petar Veliˇ ckovi´ c. Geometric Deep Learning: Grids, Groups, Graphs, Geodesics, and Gauges, 2021.

Franco Scarselli, Marco Gori, Ah Chung Tsoi, Markus Hagenbuchner, and Gabriele Monfardini. The graph neural network model. IEEE Transactions on Neural Networks , 20(1):61-80, 2008.

Alessio Micheli. Neural network for graphs: A contextual constructive approach. IEEE Transactions on Neural Networks , 20(3):498-511, 2009.

Jie Zhou, Ganqu Cui, Shengding Hu, Zhengyan Zhang, Cheng Yang, Zhiyuan Liu, Lifeng Wang, Changcheng Li, and Maosong Sun. Graph neural networks: A review of methods and applications. AI Open , 1:57-81, 2020.

Kristof T. Schütt, Huziel E. Sauceda, Pieter-Jan Kindermans, Alexandre Tkatchenko, and Klaus-Robert Müller. SchNet - a deep learning architecture for molecules and materials. The Journal of Chemical Physics , 148(24):241722, 2018.

Johannes Gasteiger, Florian Becker, and Stephan Günnemann. GemNet: Universal Directional Graph Neural Networks for Molecules, 2022.

Weihua Hu, Muhammed Shuaibi, Abhishek Das, Siddharth Goyal, Anuroop Sriram, Jure Leskovec, Devi Parikh, and C. Lawrence Zitnick. ForceNet: A Graph Neural Network for Large-Scale Quantum Calculations, 2021.

Kristof T. Schütt, Farhad Arbabzadah, Stefan Chmiela, Klaus R. Müller, and Alexandre Tkatchenko. Quantum-chemical insights from deep tensor neural networks. Nature Communications , 8(1):13890, 2017.

Johannes Gasteiger, Muhammed Shuaibi, Anuroop Sriram, Stephan Günnemann, Zachary Ulissi, C. Lawrence Zitnick, and Abhishek Das. GemNet-OC: Developing Graph Neural Networks for Large and Diverse Molecular Simulation Datasets, 2022.

Theodoros Evgeniou and Massimiliano Pontil. Regularized multi-task learning. In Proceedings of the 10th ACM SIGKDD International Conference on Knowledge Discovery and Data Mining , pages 109-117, 2004.

Theodoros Evgeniou, Charles A. Micchelli, and Massimiliano Pontil. Learning multiple tasks with kernel methods. Journal of Machine Learning Research , 6(21):615-637, 2005.

K. T. Schütt, H. E. Sauceda, P.-J. Kindermans, A. Tkatchenko, and K.-R. Müller. SchNet A deep learning architecture for molecules and materials. The Journal of Chemical Physics , 148(24):241722, 2018.

Albert P. Bartók, Risi Kondor, and Gábor Csányi. On representing chemical environments. Physical Review B , 87(18):184115, 2013.

James R. Kermode. F90wrap: An automated tool for constructing deep Python interfaces to modern Fortran codes. Journal of Physics: Condensed Matter , 32(30):305901, 2020.

G. Csanyi, S. Winfield, J. Kermode, M. C. Payne, A. Comisso, A. De Vita, and N. Bernstein. Expressive Programming for Computational Physics in Fortran 950+. Newsletter of the Computational Physics Group , pages 1-24, 2007.

Damien Garreau, Wittawat Jitkrittum, and Motonobu Kanagawa. Large sample analysis of the median heuristic, 2018.

Alessandro Laio and Michele Parrinello. Escaping free-energy minima. Proceedings of the national academy of sciences , 99(20):12562-12566, 2002.

Graeme Henkelman, Blas P Uberuaga, and Hannes Jónsson. A climbing image nudged elastic band method for finding saddle points and minimum energy paths. The Journal of chemical physics , 113(22):9901-9904, 2000.

Luigi Bonati, Daniela Polino, Cristina Pizzolitto, Pierdomenico Biasi, Rene Eckert, Stephan Reitmeier, Robert Schlögl, and Michele Parrinello. The role of dynamics in heterogeneous catalysis: Surface diffusivity and n2 decomposition on fe (111). Proceedings of the National Academy of Sciences , 120(50):e2313023120, 2023.

Manyi Yang, Luigi Bonati, Daniela Polino, and Michele Parrinello. Using metadynamics to build neural network potentials for reactive events: the case of urea decomposition in water. Catalysis Today , 387:143-149, 2022.

Karl Weiss, Taghi M. Khoshgoftaar, and DingDing Wang. A survey of transfer learning. Journal of Big Data , 3(1):9, 2016.

Justin S. Smith, Benjamin T. Nebgen, Roman Zubatyuk, Nicholas Lubbers, Christian Devereux, Kipton Barros, Sergei Tretiak, Olexandr Isayev, and Adrian E. Roitberg. Approaching coupled cluster accuracy with a general-purpose neural network potential through transfer learning. Nature Communications , 10(1):2903, 2019.

Ali Rahimi and Benjamin Recht. Random features for large-scale kernel machines. In Proceedings of the 20th International Conference on Neural Information Processing Systems , NIPS'07, page 1177-1184, Red Hook, NY, USA, 2007. Curran Associates Inc.

Giacomo Meanti, Luigi Carratino, Lorenzo Rosasco, and Alessandro Rudi. Kernel methods through the roof: Handling billions of points efficiently. In Proceedings of the 34th International Conference on Neural Information Processing Systems , NIPS'20, Red Hook, NY, USA, 2020. Curran Associates Inc.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn this work, the authors investigate the use of kernel ridge regression with feature maps of atomic environments from pre-trained graph neural network models for predicting energies of atomistic systems. They show that their method, mean embedding kernel ridge regression (MEKRR), outperforms or is competitive with pre-trained and fine-tuned SchNet models trained on the OC20 dataset. Their kernel-based method is interpretable and gives insight into the chemical evolution of systems in trajectories.\n\nSTRENGTHS:\nThe paper is well written and clear, giving a nice overview of both GNN- and kernel-based approaches to modeling potential energy surfaces.\n\nThe method is explained well and well-motivated by the outstanding problems related to transferability in the field of machine learned interatomic potentials.\n\nWEAKNESSES:\nThe authors only evaluate their method against one GNN, SchNet; there is a whole zoo of modern GNNs purpose-built for learning potential energy surfaces that incorporate additional physical priors and are fast and easy to train.\n\nOne key issue with neural interatomic potentials is transferability across geometric environments; it is not clear that MEKRR addresses this challenge in any way, as it is ultimately reliant on pre-computed feature maps.\n\nQUESTIONS:\nDo the authors think that scalability will be an issue with their method in practical workflows, i.e. for the million - billion scale atomistic simulations for solid state and biomolecular systems that are now being done with neural interatomic potentials?\n\nWhy is SchNet used as a baseline rather than Allegro, MACE, Spherical Channel Network, or some of the many other GNNs that are commonly used to model the OC20 dataset? What advantages does MEKRR provide over modern GNN approaches that can be quickly and accurately fine tuned?\n\nLIMITATIONS:\nLimitations are addressed in part, but further discussion is needed related to more modern ML interatomic potentials.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes a transfer learning algorithm that leverages the ability of graph neural networks (GNNs) together with kernel mean embeddings. A flexible kernel function that incorporates chemical species information is proposed to improve performance and interoperability. Empirical results in a series of realistic datasets of increasing complexity are provided to illustrate the generalization and transferability performance of the proposed method, MEKRR.\n\nSTRENGTHS:\n* MEKRR achieves excellent generalization and transferability performance on the introduced datasets with increasing complexity.\n* The paper is well-organized and clearly written.\n\nWEAKNESSES:\n* The proposed mean embedding kernel ridge regression (MEKRR) is claimed to have better generalization and transferability, however, no substantial theoretical evidence or intuition is provided to motivate this claim. Consequently, it is unclear why MEKRR could achieve better performance than other methods.\n* MEKRR cannot directly calculate the molecular force as the gradient flow of pre-trained feature maps is stopped, which limits its application in real-world scenarios.\n* In inference, it appears that kernel-based methods such as MEKRR usually need to evaluate the kernel function of a target molecule with all molecules in the training set. This linear scaling computational cost can be inefficient for large-scale datasets, in contrast to most GNN-based methods, which have a fixed inference complexity.\n* Concerns exist regarding experimental evaluation:\n    * The molecule scales in the proposed datasets are relatively small, making it difficult to justify the generalization capacity of MEKRR on a broader range of large molecules. Additionally, the limited data size negatively affects the performance of GNNs, meaning that comparisons on small-scale datasets may not adequately reveal the difference between MEKRR and SchNet.\n    * Comparison with only the SchNet model is insufficient, as SchNet is not the state-of-the-art (SOTA) model in the S2EF task of the OCP20 dataset [https://github.com/Open-Catalyst-Project/ocp/blob/main/MODELS.md]. It is also unclear whether SchNet is representative enough to benchmark the generalization and transferability of those SOTA models. Thus, it is recommended that the authors include more pretrained GNN models to further verify the effectiveness of MEKRR.\n    * The criterion for evaluating the generalization and transferability is unclear, making it difficult to justify the significance of the empirical results. As shown in Table 1-2, although MEKERR exhibits better performance than the baseline methods, there exist large gaps in RMSE/MSE metrics between in-distribution and out-of-distribution settings. For instance, RMSE in $D_1\\rightarrow D_2$ is 80 times larger than RMSE in $D_1$. Is the scale of RMSE in $D_1\\rightarrow D_2$ acceptable in real chemistry research?\n\nQUESTIONS:\n* In the main text, Equation 7 mentioned in line 205 is not found. Could the authors clarify this further?\n* Establishing MEKRR only on SchNet does not adequately support the universality of the proposed algorithm. It is recommended that the authors include more pretrained GNN models in the experimental evaluation.\n* Could the authors provide more clarification about the criterion for evaluating the generalization and transferability of a machine learning-based PES model? This would make the contribution of MEKRR clearer.\n\nLIMITATIONS:\nThe authors well clarify the limitations and potential negative societal impact of their work in the paper.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-06"}

comments: SUMMARY:
The authors propose to use embedding of GNNs pretrained on large and diverse molecular datasets (in their case OC20) to construct kernel embeddings for efficient transfer learning.

STRENGTHS:
This paper was a delight to read. It solves a real problem with a clever idea, implemented well, and I learned something. The presentation is excellent and I have to say, it is incredibly refreshing to read a NeurIPS paper on AI4Science by authors that clearly understand the physics and chemistry behind what they are doing.

WEAKNESSES:
The authors mention meta-learning a few times in the intro section, but meta-learning is a big field and can mean many things. This is hard to understand for the reader without looking up the references. 

Given the success of equivariant models over invariant ones like SchNet it would super interesting to see how to extend this approach to a setting where the kernel embedding is constructed from embeddings made of a direct sum of irreducible representations. The baselines chosen here are rather weak and were maybe state-of-the-art 5 years ago (I understand that pretrained equivariant potentials on OC20 may not have been available when the authors started this project and would've been computationally difficult to obtain).

QUESTIONS:
N/A

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-07-08

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents a transfer learning algorithm for interatomic potentials, leveraging GNNs and kernel mean embeddings. It overcomes data and computational challenges, achieving generalization and transferability on realistic datasets.\n\nSTRENGTHS:\nThe benchmark provided in this paper is relatively comprehensive. However, to ensure reproducibility, it would be ideal for the authors to make the complete code available.\n\nWEAKNESSES:\nThe authors utilize pre-trained GNN embeddings but only discuss a limited number of SOTA GNN methods in atomistic simulations. It is crucial for the paper to clearly highlight the novelty and enhancements of their proposed method compared to existing GNN works for atomistic simulations [1]. Additionally, conducting an ablation study in the experimental section would be essential to demonstrate the innovative aspects of their approach. Simply applying a regular GNN method does not suffice as a significant contribution.\n\nThe methodology section seems odd as it contains only one subsection (4.1) despite the authors' claim of multiple contributions. It is crucial for the authors to explicitly clarify which aspects of their work are entirely new methods and which are applications of existing methods. This distinction will help readers better understand and appreciate the original contributions made in the paper.\n\n\n[1] ​​Sriram, Anuroop, et al. \"Towards training billion parameter graph neural networks for atomic simulations.\" arXiv preprint arXiv:2203.09697 (2022).\n\nQUESTIONS:\nCould the inclusion of a flowchart improve readers' ability to follow the pipeline of this paper? I found some difficulty in comprehending the research process while reading through the manuscript.\n\nIn the motivation, the authors mentioned the importance of rotational and translational invariance. Therefore, it would be relevant to discuss methods like SE3-equivariant approaches that handle these invariances.\n\nLIMITATIONS:\nThe paper lacks in-depth discussions on limitations, such as issues regarding model robustness and ethical considerations, particularly regarding data misuse. It would be beneficial for the authors to address these aspects, as they are crucial in understanding the broader implications and potential challenges associated with the proposed models. Including a thoughtful exploration of these topics would enhance the completeness and relevance of the paper.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-28"}