Multi-task graph neural networks for simultaneous prediction of global and atomic properties in ferromagnetic systems^*

A graph G is usually represented in mathematical terms as

$$\begin{align} G = (V, \mathcal{E}) \end{align} \tag{ 1 }$$

where V represents the set of nodes and $\mathcal{E}$ represents the set of edges between these nodes [34]. An edge $(u,v) \in \mathcal{E}$ connects nodes u and v, where ${u,v}\in V$, $\ \mathcal{E} \in V \times V$. The topology of a graph can be described through the adjacency matrix, A, an N × N square matrix where N is the number of nodes in the graph, whose entries are associated with edges of the graph according to the following rule:

$$\begin{align} \begin{cases} A[u,v] = 1~\quad \text{iff} \quad (u,v)\in \mathcal{E} \\ A[u,v] = 0~\quad \text{otherwise}. \end{cases} \end{align} \tag{ 2 }$$

The degree of a node $u \in V$ is defined as:

$$\begin{align} d_u = \sum_{v\in V} A[u,v] \end{align} \tag{ 3 }$$

and represents the number of edges connected to a node. Every node u is represented by a a-dimensional feature vector $\mathbf x \in \mathbb{R}^a$ containing the embedded nodal properties and also a label vector $\mathbf y \in \mathbb{R}^b$ in tasks related to node-level predictions. In order to take advantage of the topology of the graph, many DL models include both the number of neighbors per node, as well as the length of each edge between nodes.

GCNNs embed the interactions between nodes without increasing the size of the input by representing the local interaction zone as a hyperparameter that cuts-off the interaction of a node with all the other nodes outside a prescribed local neighborhood. This is identical to the approximation made by many atomic simulation methods, including the LSMS-3 code used to generate the DFT training data, which ignore interactions outside a given cutoff range. GCNNs [35, 36] are DL models based on a message-passing framework, a procedure that combines the knowledge from neighboring nodes, which in our applications maps directly to the interactions of an atom with its neighbors.

The typical GCNN architecture is characterized by three different types of hidden layers: graph convolutional layers, graph pooling layers, and FC layers. The convolutional layers represent the central part of the architecture and their functionality is to transfer feature information between adjacent nodes (in this case atoms) iteratively. In the kth convolutional layer ($k = 0,1,\ldots, K$), message passing is performed in sequential with the following operations:

• (a)  
  Aggregate information from neighbors: the node u collects the hidden embedded features of its neighbors N(u) as well as the information on the edges (if available) via an aggregation function:

  $$\begin{align} \mathbf {h}_{N(u)}^{k+1} = \mathrm{AGGREGATE}\left(\mathbf {m}_v^{k}, \forall v\in N(u)\right), \end{align} \tag{ 4 }$$

  where $\mathbf {m}_v^k = \mathrm{MESSAGE}(\mathbf {h}_v^k,\mathbf {h}_{e_{uv}}^k)$ is a message obtained from neighboring node v and the edge e_uv that connects them. The vector $\mathbf {h}_v^k$ ($\mathbf {h}_v^k\in \mathbb{R}^{p_k}$) is the embedded hidden feature vector of node v in the kth convolutional layer. When k = 0, the hidden feature vector is the input feature vector, $\mathbf{h}_v^0 = \mathbf x$.
• (b)  
  Update hidden state information: with $\mathbf {h}_{N(u)}^{k+1}$ collected, the nodal feature of node u is updated as in:

  $$\begin{align} \mathbf {h}_u^{k+1} = \mathrm{UPDATE}\left( \mathbf {h}_u^k, \mathbf {h}_{N(u)}^{k+1}\right) \end{align} \tag{ 5 }$$

  where $\mathrm{UPDATE}$ is a differentiable function which combines aggregated messages $\mathbf {h}_{N(u)}^{k+1}$ from neighbors of node u with its nodal features $\mathbf {h}_u^k$ from the previous layer k.

