Gradient-based grand canonical optimization enabled by graph neural networks with fractional atomic existence

We show that neural networks allow us to formulate gradients for a property that is traditionally not differentiable and that these gradients allow us to formulate new optimization algorithms.

Message-Passing Graph Neural Network with node existence

Introduction

In this paper we introduce a small but powerful adaption to the graph neural network framework, which the majority of the best performing machine learning interatomic potentials (MLIP’s) employ. This allows us to treat not only the Cartesian coordinates of atoms as continous variables but also their existence. This allow us to use gradient-based optimization to determine the chemical composition of a material under different conditions.

Messsage-Passing

Message-passing is a beatifully simple but extremely expressive modelling framework. The basic idea is for the description of nodes, node embeddings, of a graph to be updated based on the description of neighbouring nodes. This is expressed by the equation below

\[ \begin{equation} x_i^k = \gamma^k\left[ x_i^{k-1}, \underset{j\in\mathcal{N}(i)}{\Large\oplus} \ \phi^k\left(x_i^{k-1}, x_j^{k-1}, \mathbf{e}_{ji}\right)\right] \label{eq:message_passing} \end{equation} \] Here \(x_i^k\) is the node embedding of the \(i\)’th node at iteration \(k\), \(\gamma\) and \(\phi\) are learnable functions (neural networks) and \(\underset{j\in\mathcal{N}(i)}{\Large\oplus}\) is a permutation invariant aggregation operator over the the neighbour \(\mathcal{N}(i)\).

As described in much more detail in the paper, node existence \(q_i \in [0, 1]\) can be included through simply as \[ \begin{equation} x_i^k = \gamma^k\left[ x_i^{k-1}, \underset{j\in\mathcal{N}(i)}{\Large\oplus} \ \phi^k\left(x_i^{k-1}, x_j^{k-1}, \mathbf{e}_{ji}\right) \cdot q_j \right] \end{equation} \]