Transfer learning for atomistic simulations using GNNs and kernel mean embeddings
Authors: John Isak Texas Falk
Luigi Bonati
Pietro Novelli
Michele Parrinello
Massimiliano Pontil
Venue: NeurIPS 2023 poster (2023)
Keywords: GNN
Mean Embedding
Kernels
Atomistic Simulations
OCP
Transfer Learning
Molecular Dynamics
Kernel Ridge Regression
Neural Networks
TL;DR: We show that combining pretrained GNNs on the OCP dataset with kernel mean embeddings transfer to other chemical datasets, laying the groundwork for foundation models in chemistry

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
John I. Falk CSML
Luigi Bonati
Pietro Novelli CSML
Atomistic Simulations
Istituto Italiano di Tecnologia Genova, Italy me@isakfalk.com
Istituto Italiano di Tecnologia
Istituto Italiano di Tecnologia
Genova, Italy
Genova, Italy
Sogg.bbonat@iit..t.
it
Michele Parrinello
Massimiliano Pontil
Atomistic Simulations
CSML
Istituto Italiano Tecnologia
Istituto Italiano di Tecnologia
Genova, Italy
Genova, Italy
mi wwwchheeppprrnnellllllooooo
University College London, U.K.
Abstract
Interatomic potentials learned using machine learning methods have been suc- cessfully 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 feature map from GNNs pre- rre-trained on the OC20 dataset and use it to learn the potential energy surface from system- -estem--pecificc datasets catalytic processes. Our method is further enhanced by incorporating into the kernel the chemical species information, result- ing 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 similar fine-tuning approaches.
Introduction
Atomistic simulations have become 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 system of atoms by integrating Newton's 's equation of motion [I] The forces acting on the particles are determined from model. for the interactions, called the potential energy surface (PES), whose accuracy the reliability of the simulation depends. For a long time, the interactions were modeled in a rather empirical, and of course, very accurate way 21 A significant step forward made with the introduction of ab initio molecular dynamics 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 g] They regress the
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

potential energy, as 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 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 1 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) -Eqiivaaqi layers [12 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 solve downstream tasks for which only 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.22 catalytic systems, totaling over 260M data points, We show how, by exploiting the availability of this heterogeneous dataset, it is possible 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 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 large datasets with fine-tuning via kernel mean embeddings. This combination allows to satisfy the physical symmetries inherent 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 facili- tates 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-distributin with respect to the GNN pre- trained representation.
Related work There is 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 223 Later, it was proposed to model with neural networks the descriptors 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 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 preprocessing step where each image is mapped through a meta-learned or pre-trained representation, e.g. (30-36). In particular, [36] employs kernel ridge regression (KRR) on top of 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-3) or to exploit the release of large and heterogeneous datasets such OC20 [40].
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 we introduce our method MEKRR. Section 5 we validate our method on variety of realistic datasets of increasing complexity. Finally, in Section wee we conclude and outline future directions.
Organization In Section we
2