Through consecutive steps of message passing, the graph nodes gather information from nodes that are further and further away. The type of information passed through a graph structure can be either related to the topology of the graph or features assigned to the nodes. An example of a topological information is the node degree, whereas an example of nodal feature in the context of this work is the proton number of the atom. A variety of GCNNs, e.g. principal neighborhood aggregation (PNA) [37], crystal GCNN [20] and GraphSAGE [38], have been developed, differing in the definitions of functions $\mathrm{AGGREGATE, MESSAGE, UPDATE}$ for message passing. One simple example of the function combination is:

$$\begin{align} \mathbf {h}_{N(u)}^{k+1} = \mathbf{W}_{neighborhood}^{(k+1)}\sum_{v\in N(u)}\mathbf {h}_v^k+\mathbf {b}^k, \nonumber \\ \mathbf {h}_u^{k+1} = \sigma\left(\mathbf{W}_{self}^{(k+1)}\mathbf {h}_u^k + \mathbf {h}_{N(u)}^{k+1}\right), \end{align} \tag{ 6 }$$

where $\mathbf{W}_{self}^{(k+1)},\mathbf{W}_{neighborhood}^{(k+1)} \in \mathbb{R}^{p_{k+1}\times p_k}$ are the weights of (k + 1)th layer of GCNN and σ is an activation function (e.g. ReLU) that introduces nonlinearity to the model.

PNA is used in this work and is one of the convolutional layers available in HydraGNN through Pytorch Geometric; PNA combines multiple aggregating techniques to reduce the risk of classifying two different graphs as identical. Batch normalizations are performed between consecutive convolutional layers along with a ReLU activation function. Graph pooling layers are connected to the end of the convolution-batch normalization stack to gather feature information from the entire graph. Global pooling layers aim at collapsing the node feature associated with each atom across a graph into a single feature. This is achieved by summing the local interactions of each atom with its neighbors and use the result to estimate global properties. For atom (node) level features such as the atomic charge transfer and atomic magnetic moment, collapsing the information from all atoms into a total system feature is not needed. FC layers are positioned at the end of the architecture to take the results of pooling, i.e. extracted features, and provide the output prediction.

Larger sizes of the local neighborhood lead to a higher computational cost to train the HydraGNN model, as the number of regression coefficients to train at each hidden convolutional layer increases proportional to the number of neighbors. Further details on the behavior of HydraGNN with different sizes of the local neighborhood have been previously reported [39].

MTL utilizes a NN to simultaneously predict multiple quantities [16] when those predicted quantities are mutually correlated and can act as inductive biases for each other. The improvement of an MTL model depends on how strongly the quantities to be predicted are mutually correlated in a particular application. This type of field specific inductive bias has been defined as knowledge-based, for which the training of a quantity can benefit from the information contained in the training signal for other quantities. Ultimately, MTL allows a direct and automated incorporation of physics knowledge into the model by extracting correlations between multiple quantities, with manual intervention by a domain expert only needed in determining which quantities to use.

In HydraGNN, each predicted quantity is associated with a separate loss function and the global objective function minimized during NN training is a linear combination of these individual loss functions. Formally, let T be the total number of physical quantities, or tasks, we want to predict. A single task identified by index i focuses on reconstructing a function $f_i:\mathbb{R}^{a}\rightarrow \mathbb{R}^{b_i}$ defined as

$$\begin{align} \mathbf{y}_i = f_i(\mathbf x), \quad i = 1,\ldots,T , \end{align} \tag{ 7 }$$

where $\mathbf{x} \in \mathbb{R}^{a}$, $\mathbf{y}_i \in \mathbb{R}^{b_i}$. The MTL makes use of the correlation between the quantities y_i , where the functions f_i in (7) are replaced by a single function $\displaystyle \hat{f}:\mathbb{R}^a\rightarrow \mathbb{R}^{\sum_i^T b_i}$ that can model all the relations between inputs and outputs as follows:

$$\begin{align} \begin{bmatrix} \mathbf{y}_1 \\ \vdots \\ \mathbf{y}_T \end{bmatrix}\, = \,\hat{f}_{\mathbf{W}_\mathrm{MTL}}(\mathbf{x}), \end{align} \tag{ 8 }$$

