SlowGuess/NeurIPS_2023_OCR huggingface dataset:
John I. Falk
CSML
me@isakfalk.com
Luigi Bonati
Atomistic Simulations
luigi.bonati@iit.it
michele.parrinello@iit.it
University College London, U.K.
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.
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  $\sim 1.2\mathrm{M}$  catalytic systems, totaling over  $260\mathrm{M}$  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
We consider systems composed of  $n$  atoms and  $S$  different chemical species, described via one-hot-encoding over  $[S]$ , where  $[S] = \{1, \dots, S\}$ . Each atom is described by its Euclidean position  $r \in \mathbb{R}^3$  and its chemical species  $z \in \{0, 1\}^S$ . We denote a state of the system by  $x = (r_i, z_i)_{i=1}^n = (R, Z) \in \mathcal{X}$  using the design matrices  $R \in \mathbb{R}^{n \times 3}$  and  $Z \in \{0, 1\}^{n \times S}$ .
The quantity we want to regress is the potential energy, which is a scalar function  $E: \mathcal{X} \to \mathbb{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=1}^{T}$  together with the corresponding labels  $(E_{t})_{t=1}^{T}$ .
The physics of the system posits that any estimator of the energy  $\hat{f}:\mathcal{X}\to \mathbb{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=1}^{T}$ . 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
$$
\hat {\mathcal {R}} (w) = \sum_ {t = 1} ^ {T} \left((1 - \gamma) \left(E _ {t} - f _ {w} \left(R _ {t}, Z _ {t}\right)\right) ^ {2} + \gamma \| F _ {t} + \nabla_ {R _ {t}} f _ {w} \left(R _ {t}, Z _ {t}\right) \| ^ {2}\right) + \lambda \Omega (w), \tag {1}
$$
for some regularizer  $\Omega$  and regularization parameter  $\lambda > 0$  and a loss weight  $\gamma \in [0,1]$ . A loss weight  $\gamma \neq 0$  is used whenever the forces  $F(x,r) = -\nabla_r E(x)$  are regressed alongside the energy. Typically we will use  $\Omega(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  $\gamma = 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: \mathcal{X} \times \mathcal{X} \to \mathbb{R}$  on the input space  $\mathcal{X}$ . For any kernel it exists a feature space  $\mathcal{H}_K$  and a feature map  $\phi: \mathcal{X} \to \mathcal{H}_K$  such that  $K(x, x') = \langle \phi(x), \phi(x') \rangle$  for all  $x, x' \in \mathcal{X}$ .
Kernel ridge regression is a supervised learning algorithm that parametrize the dependence between inputs  $x \in \mathcal{X}$  and scalar outputs  $y \in \mathbb{R}$  as  $y = \langle w, \phi(x) \rangle$  for some vector  $w$  in the feature space. Given a dataset  $(x_{t})_{t=1}^{T} \in \mathcal{X}^{T}$  of  $T$  input points, and a corresponding dataset of outputs  $(y_{t})_{t=1}^{T} \in \mathbb{R}^{T}$ , KRR learns a functional relation between inputs and outputs by solving the least squares minimization problem
$$
\hat {w} = \underset {w \in \mathcal {H} _ {k}} {\operatorname {a r g m i n}} \left\{\sum_ {t = 1} ^ {T} \left(\left\langle w, \phi \left(x _ {t}\right) \right\rangle - y _ {t}\right) ^ {2} + \lambda \| w \| _ {\mathcal {H} _ {K}} ^ {2} \right\}. \tag {2}
$$
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  $\hat{w} = \sum_{t=1}^{T} c_t \phi(x_t)$  for some  $(c_t)_{t=1}^T \in \mathbb{R}^T$ . The representer theorem and the fact that  $K(x, x') = \langle \phi(x), \phi(x') \rangle$  also imply that predicting any new point  $x$  with KRR is as simple as evaluating
$\langle \hat{w},\phi (x)\rangle = \sum_{t = 1}^{T}c_{t}\langle \phi (x_{t}),\phi (x)\rangle = \sum_{t = 1}^{T}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  $\mathcal{X} \subseteq \mathbb{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  $\mathcal{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 \in \mathbb{R}^d : i \leq n\}$ . For example, an atomic configuration is a set containing the position of each atom in the system ( $d = 3$ ). For any feature map  $\psi(r)$  and the corresponding kernel function  $k(r, r') = \langle \psi(r), \psi(r') \rangle$  on points  $r \in \mathbb{R}^d$  we can define a feature map acting on the whole set  $x$  as
$$
\phi (x) := C _ {x} \sum_ {r _ {i} \in x} \psi (r _ {i})
$$
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
$$
K (x, x ^ {\prime}) = \langle \phi (x), \phi (x ^ {\prime}) \rangle = C _ {x} C _ {x ^ {\prime}} \sum_ { \begin{array}{c} r _ {i} \in x, \\ r _ {j} ^ {\prime} \in x ^ {\prime} \end{array} } \langle \psi (r _ {i}), \psi (r _ {j} ^ {\prime}) \rangle = C _ {x} C _ {x ^ {\prime}} \sum_ { \begin{array}{c} r _ {i} \in x, \\ r _ {j} ^ {\prime} \in x ^ {\prime} \end{array} } k (r _ {i}, r _ {j} ^ {\prime}). \tag {3}
$$
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.
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  $\mathcal{G} = (\mathcal{V},\mathcal{E})$  is a collection of  $n$  nodes  $\mathcal{V}$  and edges  $\mathcal{E}\subseteq \mathcal{V}\times \mathcal{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
$$
h _ {i} \mapsto g _ {\theta} \left(h _ {i}, \sum_ {j \in \mathcal {N} (i)} \eta_ {\theta} \left(h _ {i}, h _ {j}\right)\right).
$$
Here,  $\mathcal{N}(i)$  is the neighborhood of  $i$  i.e. the set of nodes connected to  $i$  through an edge,  $g_{\theta}$  is a learnable vector-valued function such as an MLP and  $\eta_{\theta}$  is a message-passing function. A graph
Figure 1: Diagram of MEKRR.
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
$$
x \mapsto H (x) := \left(h _ {i} (x)\right) _ {i = 1} ^ {n}.
$$
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
$$
H (x) \mapsto \phi (H (x)) = C _ {x} \sum_ {h _ {i} (x) \in H (x)} \psi (h _ {i} (x)).
$$
Finally, at training time we can use the kernel mean embedding  $\phi(H(x))$  to solve the KRR problem as
$$
\hat {w} = \underset {w \in \mathcal {H} _ {k}} {\operatorname {a r g m i n}} \left\{\sum_ {t = 1} ^ {T} \left(\left\langle w, \phi \left(H \left(x _ {t}\right)\right) \right\rangle - E _ {t}\right) ^ {2} + \lambda \| w \| _ {\mathcal {H} _ {K}} ^ {2} \right\}, \tag {4}
$$
while at inference time we can get an estimation for the potential energy at  $x$  by simply evaluating  $\langle \hat{w},\phi (H(x))\rangle$ .
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)\coloneqq \{h_{i}(x)\mid$  species of  $i = s\}$  be the subset
Parameters: Training dataset  $(x_{t},E_{t})_{t = 1}^{T}$  of configurations  $x_{t}$  and potential energies  $E_{t}$ . Kernel function  $k$ . Interpolation parameter of the multi-weight formulation  $\alpha \in [0,1]$ . Tikhonov regularization  $\lambda >0$ . Pre-trained GNN feature map  $x\mapsto H(x)$ .
Algorithm 1 Mean Embedding Kernel Ridge Regression (MEKRR)
function TRAIN
return  $c$  end function
$(H(x_{t}))_{t = 1}^{T}\gets (x_{t})_{t = 1}^{T}$ $\triangleright$  Evaluate input representations  $G_{t,l}\leftarrow K_{\alpha}(H(x_{t}),H(x_{l}))$  Form the kernel matrix using (3) and/or (5)  $c\gets (G + \lambda I)^{-1}E$  Compute KRR coefficients where  $E = (E_{1},\dots ,E_{T})^{\top}$
function PREDICT(x)  
H(x) ← x  
v(x) ← (Kα(H(x), H(xt)))T_{t=1}  
return c^T v(x)  
end function
Evaluate GNN representation
Form the kernel matrix
$\triangleright$  Return the KRR prediction
of node features corresponding to atoms of the  $s$ 'th chemical species. We then consider a composite kernel
$$
K _ {\alpha} \left(H (x), H \left(x ^ {\prime}\right)\right) := (1 - \alpha) K \left(H (x), H \left(x ^ {\prime}\right)\right) + \alpha \sum_ {s = 1} ^ {S} K \left(H _ {s} (x), H _ {s} \left(x ^ {\prime}\right)\right). \tag {5}
$$
The kernel  $K_{\alpha}$  is a direct generalization of the mean embedding kernel (3) where the parameter  $\alpha$  allows to interpolate between emphasizing all atomic interactions equally  $(\alpha = 0)$  and only within-species atomic interactions  $(\alpha = 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 = 1}^{S}$  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  $\sum_{s = 1}^{S}\langle w_s + w_0,\phi (H_s(x))\rangle$ . 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.
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.