---

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-hot- encoding over [S], where [S] {I....,... S...
Each atom described by its Euclidean position r E and its chemical species z E {0, 1]s We denote state of the system by I = (Ti, 2 = (R, 2) E X using the design matrices R E R and Z E {0, 1]
The quantity we want to regress is the potential energy, which is scalar function E X' R. The target values E, 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 sequence of T atomic configurations (m=))=1 together with the corresponding labels (E,)==-
The physics of the system posits that any estimator of the energy f X - certain invariant properties reflecting the underlying physical symmetries. These invariances are ato-trraslational invariance with respect to the position vectors and permutation inaiiiionnivviancee within each group of chemical species. In addition, the energy estimator should be curl-free and smooth (27). 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.
should possess
2.2 Objective function
In order to learn potential energies, one considers models J from set of functions (e.g., repro- ducing kernel Hilbert space or neural network functions) parameterized by w and fitted to given dataset (I,, Fe) F This is achieved by minimizing objective function, which can be split into 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
Sm 2,) + 71 Vo, (.)]( + >N((w), R(w) = Σ (11 (m.
(1)
for some regularizer Ω and regularization parameter À > 0 and a loss weight E [0, 1]. A loss weight 7 * O used whenever the forces F(x,r) r) = - V, are regressed alongside the energy. Typically we will use ((w) = ||w|| as the regularizer, but in Section 4.3 we extend this to incorporate additional information. Since in the experiments discussed below focus energy prediction, let = 0.
3 Kernel ridge regression and kernel mean embeddings
Our method is designed around kernel methods, a well-established tool the heart of most non- parametric machine learning algorithms [3 44). A kernel is any positive definite function (see (see[4) Definition 4.15) K X x X + R on the input space X. For any kernel exists a feature space H. and a feature map o: X + HK such that K ((2, x²) = (o(x), o(x') for all x. x' E X.
Kernel ridge regression is supervised learning algorithm that parametrize the dependence between inputs x E X' and scalar outputs E R as y = (w, Φ(x)) for some vector w in the feature space. Given dataset (z=)/= E X of T input points, and corresponding dataset of outputs (v+) RT. KRR learns a functional relation between inputs and outputs by solving the least squares minimization problem
argmin WEH
Σ ((w, o(x1)) - y1)² Allwwl22x S
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 the training points, that is = I (12(2)) for some (c)) E TRT,. The representer theorem and the fact that K K(x, = (o(z), o(x')) also imply that predicting any new point with KRR is simple evaluating
3

---

(rb, ((x)) regress the gradient of y with respect to z, a technique used e.g. 29 concurrently learn potential energies and forces.
o(x))= E (0()z)), o(x)) = =E= (2,,
x). KRR can also be readily extended 45 to
In the standard case in which the input space is subset of Euclidean space X' C Rd many kernels have been designed; see, for instance, [4, Section 4.2. This notwithstanding, kernel functions can be defined on arbitrary imput 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 design kernels able preserve such symmetry. We now introduce the concept of kernel mean embeddings 47 which allow us to define permutationalll-inaarri kernels over atomic configurations.
Kernel mean embeddings consider the case in which inputs J are sets of points living in an Euclidean space {ri Rd S n). For example, atomic configuration is set containing the position of each atom in the system (d = 3). For any feature map (i(r) and the corresponding kernel function k(r, r') = (v)()), (rr)) on points r Id we can define feature map acting on the whole set a as
<(x) C+ Σ (T:))
with C2 > 0 positive normalization constant. We note in passing that the value of C2 depends on the chemical property that we want to learn. Indeed, setting C₂ = returns extensive properties such the potential energy, while C2 1/n is appropriate to model intensive ones. This allows to define kernels on sets x as
K K(x, = (0(x),
C.C. MER
(o6)). C.Coe b(rrrry) NEB
E
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 principled and efficient method to learn potential energy surfaces from data.
Method
Our method relies pre-trained GNN representations of chemical environments and uses them as feature maps to define a kernel mean embedding do so by leveraging pre-trained feature maps the OC20 dataset [22]. Once kernel based mean embeddings defined, we can use KRR to regress the potential energy from sampled data (see Fig. [ for diagram of the method). In this respect, 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 show the performance of pre-trained representations based on SCN 14) and SchNet [27]. any other GNN architecture is perfectly valid choice, and can be used in place of ours without the need of further adjustments.
(3) acting on atomic configurations. In practice, we
note that method shares similarities with the
4.1 GNN representations of chemical environments
SCN and SchNet are instances of graph neural networks (GNNs) 48 50 a class architectures designed learn mappings over graphs. A graph edges & C V V between pairs of nodes. Representing the chemical environment with GNNs involves a preprocessing step that turns configuration æ into graph by associating each atom to a node and constructing the edges according to the matrix of pairwise distances between atoms. Each node is then initialized to feature vector h; encoding the chemical species of the atom via (possibly learnable) embedding layer. Each GNN layer then updates the node features via a nonlinear mnesagee-pssingg scheme
(V, E) is a collection of nodes V and
h + gb,,,
Σ JEN(()
no (hs, 1t)))
Here, Aiii) the neighborhood of i i.e. the set of nodes connected to i through an edge, go is Icarnable vector- valuee function such MLP and ne is message- saassage-pssiing function. A graph

---

H
K(H, H")
H
hgs
Atomic feature maps from pretrained GNN
b. Transfer learning via KRR
Figure I: Diagram of MEKRR.
neural network is formed by concatenating multiple message-passing layers. The weights of the embedding and GNN layers are then learned an end-to-end fashion typically using first-order optimization.
Among the many GNN architectures available [51] growing number is being developed to specialize in computational chemistry settings, see for instance [6
42 52 56
4.2 Our method: mean embeddings of GNN representations
We ready combine together the techniques deseribed far and describe in full detail our method. Starting with an atomic configuration æ, we pass to a pre-trained GNN representation to get collection of node features
I ==))):=(((<))))))
In the case of equivariant GNNs we additionally need to extract invariant features from equivariant ones via a pooling method such as averaging. case of SCN [14), for example, we average the We then evaluate the kernel mean embedding of the
features over the sphere (see Appendix C.2 node features
H(x) + (H(()) C2
Σ
(hh(())).
(x()))(())
Finally, at training time we can use the kernel mean embedding (H(x)) to solve the KRR problem
Σ (ww (H(()) E.)² Allw|l }
= argmin wEHs
while inference time can get estimation for the potential energy by simply evaluating (10, of (x)))))
Some practical considerations are now in order. Given 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 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 many strategies solve the KRR problem (4) based e.g. the solution of linear system or on -redient-descent. In Algorithm we report the pscudo-code describing our implementation of the training and prediction steps of MEKRR.
4.3 Chemically-informn kernel mean embeddings
The chemical species each atom 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 configuration let now (x(x) (h.(z) species of i s} be the subset

---

Algorithm Mean Embedding Kernel Ridge Regression (MEKRR)
Parameters: Training dataset (Iz, Fu) of configurations 2 and potential energies E,. Ker-
nel function k. Interpolation parameter of the multi-weight formulation a E [0, 1]. Tikhonov regularization >> Pre-trained GNN feature map x + (z).
function TRAIN
(H(x))) + (2:)1=1 GLi + (a(((z)))H(())
> Evaluate input representations Form the kernel matrix using 3 and/or (5) Compute KRR coefficients where E (Ei....: Er)
(G)
return e end function
function PREDICT(2)
(x)
Evaluate GNN representation > Form the kernel matrix b Return the KRR prediction
( -
(H))
return e v(x) end function
of node features corresponding to atoms of the s'th chemical species. We then consider composite kernel
K. -I
+ Σ (.H(s)),
(5)
The kernel Ka is a direct generalization of the mean embedding kernel (3) where the parameter allows to interpolate between emphasizing all atomic interactions equally (a = 0) and only within-species atomic interactions (a 1).
Solving KRR (4) with the composite kernel 5 is also equivalent to considering an extension of
(4) where each chemical species is given its weight while encouraging the weights to be close to each other and of small magnitude through a regularizer. Precisely, let (wa)² be the set of displacements from a center weight wo for each chemical species. The weight of chemical species is given by w, + Wp and the potential energy function for a configuration æ takes the form Ea=1 (w., + wo, Φ (x))) In Appendix A we provide full characterization of this extension, highlighting its connection to multi-task learning 57 58 in which multiple tasks are Iearned jointly.
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 customary in the literature, we employ the root mean squared error (RMSE) normalized by the number of atoms the comparison between systems different sizes. We make the code repository available at https:// teig..tu.... www./oooo/asaaaic http://tmmiiiitcteeeo htt.ttmmssiic_ttrssfeeeeeeeeee _merr
metric, to facilitate
5.1 Baselines and MEKRR
We consider baselines spanning different categories. Firstly we examine supervised learning algo- rithms 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 radial basis function representation the interatomic distances. Spherical Channel Networks (SCN) [14] another GNN and its atom embeddings are a set of spherical functions represented via spherical harmonics. SCN
one of the state-of-the-an models on the OC20 dataset. For both SchNer 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 66 62).

---

The second category of baselines concerns transfer learning methods based on the OC20 dataset [2]] 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-(SchNeo)) 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 sestem-specifie MD datasets where our method is fine-tuned on (Cu/formate, Fe/N2).
OC20 The Open Catalyst (OC) 20 is large dataset of ab initio calculations aimed estimating adsorption energies on catalytic surfaces. It comprises ~20 millions of DFT calcula- tions, 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 compute the adsorption energy. The adsorbate is selected out of 82 molecules relevant to environmental applications, while, for each of them, up to 55 surfaces are selected, including binary and ternary compounds. We underline that, for each aesoobate-surface pair, the configurational space sampled is very limited, and especially does not cover of---uiiiiiiiiumm and reactive (e.g. bond forming or breaking) events.
We fine-tune the method and test it on two datasets that representative of reactive catalytic events, obtained by means of molecular dynamics simulations coupled with enhanced sampling methods [64 665 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 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 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 dehydrogena- tion 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/N2 (D;) The second dataset consists of molecular dynamics simulations of nitrogen molecule adsorbing on an iron (Fe) <111> surface at high temperature (T = =70 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 transfcrability of the methods across different conditions. For this reason, divide it into 4 subsets, denoted with D,:
D1 : AI-MD Ab initio molecular dynamics simulations. The resulting configurations are highly correlated and cover small portion of the configurational space related to the adsorption process, thus being the closest dataset to the OC20 one.
D2: AI-METAD Here the ab initio MD simulation is accelerated with the help of the metady- namics [64] technique. This an importance sampling method that allows rare events 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 Dj. allowing one to sample one bond- hond-breaking event.
D3 AL-METAD Dataset built from an active learning procedure using ensemble of NNs combined with metadynamics. In this simulation, multiple reaction events are observed, covering wider region of the configurational space and providing a large number of
7

---

uncorrelated samples. Hence, these configurations are far from those used to pre-train the feature map.
D1 AL-METAD-72 Same as D3 but the surface composed of 72 atoms (8 layers) to test the transferability across systems of different sizes.
5.3 Interpolating between shared and independent weights
The a parameter in the Ka kernel can vary in the range [0, 1}, which are the limiting cases between 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 E {10-3,.. ...,, 10-9) using the same datasets. Despite this simple heuristic cross- cross-validation scheme, we will see, the scheme is effective, which we believe is strength it shows that the MEKRR method is simple to tune while still having the strongest performance among the competitors.
0 and cross-validate
In Fig. 2 we show the cross- cross-validation curves for MEKRR-((ShNee related to the two datasets Cu/formate and Fe/N2. In the latter case, we perform the cross-validation only on (D2), 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 a 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₂ dataset the optimal a occurs at the boundary leading to kernel in which the weights for the two species are learned independently.
Cu/formate
Ne/Na (Dz)
2.00
Validation error [meV/ [ueV//tora] //toaa]
RMSK
1.75
MAE
1.50
1.25
1.00
0.75
0.50
10"
10 2
10a 10 10
10"
Validation error (RMSE/ / MAE) of MEKKR (SchNe) on the Cu/formate and Fe/N, (D2) datasets function of geometrically spaced on grid from 0 to 1 with optimal and error given by a bold orange point. The optimal a for the Cu/formate dataset is positive but close to zero while the optimal for the Fe/N2 is found the boundary at 1.0 leading to hard multi-embedding kemel. We see that tuning the a allows for improved performance in practice and that the multi-weight formulation (5) is practically beneficial.
Figure 2:
5.4 Potential energy regression
In this section, we consider the setting of predicting the potential energy surface. We first evaluate the performance 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 I 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 aatto--distribution for the pre-trained feature map. This is particularly evident for D3 and D4 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.
8

---

Table dateset energy prediction, metric being RMSE. The errors are in units of V/atom. Best performance given by bold number in gray cell.
F
Group
Algorithm
Cu/formate
D1
D2
D3
D4
Supervised
GAP
0.4
2.1
3.9
4.9
2.8
SchNet
0.5
4.1
5.1
6.2
6.0
SCN
0.3
5.1
7.5
7.3
2.5
Fine-tune
SchNet-FT
0.1
2.0
2.5
3.2
1.9
Ours
MEKRR-(SchNet)
0.1
1.3
2.4
3.3
1.2
MEKRR-(SCN)
0.2
0.9
1.9
2.7
1.7
Across-dataset energy prediction Here we evaluate the performance of the algorithms and MEKRR on transferring from different systems in the Fe/Nz family of datasets. To do this consider the task of zero-shot transfer learning (see e.g. 18 and references therein) where we evaluate model trained on source dataset Douurce On target dataset Dargett. While the two datasets Dsuurce and Drarget may be sampled from arbitrary systems, we consider here systems that share some characteristies as are evaluating the transfer capability of the models 688 69]. Due to the ordering of the datasets D1,-.., D₁ in increasing complexity on several axes (size, ab initio vs. active sampling, standard vs. biased dynamics, etc.) we consider transfer from simpler 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- MEKRR- (SchNet has the lowest error the remaining task. Furthermore, the relative transferability of MEKRR compared to the other methods even improves as the task becomes harder. To this respect, worth noting that D1, D2, D3 are qualitatively similar, being composed of layers of Fe and differing for the sampling method used. Instead, the atomic environments contained in Da are different they refer to slab with different number layers. This explains the different order of magnitudes in the last two columns. Despite this, MEKRR still performs very well compared to the basclines.
has the lowest eтror in four out of five tasks, while
Table 2: Transfer evaluation of algorithms on source to target: Dsource Drarget metric being RMSE. The errors are in units of me V/atom. Best performance is given by bold number in gray cell.
Group
Algorithm
D1 + D2
D1 + D3
D2 + D3
D2 + D4
D3
Da
Supervised
GAP
24.9
59.1
5.8
830
888
SchNet
13.2
15.4
6.2
93
107
SCN
22.1
29.3
9.7
139
131
Fine-tune
SchNet-FT
17.6
27.3
3.7
121
116
Ours
MEKRR-(SchNet)
8.0
9.3
2.9
27
55
MEKRR-(SCN)
7.0
6.3
2.0
40
42
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 pre-trained feature map together with the K2 kernel. However, this idea not restricted supervised learning. We can indeed leverage the similarity measure provided by the kernel for tasks beyond potential energy regression. As simple example, in Fig. we plot the kernel matrix, when using SchNet as the feature map, of a part of the trajectory of D, containing an N-N bond breaking event in the cases = and a 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 visualized on the top margin the heatmap along with the time evolution of physical quantity that signals the N-N bond- mondd-beaking... This facilitates

---

(N'N I
150
1.2
1.
150
0.8
450
Figure 3: Heatmaps of the Ka SchNNet kernel applied part of the trajectory of D2 (where reactive event occurs) and time series the distance between nitrogen atoms over time L. The cases with 0 and the two kernels as inputs we label each time-index with one of two classes, with the background color showing the class. Spectral clustering with the multi-weight kernel picks out the reactive event perfectly.
reported the left and right, respectively. Using spectral clustering with
physical interpretation, as the two classes correspond to configurations containing the reactants (the N2 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 under-represented, which typically correspond to adsorbed atoms in surfaces. This allows to give more weight the most important actors catalytic applications.
Conclusion and future work
In this work, introduced approach 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 tested 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 random features (70) 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.
catalytic processes that outside the distribution of the
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.
CRediT author statement
Acknowledgements We acknowledge the financial support from the PNRR MUR Project PE000013
CUP J53C2200301 0006 "Future Artificial Intelligence Research (FAIR),, funded by the European Union NexIGenerationEU, EU Project ELIAS under grant agreement No. 101 10123777, and the "Joint project TransHyDE_FPP Reforming ammonia transport Hz, via derivatives". funded from the German Federal Ministry of Research (BMBF), funding code: 03HY203A-F.
10

---

References
[1] Daan Frenkel and Berend Smit. Understanding Molecular Simulation: From Algorithmss to Applications, volume 1. Elsevier, 2001.
[2] J. E. Lennard-Jones. Cohesion. Proceedings of the Physical Society, 43(5):461, 1931.
[3] Roberto Car and Michele Parrinello. Unified approach for molecular dynamics and density- functional theory. Physical Review Letters, 55(22):2471, 1985.
[4] Dominik Marx and Jurg Hutter. Ab initio molecular dynamics: Theory and implementation. Modern Methods and Algorithms of Quantum Chemistry, 1(301-449::41, 2000.
[5] Walter Kohn and Lu Jeu Sham. Self-consistent equations including exchange and correlation effects. Physical review, 140(4A):A1133, 1965.
[6] Jörg Behler and Michele Parrinello. Generalized neural-network representation of high- dimensional potential-energy surfaces. Physical Review Letters, 98( (8(14))1464)) 2007.
[7] Oliver T. Unke, Stefan Chmiela, Huziel E. Sauceda, Michael Gastegger, Igor Poltavsky, Tr Schütt, Alexandre Tkatchenko, and Chemical Reviews, 121(16):10142-100 101868 2021.
Müller. Machine Learning Force Fields.
[8] Jörg Behler and Gábor Csányi Machine learning potentials for extended systems: perspective. The European Physical Journal B, 94:1-11, 2021.
[9] Albert P. Bartók, Risi Kondor, and Gábor Csányi. On representing chemical environments. Physical Review B, 87(18):184113 2013.
[10] Michael J Willatt, Félix Musil, and Michele Ceriotti. Atom-density representations for machine learning. The Journal of Chemical Physics, 150(15):154110 10, 2019.
[11] 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-911 2021.
[12] Ilyes Batatia, David P Kovaes, 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:1 1423-11436, 2022.
[13) 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.
[14] 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.
[15] Simon Batzner, Albert Musaelian, Lixin Sun, Mario Geiger, Jonathan P. Mailoa, Mordechai Kornbluth, Nicola Molinari, Tess E. Smidt, and Boris Kozinsky. E()e-viuirariant graph neural networks for data-efficient and accurate interatomic potentials. Nature Communications, 13(1):2453, 2022.
[16] Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E(n) Equivariant Graph Neural Networks, 2022.
[17] 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 (Long and Short Papers), pages
4171 1171-1 Association for Computational Linguistics, 2019.
11

---

[18] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Weerbert-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, llya Sutskever, and Dario Amodei. Language Models are Few-Shot Learners. In Advances in Neural Information Processing Systems, volume 33, pages 1877-1901, 2020.
[19] 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.
[20] 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.
[21]
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 Nic, 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.
[22] Lowik Chanussot, Abhishek Das. Siddharth Goyal, Thibaut Lavril, Muhammed Shuaibi, Mor- gane 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):05596077) 2021.
[23] 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.
[24] 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)ee 2017.
[25) Stefan Chmiela, Huziel F. 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.
[26] Linfeng Zhang, Jiequn Han, Han Wang, Wissam Saidi, Roberto Car, al. End-to-end symmetry preserving inter-atomic potential energy model for finite and extended systems. Advances in Neural Information Processing Systems, 31, 2018.
12

---

[27] Kristof Schütt, Pieter-Jan Kindermans, Huziel Enoc Sauceda Felix. Stefan Chmicla, Alexandre Tkatchenko, and 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.
[28] Johannes Klicpera, Janek GroB, and Stephan Günnemann. Directional Message Passing for Molecular Graphs. In Eighth International Conference on Learning Representations, 2020.
[29) 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.
[30] Liam Collins, Aryan Mokhtari, Sewoong Oh, and Sanjay Shakkottai. MAML and ANIL Provably Learn Representations. In Proceedings of the 39th International Conference Machine Learning, pages 4238-4310. PMLR, 2022.
[31] John Isak Texas Falk, Carlo Ciliberto, and Massimiliano Pontil. Implicit kernel meta-learning using kernel integral forms. In Proceedings of the Thirly-Eightt Conference Uncertainty in Artificial Intelligence, pages 652-662. PMLR, 2022.
[32] 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-277 271700, 2021.
[33] Ruohan Wang, Isak Falk, Massimiliano Pontil, and Carlo Ciliberto. Robust meta-repressentation learning via global label inference and classification. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2023.
[34] Guneet Singh Dhillon, Pratik Chaudhari, Avinash Ravichandran, and Stefano Soatto. A Baseline for Few-Shot Image Classification. In International Conference on Learning Representations, 2023.
[35] Jake Snell, Kevin Swersky, and Richard Zemel. Prototypical Networks for Few-shot Learning. In Advances in Neural Information Processing Systems, volume 30, 2017.
[36] Luca Bertinetto, Joao Henriques, Philip Torr, and Andrea Vedaldi. Meta-learning with differentiable closed-form solvers. In International Conference on Learning Representations, 2023.
[37] 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.
[38] Viktor Zaverkin, David Holzmüller, Luca Bonfirraro, and Johannes Kästnner Transfer learning for chemically accurate interatomic neural network potentials. Physical Chemistry Chemical Physics, 25(7):5383 (37):55333-3396 2023.
[39] 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.0400 04002, 2023.
[40] Duo Zhang, Hangrui Bi, Fu-Zhi Dai, Wanrun Jiang, Linfeng Zhang, and Han Wang. Dpa-I: Pretraining of attention-bast based deep potential model for molecular simulation. arXiv preprint arXiv:2208.082266 2022.
[41] Tian Xie and Jeffrey C. Grossman. Crystal Graph Convolutional Neural Networks for an Accurate and Interpretable Prediction of Material Properties. Physical Review Letters, 1 120(14): 2018.
[42] KristofT. Schütt, Oliver T. Unke, and Michael Gastegger. Equivariant message passing for the prediction of tensorial properties and molecular spectra, 2021.
[43] 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.
13

---

[44] Ingo Steinwart. Support Vector Machines. Information Science and Statistics. Springer, New York. NY, 2008.
[45] Lei Shi, Xin Guo, and Ding- Xuan Zhou. Hermite learning with gradient data. Journal of Computational and Applied Mathematics, 233(1 1):3046-300000 -054,, 2010.
[46] Carl Edward Rasmussen. Gaussian processes in machine learning. In Advanced Lectures on Machine Learning, pages 63-71. Springer Berlin Heidelberg, 2004.
[47] Krikamol Muandet, Kenji Fukumizu, Bharath Sriperumbudur, and Bemhard Schölkopf. Kernel
Mean Embedding of Distributions: A Review and Beyond. Foundations and Trends@ in Machine Learning, 10(1-2)11-141 2017.
[48] Michael M. Bronstein, Joan Bruna, Taco Cohen, and Petar Veličković. Geometric Deep Learning: Grids, Groups, Graphs, Geodesics, and Gauges, 2021.
[49] 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.
[50] Alessio Micheli. Neural network for graphs: A contextual constructive approach. IEEE Transactions on Neural Networks, 20(3):498-51) 2009.
[51] 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. Al Open, 1:57-81, 2020.
[52] Kristof T. Schütt, Huziel E. Sauceda, Pieter-Jan Kindermans, Alexandre Tkatchenko, and Klbber Müller. SchNet deep learning architecture for molecules and materials. The Journal of Chemical Physics, 148(24):241722, 2018.
[53] Johannes Gasteiger, Florian Becker, and Stephan Günnemann. GemNet: Universal Directional Graph Neural Networks for Molecules, 2022.
[54] 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.
[55] Kristof T. Schütt, Farhad Arbabzadah, Stefan Chmiela, Klaus R. Müller, and Alexandre Tkatchenko. Quantum-chemical insights from deep tensor neural networks. Nature Com- munications, 811):13890, 2017.
[56] 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.
[57] Theodoros Evgeniou and Massimiliano Pontil. Regularized multi task learning. Proceedings of the 10th ACM SIGKDD International Conference Knowledge Discovery and Data Mining, pages 109-117, 2004.
[58] Theodoros Evgeniou, Charles A. Micchelli, and Massimiliano Pontil. Learning multiple tasks with kernel methods. Journal of Machine Learning Research, 6(21):615-63) 2005.
[59] K. T. Schütt,, H. E. Sauceda, P.-J. Kindermans, A. Tkatchenko, and K.-R. Müller. SchNet deep learning architecture for molecules and materials. The Journal of Chemical Physics, 148(24):24172,, 2018.
[60] Albert P. Bartók, Risi Kondor, and Gábor Csányi. On representing chemical environments. Physical Review B,
2013.
[61] James R. Kermode. F90wrap: An automated tool for constructing deep Python interfaces to modern Fortran codes. Journal of Physics: Condensed Matter, 32030) 2020.
[62] G. Csanyi, S. Winfield, J. Kermode, M. C. Payne, A. Comisso, A. De Vita, and N. Bern- stein. Expressive Programming for Computational Physies Fortran 950+. Newsletter of the Computational Physics Group, pages 1 1-24, 2007.
14

---

[63] Damien Garreau, Wittawat Jitkrittum, and Motonobu Kanagawa. Large sample analysis of the median heuristic, 2018.
[64] Alessandro Laio and Michele Parrinello. Escaping free-energy minima. Proceedings of the national academy of sciences, (9(20))12522-123))) 2002.
[65] 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, (13(22)9901199044 2000.
[66] Luigi Bonati, Danicla 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.
[67] 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.
[68] Karl Weiss, Taghi M. Khoshgoftaar, and DingDing Wang. A survey of transfer learning, Journal of Big Data, 3(1):9, 2016.
[69] Justin S. Smith, Benjamin T. Nebgen, Roman Zubatyuk, Nicholas Lubbers, Christian Devereux, Kipton Barros, Sergei Tretiak, Olexandr Isayev, and Adrian Roitberg. Approaching coupled cluster accuracy with general-purpose neural network potential through transfer learning. Nature Communications, 10(1):2903, 2019.
[70] Ali Rahimi and Benjamin Recht. Random features for large-scale kernel machines. In Proceed- ings of the 20th International Conference Neural Information Processing Systems, NIPS'07, page 1177-1184, Red Hook, NY, USA, 2007. Curran Associates Inc.
[71] 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
15


Reviews:
[{"review_id": "SY5gzc5K1k", "reviewer": "Reviewer_7Ada", "summary": "This 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.", "questions": "Could 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", "limitations": "The 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.", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "The 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. ", "weaknesses": "The 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": ">…it would be ideal for the authors to make the complete code available.\n\nThe link to the code can be found in the supplementary material (https://anonymous.4open.science/r/mekrr-76F1/).\n\n> Novelty of the proposed method and comparison to to existing GNN works…\n\nWe will expand the contribution section with the comments made in the global reply above, highlighting the key points of novelty. Nevertheless, we wish to point out that [1] concerns parallelizing GNN training which is not the aim of this work (we run all of our experiments using 1 GPU).\n\n> Additionally, conducting an ablation study in the experimental section would be essential…\n\nWe have provided an ablation study in the Appendix, C.3. We note that since the GNN feature map is fixed there is nothing to ablate with respect to it. Instead we focus on the kernel embedding part of the pipeline, showing that the Gaussian kernel is integral to good performance by comparing it with the linear kernel on the same GNN as the feature map. Finally we include a study on how the representation layer of the GNN affects performance.\n\n> The 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\nWe will rewrite the contributions section building on the paragraph in the global response. In particular, we will stress that we are **not** proposing a new GNN architecture for atomistic systems, but rather our goal is transferring pretrained GNNs to new systems effectively using a novel combination with kernel mean embeddings.\n\nAdditionally, we will reorganize the methods section to reflect the inclusion of equivariant models in the experiments. In section 4.1 we will first describe the extraction of the feature vectors from GNN node features of invariant models like SchNet as well as from equivariant ones like SCN. Then in section 4.2 we will discuss the MEKRR method to fine tune the representation on a new dataset, and finally in section 4.3 we will introduce the more general version with a chemical-species dependent term. \n\n> Could the inclusion of a flowchart improve readers' ability to follow the pipeline of this paper? \n\nThanks for the suggestions, we will prepare such a flowchart to be integrated with the revised methods section discussed above, to more clearly highlight the pipeline of our method. \n\n> In 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\nWe will discuss them in the introduction as well as in the methods section, reflecting the inclusion of equivariant models in the experiments. \n\n> The paper lacks in-depth discussions on limitations, such as issues regarding model robustness and ethical considerations, particularly regarding data misuse. \n\nWe will include a more in-depth discussion of limitations in the revised version. In particular, with respect to model robustness we note that MEKRR transfers much more effectively compared to other algorithms (see Table 2), but that it is integral to stress test the final model to confirm that it is up to task before deploying it to production. For ethical considerations and data misuse we are unsure what the reviewer means, in general MEKRR inherits the same ethical problems as the field at large.\n\n"}, {"review_id": "lL2qkw5rSV", "reviewer": "Reviewer_qEyK", "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. ", "questions": "N/A", "limitations": "N/A", "rating": 7, "confidence": 5, "soundness": 4, "presentation": 3, "contribution": 3, "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. \n\nGiven 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).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your positive comments.\n\n> 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.\n\nWe will clarify this point further in our revision. Our method is similar in spirit to meta-learning a data representation which is fine tuned on novel tasks. A core insight behind meta-learning  comes from few-shot image classification (e.g. refs. [29–33] in the submission), where we want to learn an algorithm that adapts quickly to new tasks given a collection of image classification tasks, and  an integral part of this approach is to add a preprocessing step where each image is mapped through a meta-learned or pre-trained representation.\n\n> …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 \n\nWe have added an experiment using an equivariant model (SCN), see the general answer above in addition to Tables 1 and 2 in the pdf. A more formal investigation of how to use equivariant representations in the MEKRR algorithm (e.g. when dealing with the direct sum of irreducible representations as atom embeddings) is an important direction but outside of the scope of the current paper.\n"}, {"review_id": "97LWZBW3mt", "reviewer": "Reviewer_D8fw", "summary": "The 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. ", "questions": "* 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.", "limitations": "The authors well clarify the limitations and potential negative societal impact of their work in the paper.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "* MEKRR achieves excellent generalization and transferability performance on the introduced datasets with increasing complexity.\n* The paper is well-organized and clearly written.", "weaknesses": "* 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> … no substantial theoretical evidence or intuition is provided to motivate this claim. \n\n\nWhile we do not provide formal evidence we believe the experimental results speak for themselves. Theoretical justification could be made using results from transfer learning and kernel ridge regression. A good starting point could be the work by [1].\n\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\nThis is not correct. Although in this paper for simplicity we only focused on energy prediction, the forces can be simply calculated via the chain rule, multiplying the derivatives of the kernel with respect to the feature vector by the derivatives of the feature vector with respect to atomic positions; see [2] for how to do this effectively in practice using automatic differentiation. \n\n> 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\nFirst, we would like to stress that our focus is to perform transfer learning to a target atomistic system, hence the training set will typically be rather small. Indeed, datasets used for modeling realistic potentials for molecular dynamics applications are of the order of $10^3-10^4$ samples.\n \nSecond, there are several techniques that can be used to remove the linear scaling of kernel methods (such as random fourier features, Nystrom sampling etc.) which could be studied in future work. We will add this comment in the discussion of the limitations.\n\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\nWe disagree, and indeed, this is exactly the point of our work: to show that the GNN-based representations trained in OC20 transfer well to small datasets when used with kernel mean embeddings and kernel ridge regression. As we discuss in the general answer, we consider a real-life scenario in which we want to specialize our pre-trained model to a new system, using a transfer learning approach. This is motivated by the fact that labeling of new configurations is very expensive and so in the practice one has to work with the constraint of limited data availability.\n\nFor this reason, we should test the generalization ability of the model not to larger molecules but rather to different regions of the configurational space, which is exactly what we do in the transfer-learning experiments reported in Table 2.\n\nFurthermore, our work is in synergy and not in competition with the GNN approaches. Indeed, it is thanks to the representations learned by GNN architectures that a conceptually simple method like kernel ridge regression works so well!\n\n> In the main text, Equation 7 mentioned in line 205 is not found. Could the authors clarify this further? \n\nThis should be Equation (4) (see the supplementary material for the corrected cross-reference). We will update this in the final version of the paper.\n\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\nAs described in the general answer, we repeated the experiments using SCN, an equivariant model which was devised specifically for the OC20 dataset and indeed it is among SOTA models. As discussed in the general answer (see Tables 1 and 2) all the conclusions obtained when using SchNet hold also for this more elaborate GNN model, supporting the generality of the proposed approach.\n\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\nAs discussed above, a crucial property is transferability across geometric environments. To this extent, we consider the transfer learning across the $D_i$ datasets. They have been taken from a single dataset used to study a realistic catalytic application, and divided according to the sampling technique used, covering an increasingly larger portion of the configuration space. As detailed in section 5.2, they represent by construction complementary datasets of increasing complexity, for which it is natural to obtain gaps in validation metrics between in-distribution and out-of-distribution settings. \n\nThe fact that MEKRR obtains consistently better results on this task (sometimes by one order of magnitude) is a clear indication of its better transferability to wider portions of the configuration space. \n\n> … For instance, RMSE in D1→D2  is 80 times larger than RMSE in D1. Is the scale of RMSE in D1→D2 acceptable in real chemistry research?\n\n$D_1$ represents a limiting case, where the configurations are collected via DFT (oracle)-based molecular dynamics, and as such they are composed of very correlated and near-equilibrium. configurations, meaning that the effective sample size is very small. Indeed, collecting configurations only from ab initio molecular dynamics is not sufficient [4] and needs to be complemented with other sampling schemes. We will expand the discussion of the datasets clarifying this.\n\n[1] Tripuraneni, Jordan, and Chi Jin, Advances in neural information processing systems 33 (2020): 7852-7862.\n\n[2] Schmitz, Klaus-Robert Müller, and Chmiela, The Journal of Physical Chemistry Letters 13.43 (2022): 10183-10189.\n\n[3] Chatalic, Antoine, et al, International Conference on Machine Learning. PMLR, 2022.\n\n[4] Unke, Oliver T., et al., Chemical Reviews 121.16 (2021): 10142-10186.\n\n"}, {"review_id": "bNRkzNn2Nj", "reviewer": "Reviewer_f9vG", "summary": "In 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. ", "questions": "Do 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", "limitations": "Limitations are addressed in part, but further discussion is needed related to more modern ML interatomic potentials.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The 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", "weaknesses": "The 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> The 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\nAs discussed above we do not simply test our method against a GNN, but we synergistically combine the representation power of GNN with the simplicity of kernel ridge regression. To further support the generality of our approach, we added new experiments (see Table 1 and Table 2) in which we both extract the feature map from SCN and use this method as a baseline.\n\n> One 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\nThis is not correct. Indeed, the experiment in Table 2 measures exactly the transferability across geometric environments. Indeed, each of the $D_1$, … , $D_4$ datasets cover increasingly larger regions of the phase space (see detailed description in section 5.2). MEKRR performing better on this task is a clear sign of improved transferability than using GNN methods alone. This is another key contribution of our work: to show that the representation learned from the OC20 dataset (consisting of a few correlated samples of many chemical systems) can provide an effective representation also for out-of-sample configurations of chemical interest. The datasets tested, which are taken from real-life applications to molecular dynamics, indeed contain reactive configurations in which bonds are broken or formed which are severely out-of-sample for the pre-trained feature map.\n \n> Do 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\nIf scalability refers to predicting energy for systems with large numbers of atoms, we see no obstacles in principle. In fact, recent work [1] has shown how, with proper adjustments, simulations of 0.5 trillion atoms can be performed using a Gaussian process, which shares the same limitations as our method.\n\n> Why 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\nThe goal of the paper is not to perform well on OC20, but how to effectively transfer GNNs pretrained on OC20 to systems qualitatively different from OC20, see the general answer and the final comment to Rev. D8fw concerning evaluation criteria. As suggested by the reviewer, we have added SCN to the experiments (Table 1 and Table 2) as a baseline, GNN fine-tuning and a feature map for MEKRR. We chose SCN since it was implemented in the OCP codebase and thus easy to integrate into our codebase and had pretrained weights available.\n\nRegarding fine-tuning of a pretrained GNN representation with MEKRR vs standard fine-tuning, we invite the reviewer to look at Table 3, in Appendix, section B “Timings”, for some wallclock comparison of computational complexity of training a GNN from scratch vs fine-tuning vs MEKRR where MEKRR is order of magnitudes faster. Indeed, for the new SCN pre-trained model we were not able to fine-tune it to the new dataset with an acceptable quality, see the general answer above. \n\nFurthermore, in realistic applications one often has to resort to active learning strategies, and retrain the potential whenever new configurations are labeled. In this scenario, kernel ridge regression methods are very appealing.\n\n> Limitations are addressed in part, but further discussion is needed related to more modern ML interatomic potentials.\n\nFollowing your and other reviewers' suggestion we will include in the revised manuscript an expanded discussion on limitations.\n\n[1] Johansson et. al., arXiv:2204.12573 (2022)\n"}]

Author rebuttal:
We thank the reviewers for their useful comments, which we will incorporate in our revision. Before addressing each review in detail, in this global message we answer some key suggestions/questions from the reviewers that we implemented and summarize our contributions.

Note that we have updated Table 1 and Table 2 with results using SCN which may be found in the attached pdf file.

__Comparison to more recent baselines__

While the original experiments only focused on SchNet as a pretrained feature map, we have conducted additional experiments using SCN, which is a state-of-the-art architecture on the OC20 S2EF task and an equivariant GNN, as suggested by several reviewers. In particular, we used the feature map extracted from SCN to build MEKRR, both for the standard and the multi-species $\alpha$ version, and added the baselines of training SCN from scratch or fine-tuning the feature map on the new dataset.

We report the results of these additional experiments in the updated Table 1 and Table 2 (see attached pdf). From these results, we can draw the following conclusions:
* We confirm the original conclusions with MEKRR outperforming training the GNN from scratch or standard fine-tuning, while being orders of magnitudes faster to train.
* Using MEKRR with the pre-trained features obtained with SCN as opposed to SchNet in general leads to a slight improvement in performance. 
* Tuning the multi-species $\alpha$ parameter leads, again, to further improvements.

These results support the generality of our approach.

We want to highlight the fact that we could not properly fine tune SCN (it should probably be run for more epochs, we run it for 200 epochs for Table 1) due to the time constraints. Below we show this explicitly by comparing the time needed to train MEKRR-SCN vs time to train SCN-FT for one epoch both in seconds on datasets $D_1$ to $ D_4$. Both methods use 8 as the representation layer.


| Algorithm                    | $D_1$ | $D_2$ | $D_3$ | $D_4$ |
| ---------------------------- | ----- | ----- | ----- | ----- |
| SCN-FT (seconds / epoch)     | 50    | 88    | 43    | 61    |
| MEKRR-SCN (seconds to train) | 60    | 119   | 53    | 85    |

Note how MEKRR-SCN __fully trains__ in the time it takes to train SCN-FT for __one epoch__.

__Our contributions:__

Several of the reviewers asked for clarification about the contributions and key points of novelty of our approach. Here we briefly summarize them and we will include an expanded discussion in the manuscript.

Our proposed method, MEKRR, confronts the open problem of learning interatomic potentials in settings where high quality data samples are scarce and/or extremely costly. The _sample complexity_ of an algorithm  — that is, the number of data points necessary to guarantee a small generalization error — is heavily influenced by the representation of the data, and MEKRR is designed around this fact.

Our construction combines a feature map pre-trained on a large dataset (OC20) with a fine-tuning scheme based on kernel mean embeddings. In this way we aim at extracting the best out of the representation power of GNNs and kernel methods for interatomic potential learning.

To the best of our knowledge, chaining kernel mean embeddings to deep-learned feature maps is new in this context, and we precisely choose kernel mean embeddings to satisfy the permutational symmetry inherent in atomistic datasets. GNN features, on the other hand, automatically takes care of roto-translational invariance/equivariance, making our method a natural fit for the problem of interatomic potential learning.

We point out that our approach shows __excellent out-of-distribution performance__. Indeed, the OC20 data points, on which the feature maps were pre-trained, are __definitely not__ sampled from physically relevant distributions (see e.g. section 5.2 of the draft). Despite this, MEKRR still outperforms the other baselines on the realistic datasets for the two specific systems Cu/Formate and Fe/N$_2$.


Meta-review: This paper proposes a transfer learning algorithm for modeling the energy surface of atomic systems, where the representations obtained from pre-trained GNNs are combined with kernel mean embeddings. The paper demonstrated the generalizability and transferability of the proposed method, offering significant gains over baselines. Most of the reviewers are happy with the quality and significance of the present work.
Decision: Accept (poster)