where $\mathbf{W}_\mathrm{MTL}$ represents the weights to be learned.

The global loss function $\ell_\mathrm{MTL}:\mathbb{R}^{N_\mathrm{MTL}}\rightarrow \mathbb{R}^+$ to be minimized in MTL is a linear combination of the loss functions for the single tasks:

$$\begin{align} \ell_\mathrm{MTL}(\mathbf{W}_\mathrm{MTL}) = \sum_{i = 1}^T \alpha_i \lVert \mathbf{y}_{\textrm{predict},i} - \mathbf{y}_i \rVert_2^2, \end{align} \tag{ 9 }$$

where $\mathbf{y}_{\textrm{predict},i}\, = \,\hat{f}_{\mathbf{W}_\mathrm{MTL},i}\left(\mathbf x\right)$ is the vector of predictions for the $i\textrm{th}$ quantity of interest and α_i (for $i = 1,\ldots,T$) are the mixing weights for the loss functions associated with each single quantity. The values of the α_i 's in equation (9) are hyperparameters of the surrogate model and thus can be tuned. In this work we assigned an equal weight to each property being predicted because we are equally interested in all of them; however, this definition of the loss function enables one to modify the values of the α_i to purposely favor the training toward one property of interest.

As mentioned above, the multiple quantities in MTL can be interpreted as mutual inductive biases because the error of a single quantity acts as a regularizer with respect to the loss functions of other quantities. In order to clarify why MTL can be seen as a regularizer, let us start with the formulation of a regularized training as a constrained optimization problem:

$$\begin{align} \begin{cases} \underset{\mathbf{w}}{\operatorname{argmin}} \quad \lVert \mathbf{y}_{\textrm{predicted}}(\mathbf{w}) - \mathbf{y} \rVert_2^2 \\ \mathbf{c}(\mathbf{w}) = \mathbf{g} \end{cases}, \end{align} \tag{ 10 }$$

where $\mathbf{w}\in \mathbb{R}^N$ is the vector of regression coefficients of the model, y are the target values, $\mathbf{y}_{\text{predict}}(\mathbf{w})$ are the predictions produced by the DL model, $\mathbf{c}(\mathbf{w})\in \mathbb{R}^c$ and $\mathbf{g}\in \mathbb{R}^c$ are quantities used to express a constraint. The formulation in (10) imposes the constraint $\mathbf{c}(\mathbf{w}) = \mathbf{g}$ in a strong form. The values of $\mathbf{c}(\mathbf{w})$ depend on the model configuration identified by the parameter vector w, whereas the reference values g are assumed to be given. The constrained optimization problem in (10) can be reformulated so that the constraint is incorporated in the definition of the objective function itself as follows:

$$\begin{align} \underset{\mathbf{w}}{\operatorname{argmin}} \quad \bigg \{ \lVert \mathbf{y}_{\text{predicted}}(\mathbf{w}) - \mathbf{y} \rVert_2^2 + \lambda \lVert \mathbf{c}(\mathbf{w}) - \mathbf{g} \rVert_* \bigg\}, \end{align} \tag{ 11 }$$

where $\lVert \cdot \rVert_*$ may denote the $\ell^1$-norm or $\ell^2$-norm as explained below. The objective function to be minimized in (11) interprets the constraint as a penalization term with the penalization multiplier λ. This is a weak formulation of the constraint added to the original objective function. Standard $\ell^1$ or $\ell^2$ regularizations correspond to choosing $\mathbf{c}(\mathbf{w}) = \mathbf{w}$ and $\mathbf{g} = \mathbf{0}$, and they recast the constraint term in (11) from the strong formulation to the weak formulation using the $\ell^1$-norm and the $\ell^2$-norm, respectively. In MTL, $\mathbf{c}(\mathbf{w})$ are the predictions of additional target properties whose target values are stored in g, which allows to recast the global objective function used in equation (11) as the global loss function for MTL defined in equation (9). For a fair comparison and to determine the benefit of using other tasks as a mutual regularizer, we do not use additional $\ell^1$ and $\ell^2$ regularizers for the STL training in this work.

