Curvature-informed multi-task learning for graph networks
Abstract
Properties of interest for crystals and molecules, such as band gap, elasticity, and solubility, are generally related to each other: they are governed by the same underlying laws of physics.
However, when state-of-the-art graph neural networks attempt to predict multiple properties simultaneously (the multi-task learning (MTL) setting), they frequently underperform a suite of single property predictors.
This suggests graph networks may not be fully leveraging these underlying similarities.
Here we investigate a potential explanation for this phenomenon – the curvature of each property’s loss surface signiﬁcantly varies, leading to inefﬁcient learning.
This difference in curvature can be assessed by looking at spectral properties of the Hessians of each property’s loss function, which is done in a matrix-free manner via randomized numerical linear algebra.
We evaluate our hypothesis on two benchmark datasets (Materials Project (MP) and QM8) and consider how these ﬁndings can inform the training of novel multi-task learning models.
1. Introduction
Graph networks (GNs) (Battaglia et al., 2018) are considered state-of-the-art machine learning (ML) methods for many scientiﬁc problems, including property prediction of both inorganic crystalline materials (Xie & Grossman, 2018;
Park & Wolverton, 2020) (hereafter “crystals”, e.g., Figure 1) and small organic molecules (Gilmer et al., 2017;
Gasteiger et al., 2020) (hereafter “molecules”, e.g., Figure 2); the same GN architectures have shown success in both domains (e.g., (Chen et al., 2019), who suggest the “crystal” vs. “molecule” terminology and use “material” to refer to both).
A common use case might involve training a model from materials with known properties and then screening a separate dataset lacking those properties to obtain potential candidates to further investigate.
Typically, property-prediction tasks are formulated as single-target regression problems, where the goal is to predict a scalar property like band gap, formation stability, elasticity, or solubility.
However, practical materials design requires optimizing multiple properties of interest.
For example, when designing a cell phone screen, a material with both high optical transparency and hardness would be desired.
Predicting a single property can be posed as a single task for an ML model; using a single ML model to predict multiple properties is thus a type of multi-task learning (MTL).
The most common family of MTL approaches is hard parameter sharing, in which a single representation space is shared across multiple tasks of interest, and task-speciﬁc networks map from that space to output predictions (Caruana, 1997).
However, the latest research in novel MTL techniques largely has been focused on classiﬁcation problems in the natural language processing (NLP) and computer vision (CV) domains – for example, the MultiMNIST problem (Sener & Koltun, 2018), where a model must identify two separate digits in the same picture.
With some recent exceptions – e.g., (Yang et al., 2021; Tan et al., 2021; Kong et al., 2021) – there has been little work exploring the application of MTL to multi-property GN regression problems.
Some existing papers have used hard parameter sharing for material property prediction, in which a single model predicts a set of properties.
However, results suggest that this approach has degraded performance compared to a set of single-property models – see, e.g., (Gasteiger et al., 2020), which predicts twelve quantum mechanical properties of molecules and ﬁnds that a single multi-output model performs worse than a set of single-output models.
Works have explored reasons for this: (Yu et al., 2020) identiﬁes a “tragic triad” of reasons that MTL models might perform worse compared to single-task ones, one of which is that multi-task loss functions can have high curvature, as characterized by the Hessian of the loss.
However, (Yu et al., 2020) rely on ﬁrst-order approximations to curvature and do not directly incorporate second-order information.
Concurrently, several works have applied techniques from randomized numerical linear algebra to directly probe the
Hessians of large ML models (Sagun et al., 2017; Alain et al., 2019; Yao et al., 2019; Ghorbani et al., 2019; Papyan, 2020).
These rely (1) on the fact that the product of a model’s Hessian and another vector (a Hessian-vector product) may be efﬁciently obtained, without explicitly calculating the full Hessian, and (2) that given a matrix-free
Hessian-vector product function, several key spectral properties of the Hessian, including its eigenvalues and trace, may be estimated.
However, these papers focus only on standard image classiﬁcation problems, and do not consider the MTL setting, GNs, or regression problems.
In this paper, we formulate the application of curvatureinformed techniques to MTL models (Section 2.1), describe the application of graph networks to the problem of multiproperty prediction in the domains of materials science and chemistry (Section 2.2 and Section 2.3), and analyze training dynamics of multi-property prediction models using methods based on assessing local curvature of loss function surfaces (Section 3).
Our results suggest how novel multi-property prediction models might inherently account for differences in curvature to enhance learning efﬁciency.
2. Approach
Let G be a space of directed multigraphs corresponding to crystal or molecule structures.
A directed graph g =
( V , E ) consists of a set of nodes V = { v } and edges E
= { ( v, v (cid:48) , k ) } .
Each node v has a node embedding x v ∈ R V , and each edge ( v, v (cid:48) , k ) has an edge embedding e v,v (cid:48) ,k ∈ R E .
A pair of nodes v and v (cid:48) may have multiple edges connecting them, and these multi-edges are indexed by k = 1 , . .
, K v,v (cid:48) .
Furthermore, a graph g has T properties
y =
( y 1 , . . .
, y T )
∈ Y , here assumed to all be scalars.
We require a model that can predict all targets y for a given g .
This entails minimizing a loss with respect to a set of neural network weights { θ sh , θ 1 , . .
, θ T } , where θ t is used only in predicting property t , and θ sh is shared across predictions.
The form of this problem and its neural network is discussed in Section 2.2.
First we consider the loss’s curvature.
2.1. Curvature assessment techniques
We view the shared parameters θ sh of a graph network as a ﬂattened vector in R P ; let L t : R P → R be a propertyspeciﬁc loss function of θ sh , where the property-speciﬁc parameters θ 1 , . .
, θ T are held constant.
Typically, machine learning models consider only ﬁrst-order derivative information of L t when training neural networks.
However, secondorder information is also useful in optimization problems – consider how a quasi-Newton method like limited-memory
L-BFGS (Liu & Nocedal, 1989) is more efﬁcient than a ﬁrst order method like gradient descent.
This suggests that additional insight into the training problems of ML models can be found in using second-order information.
The curvature of a function L t is characterized by its Hessian H t = ∇ 2 θ sh L t ∈ R P × P for a weight vector θ sh ∈ R P .
For typical deep neural networks, P ≈ 10 6 , and calculating
H directly is computationally infeasible due to storage contraints ( H t will have 10 12 entries).
However, the product of H t with a vector v is computable in approximately the same time-complexity as calculating the gradient of L t with respect to θ (Pearlmutter, 1994).
This enables calculation of properties of H t that can be obtained from observing its action on a chosen vector v (Yao et al., 2019; Ghorbani et al., 2019).
In particular, the trace of H t , and an approximate probability density function of its eigenvalues can be efﬁciently estimated.
The trace of H t is both the sum of its eigenvalues and L t ’s
Laplacian (i.e., tr H t = (cid:80) p ∂ θ sh,p θ sh,p L t ( θ ) ).
Thus, the trace of H t can be viewed as a measurement for the complexity of its curvature (Yao et al., 2019).
The sign of the trace also indicates the sign of L t ’s local curvature.
The trace of H t may be estimated with the expectation
[Equation]
where P R is the distribution of random variables in R P with iid Rademacher components (Avron & Toledo, 2011).
Having access to all of H t ’s eigenvalues gives us a full picture of L t ’s curvature.
These eigenvalues are often represented as a function ψ called the spectral density or density of states that is given by ψ ( t ) = 1 P (cid:80) p δ ( t − λ p ) , where λ 1 ≤ · · · ≤ λ P are the eigenvalues of H t , and δ is the Dirac distribution (Lin et al., 2016).
However, the weight vector θ sh is too high-dimensional for all of H t ’s eigenvalues to be directly calculated in practice.
Methods like power iteration can be used to estimate the large-magnitude eigenvalues of a Hessian (Yao et al., 2019), but they scale poorly as its dimensionality increases.
A solution is to ﬁrst use the Lanczos algorithm to estimate
P (cid:48) (cid:28) P eigenvalues of H t and then smooth these estimated values into a continuous approximation to the discrete density ψ by convolving the values with a Gaussian kernel (Golub & Welsch, 1969; Ghorbani et al., 2019).
This yields a continuous function ˆ ψ ( t ) that approximates ψ .
See Appendix E for an example of estimating ˆ ψ for a simpliﬁed loss function.
Interpreting estimated spectral densities remains a challenging task: (Ghorbani et al., 2019) and (Yao et al., 2019) claim that a spectral density should be “smooth” (i.e., concentrated within a dense region) and consider the effect of different network architectures on density smoothness.
(Alain et al., 2019) argues that having a large concentration of negative eigenvalues can lead to inefﬁcient training, because most optimizers fail to leverage this local negative curvature.
(Papyan, 2020) relate the outlier structure of a
Hessian’s eigenvalues to the number of classes in a dataset.
Compared to CV problems, graph-structured datasets have a number of interesting attributes.
For example, a graph input has a variable number of node and edge features, which complicates the GN learning problem.
Furthermore, in contrast to the use of very deep networks in CV problems, increasing depth in graph networks has often not been found to be effective (see, e.g., discussion in (Godwin et al., 2022)).
Hessian-based information has the potential to inform some of these comparisons.
Note that the estimations of both the trace and the spectral densities are based on random quantities.
For trace, multiple vectors v must be sampled and the quantity v T H t ( θ ) v computed, and then convergence in mean can be assessed.
In practice, we ﬁnd that approximately 500 samples are sufﬁcient.
For spectral density, each Lanczos iterate is initialized as a random Gaussian vector.
We adapt an implementation of the Lanczos algorithm and smoothing process developed by (Ghorbani et al., 2019), in which we reorthogonalize the
Lanczos iterates during each step of estimation to promote numerical stability.
We follow (Yao et al., 2019) and use
P (cid:48)
= 100 , which is more than the P (cid:48)
= 90 that (Ghorbani et al., 2019) validate by comparing to an exact calculation of every eigenvalue of a smaller network.
2.2. Graph networks for property prediction
Let ˆ y : G → Y be a neural network parameterized by θ .
We follow (Sener & Koltun, 2018) in splitting θ into a set of shared weights θ sh , and a set of property-speciﬁc weights θ 1 , . . .
, θ T .
Then the predicted t th property is ˆ y t ( g ; θ sh , θ t ) =
f t ( z ( g ; θ sh ); θ t ) , where z : G → R D is a task-independent GN (Battaglia et al., 2018), and each f t : R D → R is a task-speciﬁc feedforward network.
We brieﬂy outline the functionality of our graph network z , which maps an input graph g to a graph-level representation vector u M .
First, the multigraph g is given a global feature u 0 ∈ R U that is initialized as a vector of ones, and, similarly to (Chen et al., 2019), the node features x v and edge features e v,v (cid:48) ,k are linearly projected:
x 0 v =
W V x v and e 0 v,v (cid:48) ,
k = W E e v,v (cid:48) ,k , where W V and W E are matrices learned during training.
Then information is propagated across the multigraph during m = 1 , . . .
, M messagepassing steps.
The m th message-passing step proceeds as follows: First, the edge features are updated:
[Equation]
where φ E is a feed-forward neural network, and cat is the concatenation operator.
Then node updates are collected.
For every node v , ﬁrst edge updates are calculated:
[Equation]
,v,k , where the ﬁrst summation is over the tuples ( v (cid:48) , k ) such that ( v, v (cid:48) , k ) is an edge for g , and the second summation is deﬁned similarly.
With these updates, the new node representation for v is given by
[Equation]
where φ V is a feed-forward neural network.
Finally, the global feature is updated:
[Equation]
where φ U is a feed-forward neural network, and each feature vector x mv , e mv,v (cid:48) ,k , u m is processed with a layer normalization (Ba et al., 2016) layer.
After M steps, the ﬁnal global state u M is fed into each property-speciﬁc network f t .
Our graph networks are implemented in the jraph 5 framework that is built on jax 6 .
We use flax 7 to implement component neural networks φ V , φ E , and φ U .
Because we seek to evaluate second-order derivative information, we ensure that our neural networks φ V , φ E , φ U are smooth functions with respect to their weights by using tanh activation functions.
Further details on speciﬁc hyperparameters used to instantiate models are given in Appendix A.
We impose on each property a loss function L t (here assumed to be mean squared loss for all t ), and we collect a training dataset of multigraph data points D
= { ( g n , y n ) } Nn =1 .
Then we solve the minimization problem
[Equation]
where
[Equation]
are task-speciﬁc losses and θ = { θ sh , θ 1 , . .
, θ T } .
We solve this minimization with stochastic gradient descent using the optax 8 implementation of AdamW (Loshchilov & Hutter, 2019) to choose stepsizes; we handle the potential difference in scales across properties by normalizing all targets prior to training and set µ t = 1 for all t .
Further details on the speciﬁcs of model training are given in Appendix B.
2.3. Data sources
We evaluate our graph network curvature assessment techniques on datasets from two scientiﬁc domains: materials science (crystals) and chemistry (molecules).
For crystals, our data featurization follows (Park & Wolverton, 2020); for molecules, it follows (Gilmer et al., 2017).
To reduce the complexity of our considered problem space, we choose to use simple featurizations (one-hot encoding of atom-types as node featurizations), but our methods are compatible with https other featurizations (e.g., hand-crafted node descriptors like those provided by rdkit ).
Further details about the data used are available in Appendix C.
For materials science, we collect data from Materials Project (MP) (Jain et al., 2013), which contains results of density functional theory (DFT) calculations for different crystals.
crystal records contain compositional and structural information, as well as some related properties.
Each crystal’s structure information is captured in a Crystallographic Information File (CIF), which we convert into a multigraph using the VoronoiNN function from pymatgen 9 , following the approach of (Park & Wolverton, 2020).
As crystals are periodic structures, this process yields multiple edges between many given pairs of nodes (Xie & Grossman, 2018;
Park & Wolverton, 2020; Sanyal et al., 2018)
(e.g., Figure 1).
As targets, we use several assessments of a crystal’s elastic properties: the Voigt-Reuss-Hill (Hill, 1952) calculation for shear ( G V RH ) and bulk ( K V RH ) modulus, both measured in units of GPa, as well as the isotropic Poisson ratio µ , a dimensionless quantity.
These properties are inherently coupled – µ can be calculated as a function of G V RH and K V RH .
Thus, an ML model that predicts all three serves as a test for how well it can learn underlying physical relationships.
A one-hot encoding of atom-type is used to obtain initial node features x v , and four bond-related properties calculated by pymatgen are used as edge features, resulting in a dataset of 10,500 crystal structures with known elastic properties.
For chemistry, we use the QM8 (Ramakrishnan et al., 2015) dataset of organic molecules.
We use the rdkit packageto convert SMILES strings (Weininger, 1988) into molecular structure graphs (see, e.g., Figure 2).
As targets, we use the ﬁrst and second transition energies, E 1 and E 2 , and the ﬁrst and second oscillator strengths f 1 and f 2 .
QM8 contains several versions of transition energies and oscillator strengths that have been calculated via different levels of
DFT.
We use the predictions of the approximate coupledcluster (CC2) (H ¨ attig & Weigend, 2000) method as our targets, as these are treated as the “ground truth” in (Ramakrishnan et al., 2015).
A one-hot encoding of atom-type is used to obtain initial node features x v , and a one-hot encoding of bond-type is used to obtain initial edge features e v,v (cid:48) ,k .
Note that, in Section 2.2, we describe our input data points g as directed multigraphs.
However, for our actual data points, the edges are not inherently directed – bonds are symmetric.
Thus, we duplicate each edge feature (i.e., e v,v (cid:48) ,k = e v (cid:48) ,v,k for all edges) during data preprocessing.
3. Results
We demonstrate the application of our curvature-assessment techniques (Section 2.1) on trained GNs (Section 2.2) in two application domains: materials science and chemistry (Section 2.3).
Although here we focus on GNs (Battaglia et al., 2018), our assessment techniques are applicable to the broad family of graph neural networks used in property prediction tasks.
Figure 3 shows training and test set errors for each dataset.
These results do not necessarily align with domain intuition, suggesting that the models do not leverage scientiﬁc knowledge in their learned representations.
For example, for MP, the test set error for Poisson ratio is higher than that of G V RH or K V RH , despite the fact that Poisson ratio can be calculated as a function of G V RH and K V RH .
This is interesting and suggests that the GNs are not fully leveraging known scientiﬁc knowledge.
For QM8, the oscillator strengths f 1 and f 2 have the lowest and highest, respectively, test set errors.
Next we summarize the curvature of the loss functions using the trace of each task’s loss’s Hessian.
In Figure 4, we show that, despite the general decreasing training error, the curvature of each property’s loss surface, as measured by the trace of its Hessian, is highly variable.
In particular, both datasets show high heterogeneity across properties and across training epochs in their estimated traces.
Our observations here do not align with prior work that has analyzed
Hessian traces for CV problems (Yao et al., 2019).
In those works, traces increased monotonically during training and were consistently positive.
Here, our estimated traces often ﬂip between being positive and negative.
The Hessian trace is a high-level summary of the curvature of a loss surface; for a more granular examination, we estimate the spectral densities of each property’s Hessian.
In Figure 6a in D, we show that the spectral densities are similarly variable and often feature the presence of outliers that brieﬂy occur during training.
These outliers vary across properties.
For the QM8 dataset, we zoom in on the range of eigenvalues where most density is concentrated in Figure 5.
Similar to (Yao et al., 2019; Ghorbani et al., 2019), we see that most density is concentrated near 0.
This suggests that there is a great deal of redundancy in the latent spaces learned by these models, and their true dimensionality is likely signiﬁcantly less than the full P parameters of the θ sh weight vector.
Unlike these works, the spectral densities are more symmetric around 0.
This matches Figure 4, since the
Hessian trace (sum of all its eigenvalues) varies between being very positive and very negative.
The exact spectral densities vary across properties, even at the end of training.
[Figure]
4. Discussion
In this work, we have posited that the performance of multitask GNs for property prediction may be stymied by underlying variation in property-speciﬁc loss-surface curvatures.
Some existing work considers this question (Yu et al., 2020) in a simpliﬁed theoretical setting, but, to the best of our knowledge, no one has empirically investigated the Hessian properties of multi-task GNs.
In two domains – chemistry and materials science – our results suggest that loss surface curvature varies across each modeled property.
In order to assess curvature without calculating the full Hessian, we build on recent work that uses matrix-free methods to estimate Hessian properties (Alain et al., 2019; Yao et al., 2019; Ghorbani et al., 2019; Papyan, 2020).
We extend their results by considering a novel type of learning problem – multi-output scalar regression – and a novel class of neural network – graph networks.
Our results echo some previous results – the majority of a Hessian’s spectral density is concentrated near 0 – and diverges in other respects – Hessian properties appear far noisier and more variable for GNs than for CV tasks.
We leave further investigation of this question to future work.
Potential explanations include the difference in loss functions for regression vs. classiﬁcation tasks, a higher level of noise in the datasets we examined, and some special characteristic of GN vs. other neural network architectures.
We have here focused on a speciﬁc but representative subset of GN prediction problems, but considerable potential variation exists.
For example, many recent GN architectures incorporate a notion of equivariance (Satorras et al., 2021;
Gasteiger et al., 2020) into their feature-extraction models, and this might impact their curvature in different ways.
In addition, we have not evaluated how the choice of optimizer (e.g., stochastic gradient descent vs. Adam (Kingma & Ba, 2014) vs. AdamW) impacts the curvature properties of a learned loss surface.
Our curvature assessment enable several future research directions.
First, our analysis here was primarily empirical.
The phenomena we identify here (a diversity of curvatures across multi-task loss functions) and the previouslyobserved degradation of performance for multi-task models (Gasteiger et al., 2020) could be connected by a theoretical justiﬁcation.
Similarly, we ﬁnd that curvature properties of GNs appear to be much noisier and more variable than the properties of other network architectures, and we currently lack a theoretical justiﬁcation of why.
Intermediate steps might entail investigating the curvature properties of MTL methods that, in other domains, do out-perform single-task models (e.g., (Sener & Koltun, 2018)).
Existing work in analyzing curvature for computational geometry (e.g. (Goldman, 2005)) might provide techniques to build upon. Acknowledgements
This work was supported by internal research and development funding from the Research and Exploratory Development Mission Area of the Johns Hopkins University
Applied Physics Laboratory.
Thanks to Christopher Ratto for assistance in editing and improving the work.
A. Model speciﬁcs
We use M = 5 message-passing steps.
Node, edge, and global features are projected into a 64-dimensional feature space.
φ V and φ E have two layers with 256 units and a skip connection, followed by an output layer of 64 units.
φ U has two layers with 192 units and a skip connection, followed by an output layer of 64 units.
All networks use tanh activations.
B. Training speciﬁcs
Models are trained for 512 epochs using the AdamW (Loshchilov & Hutter, 2019) optimizer and the default hyperparameters used in the optax implementation.
The initial rate is set as 10 − 3 , with an exponential decay rate of 0.997 applied every epoch after the ﬁrst 256 epochs.
C. Data
We summarize our datasets in Table 1.
We scraped Materials Project for crystal records present in it as of October 2020 using the MPRester class from the pymatgen package.
The Supplemental Information contains a table of MP IDs used in this study.
The raw entries of
G V RH , K V RH , and Poisson ratio µ contained several anomalously small and large values, so we removed entries with values less than the 5th percentile or greater than the 95th percentile of obtained elastic properties.
This left us with a total dataset of 10,500 crystals.
Summary statistics for the targets used are given in Table 2.
Furthermore, G V RH and K V RH were log-transformed to create a more normal distribution.
Roughly 70% of the dataset (7,400 crystals) was used as training data, and all targets were standardized prior to training.
As initial node features, we used a one-hot encoding of atom type.
For initial edge features, we used four features calculated by pymatgen : area , face dist , solid angle , and volume .
Following similar work in applying graph neural networks to crystals (Xie & Grossman, 2018), we discretize these four features based on deciles.
We use the version of QM8 hosted by MoleculeNet (Wu et al., 2018), except that we drop molecules with negative oscillator energies.
Summary statistics for the targets are given in Table 3.
As initial node features, we use one-hot encodings of atom type.
As initial edge features, we use one-hot encodings of bond type.
Roughly 70% of the dataset (15,300 molecules) was used as training data [Table] and all targets were standardized prior to training.
[Table]
[Table] [Table]
D. Supplemental ﬁgures
[Figure]
E. Example spectrum estimation
We consider a loss function deﬁned by L ( x ) = 1 2 x T Ax , where A = 1 2( B + B T ) , for a matrix B with entries sampled iid from a standard normal distribution.
In this case, the Hessian is given by A
and so its estimated eigenvalues can be compared to its true eigenvalues.
In Figure 7, we plot the results of a sample calculation, for a 1 , 000 × 1 , 000 matrix.
We show that the estimated density has reasonable similarity to the true distribution, and the estimated trace also matches the true trace.