Figure 2 shows the topology of an HydraGNN model for multi-tasking learning to model mixing enthalpy, atomic charge transfer, and atomic magnetic moment. The architecture of HydraGNN for MTL is organized so that the first hidden layers are shared between all tasks, while keeping several task-specific output layers. This approach is known in literature as hard parameter sharing.

The architecture of the HydraGNN models has 6 PNA [37] convolutional layers with 20 neurons per layer. A radius cutoff of 7 Å is used to build the local neighborhoods used by the graph convolutional mask. Every learning task is mapped into separate heads where each head is made up of two FC layers, with 50 neurons in the first layer and 25 neurons in the second. This choice for the GCNN architecture is driven by the following considerations. On the one hand, using a small GCNN architecture leads to underfitting because the neural network is not sufficiently complex to learn all the relevant features described in the dataset. On the other hand, an overly complex neural network may trigger overfitting, especially when the dataset is relatively small in volume, as it is for this work. Moreover, a complex MTL neural network would result in similar training to STL because it separates too much the quantities of interest. However, the goal of the MTL is to extract physics correlations, which requires the quantities to overlap sufficiently within the architecture. The DL models were trained using the Adam method [45] with a learning rate equal to 0.001, batch sizes of 64, and a maximum number of epochs set to 200. Early stopping is performed to interrupt the training when the validation loss function does not decrease for several consecutive epochs, as this is a symptom that shows further epochs are very unlikely to reduce the value of the loss function. The training set for each of the NN represents 70% of the total dataset; the validation and test sets each represent half of the remaining data. As discussed in section 3, compositional stratified splitting was performed to ensure that all the compositions were equally represented across training, validation, and testing datasets. Equal loss function weights were used for all properties to define the global loss function for MTL. The training of each DL model was performed on Summit with one model per NVIDIA V100 GPU across two nodes, resulting in ensembles of 12 models per MTL/STL setup discussed in the next section.

We compute and analyze not only the accuracy (root mean-squared error (RMSE)) of each model, but also the standard deviation of RMSE for an ensemble of 12 models. This simple metric enables some quantification of the uncertainty associated with each MTL or STL prediction and an understanding of the reliability of the GCNN models. Note that the RMSE for MTL may be higher compared to STL for two reasons. First, the number of graph convolutional layers and the number of nodes per layer are the same for MTL and STL, but MTL forces these parameters to be shared among the multiple predicted quantities (STL and MTL differ only in the split hidden layers for the heads). Second, MTL introduces an inductive bias in the predictions of a quantity under the influence of other quantities simultaneously predicted.

We first show figure 4 parity plots for the predictions generated by HydraGNN against the DFT data when MTL is used to simultaneously predict mixing enthalpy, charge transfer, and magnetic moment. HydraGNN with MTL predicts all three properties well for most of the samples over the entire dataset, as shown by the alignment of data near the diagonal. The colormap highlights that more sample points for all three properties are tightly clustered near zero, with larger variations as the property magnitudes increase. As expected, this non-uniform concentration of values for the target properties across the data affects the predictive performance of the HydraGNN models, with more accurate predictions in regions with higher concentration of data.

Zoom In Zoom Out Reset image size

In table 1, we compare the predictive accuracy of HydraGNN used to simultaneously predict mixing enthalpy, atomic charge transfer, and atomic magnetic moment with MTL, MTL with each pair of properties, and STL for each individual property. The table shows the average RMSE and the standard deviation from 12 independent runs with random initialization on the test set. Figure 5 shows the probability distribution functions (PDFs) of the results from table 1 for each combination of properties. The lines mark the average values and the filled area indicates one standard deviation. Models with a lower average RMSE are interpreted as more accurate and lower standard deviation as more reliable, as the model predictions are similar across different trainings due to more relevant features extracted that better characterize the dataset. The results show that adding the magnetic moment as a physical constraint improves the both the predictive accuracy and reliability of MTL models over the STL predicting mixing enthalpy only. Addition or replacement of magnetic momement for charge density somewhat reduces the accuracy or reliability, but much less than using STL.

Zoom In Zoom Out Reset image size

Table 1. Test RMSE of HydraGNN to predict physical properties for an FePt alloy. The training method states whether multi-tasking (MTL) or single-tasking (STL) was used and which quantities were predicted: mixing enthalpy (H), charge transfer (C), and/or magnetic moment (M).

Training method	Test RMSE
	Mixing enthalpy	Charge transfer	Magnetic moment
MTL, HCM	$7.54 \times 10^{-03} \pm {8.70\times 10^{-04}}$	${6.77\times 10^{-03}} \pm {3.59\times 10^{-04}}$	${1.04\times 10^{-02}} \pm {4.94\times 10^{-04}}$
MTL, HC	$7.33\times 10^{-03} \pm {4.77\times 10^{-04}}$	${7.36\times 10^{-03}} \pm {3.23\times 10^{-04}}$	−
MTL, HM	$6.64\times 10^{-03} \pm {5.08\times 10^{-04}}$	−	${1.02e-02} \pm {5.23\times 10^{-04}}$
MTL, CM	−	${5.94e-03} \pm {3.02\times 10^{-04}}$	${9.30\times 10^{-03}} \pm {4.12\times 10^{-04}}$
STL, H	$1.02\times 10^{-02} \pm {1.16e-03}$	−	−
STL, C	−	$5.94\times 10^{-03} \pm {4.39\times 10^{-04}}$	−
STL, M	−	−	$8.77\times 10^{-03} \pm {3.18\times 10^{-04}}$

These results suggest that there is a strong correlation between mixing enthalpy and magnetic moment, which has also been documented in the literature for ferromagnetic materials from DFT calculations and experiments [46–48], as well as previous DL studies using multilayer perceptrons for MTL on the same dataset [49]. The total magnetization of the binary alloy FePt is strongly correlated with the concentration of Fe in the alloy, as shown in figure 6 to the left, because only the Fe atoms have a non-negligible magnetic moment. This results in a strong (albeit nonlinear) correlation between the mixing enthalpy and the total magnetization of the alloy, as shown in figure 6 to the right. While this correlation between total properties is useful in determining whether the MTL training should be useful in this case, we note that the correlations between the global and atomic properties of the system are much more complex. The HydraGNN results indicate this correlation indeed helps the MTL predict the enthalpy with the addition of magnetic moment, as well as charge density.

Zoom In Zoom Out Reset image size

The error distributions for MTL are slightly skewed and not centered at zero because the physical properties operate as mutual inductive biases, but this does not affect the accuracy of MTL. While MTL stabilizes the training and prediction and improves accuracy of the global graph mixing enthalpy, the same is not true of the atomic-level properties. The STL models for charge density and magnetic moment are slightly more accurate (as discussed at the beginning of the section), but also slightly more reliable for magnetic moment. The difference in terms of stabilization and accuracy obtained by MTL on the mixing enthalpy and atomic-level properties is possibly due to the different dimensionality of each property (scalar vs. high-dimensional vector).

We finally note that while our results for the mixing enthalpy are less accurate than the best reported GCNN results for DFT data (∼0.21 eV atom⁻¹ in this work versus 0.022–0.067 eV atom⁻¹ [22, 23]). This is partially due to the model used, as PNA does not explicitly include bond angle or crystallographic information, as well as the highly non-equilibrium nature of this dataset (although that the literature results include a much more broad range of chemistries should be taken into account). However, the focus of this work is on the MTL capabilities of HydraGNN and comparisons with STL.

In figure 7, we report the average wall-clock time and standard deviation to train all the MTL and STL models for the same 12 averaged runs from section 4.2 for each combination of properties. In general, STL takes approximately 10% less time to train than MTL for two properties, which in turn takes approximately 10% less than MTL to predict all thee material properties. This is due to the fact that each head introduces additional parameters into the neural network, increasing the total computational cost. However, the training for MTL to simultaneously predict all three material properties is still more than 2.2x faster than to train three separate HydraGNN models using STL.

Zoom In Zoom Out Reset image size