# Optimizing Brain Tumor Classification: A Comprehensive Study on Transfer
Learning and Imbalance Handling in Deep Learning Models

Raza Imam Aligarh Muslim University, Aligarh, India  Mohamed Bin Zayed
University of Artificial Intelligence, Abu Dhabi, UAE  Mohammed Talha Alam
Jamia Hamdard University, New Delhi, India  Mohamed Bin Zayed University of
Artificial Intelligence, Abu Dhabi, UAE

###### Abstract

Deep learning has emerged as a prominent field in recent literature,
showcasing the introduction of models that utilize transfer learning to
achieve remarkable accuracies in the classification of brain tumor MRI images.
However, the majority of these proposals primarily focus on balanced datasets,
neglecting the inherent data imbalance present in real-world scenarios.
Consequently, there is a pressing need for approaches that not only address
the data imbalance but also prioritize precise classification of brain cancer.
In this work, we present a novel deep learning-based approach, called Transfer
Learning-CNN, for brain tumor classification using MRI data. The proposed
model leverages the predictive capabilities of existing publicly available
models by utilizing their pre-trained weights and transferring those weights
to the CNN. By leveraging a publicly available Brain MRI dataset, the
experiment evaluated various transfer learning models for classifying
different tumor types, including meningioma, glioma, and pituitary tumors. We
investigate the impact of different loss functions, including focal loss, and
oversampling methods, such as SMOTE and ADASYN, in addressing the data
imbalance issue. Notably, the proposed strategy, which combines VGG-16 and
CNN, achieved an impressive accuracy rate of 96%, surpassing alternative
approaches significantly. Our code is available at Github.

##  1 Introduction

Impactful solutions are being offered by the context-aware deployment of deep
learning methodologies to enhance medical diagnostics. The World Health
Organization (WHO) states that a correct diagnosis of a brain tumor entails
its discovery, localization, and classification based on its degree, kind, and
severity. This research comprises finding the tumor, grading it according to
type and location, and classifying it according to grade in the diagnosis of
brain tumors using magnetic resonance imaging (MRI). This approach has
experimented with using several models rather than a single model for
classification task in order to categorize brain MRI data [Veeramuthu et al.,
2022]. Deep learning and transfer learning models have been proposed with
higher accuracies in recent literature for classifying Brain Tumor MRI images.
However most of such proposals are focused on balanced data. Hence, approaches
to account for the imbalance in data, as well as focusing on the precise
classification of brain cancer in real world scenarios, are needed. We
introduce a ’Transfer Learning + CNN model’ rather of just ’Transfer Learning
model with Fine Tuning’ and compared 8 of such Transfer Learning models using
various approaches on imbalance MRI images [Li et al., 2022].

Figure 1: Conventional Classification of Brain Tumors from MRI images using
Convolutional Neural Networks

Using three pathogenic forms of brain tumor (glioma, meningioma, and pituitary
tumor), we aim to propose an accurate and automated classification scheme in
this work. Towards this classification task, our goal is to empirically assess
how well the most common benchmark models perform. Our goal of offering the
finest classification model will open up new research directions in terms of
choosing the right model for real-world brain tumor classification deployment.
Utilizing tested models, the acquired characteristics are classified. The
approach of our experiments that we will be acquiring is also represented in
Figure 1. Subsequently, a thorough assessment of the suggested system is
performed utilizing several effective evaluation metrics of classification
tasks along with comparing several analytical factors, such as how well each
model performs with less training samples from practicality aspect and how
overfitting with lower training samples affects performance of the classifier.

###  1.1 Motivation

Early brain tumor identification and classification represent a significant
area of research in the field of medical imaging. It helps in choosing the
most appropriate line of action for treatments to save patients’ lives. In
both children and adults, brain tumors are regarded as one of the most severe
disorders. Brain tumors account for 85% to 90% of all major malignancies of
the Central Nervous System (CNS). An estimated 11,700 people receive a brain
tumor diagnosis each year. For those with a malignant brain or CNS tumor, the
5-year survival rate is around 34% for males and 36% for women [Wang et al.,
2020]. It is difficult to treat brain tumors as we know that our brain has a
very complex structure having tissues that are linked to each other in a
complex manner. Often, producing MRI results is very difficult and time-
consuming in underdeveloped nations due to a shortage of skilled medical
professionals and a lack of understanding of malignancies.

Depending on the tumor’s severity—that is, its location, size, and
type—different treatment methods are possible. The most common technique for
treating brain tumors at the moment is surgery since it has no adverse
implications on the brain. There are different medical imaging techniques that
are used to view the internal structures of a human body in order to discover
any abnormalities Imam et al. [2022]. The most often used of them to identify
brain tumors is Magnetic Resonance Imaging (MRI) since it can show
abnormalities that may not be seen or just dimly visible on computed
tomography (CT) [Kibriya et al., 2022]. However, the rush of patients makes it
difficult, chaotic, and perhaps error-prone to manually review these images.
Automated classification methods based on machine learning and artificial
intelligence have regularly outperformed manual classification in terms of
accuracy in order to solve such issues. Therefore, recommending a system that
does detection and classification utilizing deep learning algorithms employing
the above-mentioned benchmark models would be helpful for radiologists and
other medical professionals.

###  1.2 Contributions

In this work, we focus on imbalance problem of Brain Tumor Classification as
it is a real-world scenario in deployment. To solve this imbalance problem, we
experiment with several loss functions including focal loss. Along with this,
a comparative study is conducted on the performance of different oversampling
methods including augmentation, SMOTE, and ADASYN. The models are experimented
on are 8 Transfer Learning-CNN models that incorporates the base Transfer
learning model, followed by the integration of proposed CNN layers. Following
a detailed evaluation of each Transfer Learning-CNN model, we finally conclude
VGG16-CNN is the novel proposal of this study.

  * •

Developed a Transfer learning-CNN framework for brain tumor MRI
classification, utilizing pre-trained models, where the produced weights
transfer to an 8-layer CNN head for effective training.

  * •

Employed 8 different transfer learning models integrated with CNNs to increase
the classification accuracy on different types of brain cancer (no tumor,
glioma, meningioma, and pituitary cancer).

  * •

Experimented with 5 different approaches to deal with imbalanced datasets such
as- Changing loss functions: (1) Focal loss (2) Cross Entropy, and
Oversampling methods: (3) Data Augmentation, (4) SMOTE (5) ADASYN.

  * •

Assessed empirical evaluation of the models under the different approaches on
metrics including Accuracy, Precision, Recall, and F1 Score.

##  2 Related Works

The classification of brain tumors using MRI data has been the subject of
numerous studies based on convolutional neural networks in recent years. Many
of these methods make use of hybrid approaches, and many also offer technical
variations on widely used deep learning models Alam et al. [2021]. In [Deepak
and Ameer, 2019], the authors describe a classification method for the 3-class
classification issue that combines transfer learning with GoogleNets. They
used a number of evaluation metrics, with a mean classification accuracy of
98%, including area under the curve (AUC), precision, recall, F-score, and
specificity. Abd El Kader et al. [2021] have made use of the benefit of
differential CNN by generating extra differential feature maps. With the
capacity to categorize a sizable database of pictures with high accuracy of
about 98%, their approach demonstrated a considerable improvement for the
brain MRI classification problem. In [Raza et al., 2022], the authors proposed
a hybrid architecture by adopting GoogleNet as a based CNN model while
tweaking the last few layers for the specific Brain Tumor Classification.
Their proposal attained the classification accuracy of 99.67%. Moreover, Irmak
[2021] conducted a multi-class study of Brain Tumor MRI Images as they propose
a CNN model for early diagnoses purposes with fully optimized framework.
Compared to the conventional CNN models, their solution attained an accuracy
of 98.14%.

Figure 2: Difference between traditional Transfer Learning vs Transfer
Learning-CNN

##  3 Methodology

###  3.1 Transfer Learning and CNN

We observed that using the most recent transfer learning models that have been
presented in the most recent literature produced relatively decent results
when trained on the balanced Brain MRI datasets. This empirical evaluation of
brain tumor classification is focused on the Imbalance problem. However, in
training such Transfer Learning models with fine tuning on the imbalance
dataset did not show as good results in comparison to the Transfer models
integrated with newly added CNN layers [Li et al., 2022]. Since the later case
demonstrated improved predicted accuracies on the unbalanced dataset, we
experiment with 8 of these Transfer Learning models that are combined with
newly added CNN layers. The ImageNet training dataset was used to train the
traditional pre-train models, namely VGG16, EfficientNetB0, EfficientNetB3,
ResNet50, DenseNet201, MobileNet, GoogleNet, XceptionNet. For effective
training, each model’s trained parameter weights were transferred to the newly
added layers of the additional CNN model. The CNN model was then fine-tuned
using the brain tumor augmented MRI data set for final classification into 4
classes. CNN, VGG16-CNN, EfficientNetB0-CNN, EfficientNetB3-CNN, ResNet50-CNN,
DenseNet201-CNN, MobileNet-CNN, GoogleNet-CNN, XceptionNet-CNN are the
resultant models we empirically evaluated the training performance on the
imbalanced dataset using various approaches.

###  3.2 Proposed Method

The proposed Transfer Learning-CNN (VGG-CNN) model incorporates the base
Transfer learning model, i.e., VGG16, followed by the basic structure of the
CNNs. It might take days to weeks to train a raw CNN completely from scratch,
making it a difficult task [Raza et al., 2022]. Therefore, it would be
preferable to train the suggested deep learning approach using a pre-trained
classifier rather than creating a new deep learning classifier from start
[Petmezas et al., 2022]. Additionally, in terms of predictive accuracy, our
Transfer Learning models combined with CNN performed significantly better on
the imbalanced Brain MRI dataset than the base CNN as well as the base
Transfer Learning model with Fine tuning. In order to do this, we chose the
most accurate current Transfer Learning model on the balanced dataset, VGG16,
as our foundation model. CNN architecture VGG16 was employed to win the 2014
ILSVR (Imagenet) competition. It is regarded as one of the best vision model
architectures created to date. The most distinctive feature of VGG16 is that
it prioritized having convolution layers of 3x3 filters with a stride 1 and
always utilized the same padding and maxpool layer of 2x2 filters with a
stride 2. Throughout the entire architecture, convolution and max pool layers
are arranged in the same manner. It concludes with two fully connected layers
(FC) and a softmax for output.

Figure 3: The proposed framework for the classification of brain tumor MRIs.
Here, the pre-trained CNN models (VGG16-CNN, EfficientNetB0-CNN,
EfficientNetB3-CNN, ResNet50-CNN, DenseNet201-CNN, MobileNet-CNN, GoogleNet-
CNN, and XceptionNet-CNN) are trained using the Imagenet dataset, and the
produced weights of these pre-trained models are individually transferred to
the suggested CNN model for effective training

Moreover, the final four layers of VGG-16 were dropped from the planned VGG-
CNN architecture, and 8 new levels were added in their place. After these
adjustments, there were 148 layers overall instead of 144. Addition to the
layers of transfer learning model, the first convolution layer of the CNN
portion, uses a filter size of 3x3 with the depth of 32, which immediately
reduces the image size followed by a max pooling layer of 2x2 filter. The
second convolution layer has a depth of 64 with the same filter size of 3x3,
followed by a max pooling layer of 2x2 filter. Again, a 2D convolution layer
of depth 32 with filter size 3x3 is incorporated followed by a Flatten layer.
Dense layer of 16 hidden units with relu activation and the final output layer
of 4 units with softmax activation are used. Compared to the initial 4 layers
of transfer learning models, adding the extra convolutional layers in the CNN
portion gave us more detailed, accurate, and robust features. These 3 later
convolutional layers extracted high-level features compared with the initial
layer, which extracted low-level features. The choice of this additional
8-layer architecture as the head on the pre-trained network was based on
several considerations. We aimed to strike a balance between model complexity
and performance. This 8-layer architecture was found to be effective in
capturing the necessary features for our specific task, while not being overly
complex, which could lead to overfitting or increased computational
requirements. Our method’s superior robustness is supported by evaluating
various architectures during experimentation, where the 8-layer head
consistently demonstrated better performance across multiple metrics and
outperformed other architectures in diverse datasets and scenarios. Hence,
compared to the traditional Transfer learning model, the proposed Transfer
Learning-CNN model achieves higher accuracies in comparison as more intricate,
exclusionary, and deep features have been acquired in the proposed method. The
specific architecture of the proposed methodology is also shown in Figure 3.

###  3.3 Experiments

The experiments have been performed on a combination of three distinct
datasets taken from Kaggle namely figshare, SARTAJ and Br35H. This accumulated
custom imbalance dataset contains about 4200 images of human brain MRI images
which are classified into 4 classes: no-tumor (1760), glioma (858), meningioma
(1265) and pituitary (341) cancer, whereas no-tumor class images were taken
from the Br35H dataset. The dataset images have been normalized and resized as
a part of pre-processing according to the input size of the proposed model,
i.e., 128x128. In addition, the accumulated dataset is divided into train and
test with 90:10 ratio for validation purpose. The results mentioned in this
study only apply to the test set. The testing dataset is then used to evaluate
all of the predictive metric claims made about the models in the results and
discussion section.

Figure 4: Data distribution of different classes for imbalanced Brain MRI
dataset Figure 5: Train-test split of the original dataset

Predictive modeling is challenged by imbalanced classifications because the
majority of machine learning methods for classification were devised on the
premise that there should be a similar number of samples in each category. The
distribution of examples among the recognized classes is skewed or biased in
imbalance classes. These kind of issues with a strong to weak bias in the
dataset are typical in real-world applications. As a result of such imbalance,
models perform poorly in terms of prediction, particularly for the minority
class [Deepak and Ameer, 2019], [Deepak and Ameer, 2022]. This is a
complication since, in general, the minority class is more significant and, as
a result, the issue is more susceptible to errors in classifying for the
minority class than the classes with the higher samples. In the real-world
scenarios of a model deployment, imbalanced classification is a huge
complication [Badža and Barjaktarović, 2020]. Hence, in our experiments, we
tried to empirically evaluate majority of the existing solutions that can be
used for imbalanced or long tail classification problems.

Figure 6: Distribution of the incorporated approaches to resolve data
imbalance problem Figure 7: Implementation of data imbalance approaches
throughout the phases of experimentation on models

Focal Loss and Cross Entropy. For each of the 8 deep learning models, we
comparatively examined 5 distinct approaches on each of them in order to
address the imbalance scenarios that arise during the classification of brain
tumors. We explore with changing the loss function in our modeling because
it’s one of the main reasons that simple examples will divert training in an
unbalanced class scenario. We compared the results of test accuracy on our
models by changing Cross Entropy with Focal loss in order to obtain better
results. Focal Loss, one of the effective solutions in terms of loss
functions, deals with this imbalance issue and is created in a way that allows
the network to concentrate on training the challenging cases by lowering the
loss (or "down-weight") for the simple examples. In other words, Focal loss
reduces the importance of the simple examples and emphasizes the difficult
ones; thus, smaller class counts would have heavier weights in errors than
cross-entropy. The Cross-Entropy loss is multiplied by a modifying factor in
focal loss [Chatterjee et al., 2022]. When a sample is incorrectly classified,
the modulating factor is close to 1, p is low, and the loss is unchanged. The
modulating factor gets closer to zero as p approaches 1 and the loss for
correctly categorized samples gets down-weighted.

$FL(P_{t})=-\alpha_{t}(1-P_{t})^{\gamma}log(P_{t})$

Here, $P_{t}$ = P if $\gamma$ = 1, else (1-P); where $\gamma$ = 0,
$\alpha_{t}$ = 1 then FL is Cross Entropy Loss.

In addition, oversampling, which uses artificial data creation to increase the
number of samples in the data set, is one of the most fundamental approaches
in the state of the art to handle the imbalance problem, other than loss
functions [Xie et al., 2022]. By producing synthetic observations based on the
minority observations already there, oversampling aims to grow the minority
class so that the data set becomes balanced. In our experimental evaluations,
we used 3 oversampling approaches, namely Augmentation, SMOTE, and ADASYN
methods.

Data Augmentation. Data augmentation is used to expand the volume of data by
introducing slightly altered versions of either existing data or brand-new
synthetic data that is derived from available data. It also serves as a
regularization term and aids in lowering overfitting when a model is being
trained [Irmak, 2021]. In our experiments, we tried to augment the samples
from minority classes to achieve the number of the samples equal to the
majority class. The augmentation tasks we applied to such minority class
samples are brightness, contrast, and sharpness alteration each with a random
intensity of 80% to 120%. This was followed by the normalization process as a
part of preprocessing.

SMOTE. Second oversampling method we implemented on our dataset was SMOTE
(Synthetic Minority Oversampling Technique). By creating artificial data
samples that are marginally different from the current data points on the
basis of the existing data points, SMOTE performs oversampling. Following is
how the SMOTE algorithm operates:

  * •

A random sample is picked from the minority class.

  * •

Locate the $k$ nearest neighbors for the observations in this minority class
sample.

  * •

The vector (line) between the existing data point and one of those neighbors
will then be determined using that neighbor.

  * •

The vector is then multiplied by a randomized range unit between 0 and 1.

  * •

We combine this with the existing data point to get the synthetic data sample
in the space.

ADASYN. ADASYN is an improvised version of SMOTE. It accomplishes the same
things as SMOTE, albeit slightly better. It then adds a random tiny value to
the points to make it more realistic after creating those synthetic samples
[Raza et al., 2022]. In other words, the sample’s variance, or degree of
dispersion, is a little higher than its linear correlation to the parent.
Moreover, ADASYN adaptively change the weights of different minority samples
to compensate for the skewed distributions. Difference between SMOTE and
ADASYN oversampling approaches is also shown in Figure 8.

Figure 8: Working of SMOTE vs ADASYN algorithm

We examined each of our 8 TL-CNN models on each of these above-mentioned 5
approaches - Focal Loss, Cross Entropy, Data Augmentation, SMOTE, and ADASYN,
with an aim to achieve higher test accuracy on the imbalance dataset. Focal
loss is initially compared with the Cross Entropy loss on unbalanced dataset,
whereas Data Augmentation, SMOTE, and ADASYN are compared on the unbalanced
dataset too but with cross entropy as these later 3 approaches are
oversampling approaches [Abd El Kader et al., 2021]. With the parameters
listed in Table 1, we conducted tests using a trial-and-error methodology. In
order to determine the best convergence for each TL-CNN model, we continuously
tracked the development of the training validation accuracy and error. With a
mini-batch size of 20 images and an initial learning rate ranging between
0.001 and 0.0001 depending on the model, we utilized mini-batch GD to train
the TL-CNN models. To get the results that are explained in the following
section, the evaluated models were trained on 5–12 epochs for brain tumor
classification.

Table 1: Several parameters utilised during the training phase

Parameter | Values |
---|---|---
Mini Batch Size | 20 |
Number of Epochs | 5 to 12 |
Learning Rate(s) | 0.001 to 0.0001 |
Shuffle | True |
Steps Per Epoch | Train Size (1600) / 20 |
Dense Activations | RELU |
Optimizer | ADAM |

###  3.4 Evaluation

For the typical assessment of a classifier, numerous performance metrics are
specified, including Classification Accuracy, Precision, Recall, and F1-Score
Alam et al. [2022]. The evaluation employs criteria besides the total
classification accuracy due to the unbalanced dataset. The mean of correctly
classified samples from each class is used to compute balanced accuracy. A
class’s F-score is determined by taking the harmonic mean of its precision and
recall values.

$Precision=\frac{\text{True Positive}}{\text{True Positive}+\text{False
Positive}}$

$Recall=\frac{\text{True Positive}}{\text{True Positive}+\text{False
Negative}}$

$Accuracy=\frac{\text{TN}+\text{TP}}{\text{TN}+\text{FP}+\text{TP}+\text{FN}}$

$F1\text{
Score}=2\times\frac{\text{Precision}\times\text{Recall}}{\text{Precision}+\text{Recall}}$

When the test dataset has an identical amount of samples from each category,
classification accuracy is a useful metric to assess performance.
Nevertheless, the dataset we aim to use for this categorization problem under
discussion is adequately imbalance. This calls for a more thorough assessment
of the suggested system using more evaluation metrics. To evaluate the
effectiveness of our tumor categorization method, we employed other metrics
including Precision, Recall, and F1-Score. Tables 2, 3, 4, and 5 presents a
summary of such metrics on each of the 8 models in terms of various approaches
we employed to resolve imbalance classification. The effectiveness of the
relevant implemented models that have been implemented in our categorization
are empirically evaluated using these classification metrics. Due to the
existence of imbalance among the 4 classes, we also showed individual
precision values for each class that have been implemented using different
approach.

Table 2: Evaluation of Integrated frameworks performance on imbalanced dataset
when trained with Cross Entropy Loss. The space complexity of each model is
the number of trainable parameters. M = Million. The space complexity
increases with increasing number of trainable parameters. The time complexity
is the training time (in hours) of the models

Model |  Space Complexity |  Time (hrs) Complexity | LR | Acc. | Precision | Recall | F1 Score
---|---|---|---|---|---|---|---
P | M | G | NT | P | M | G | NT | P | M | G | NT
CNN | 440k | 0.091 | 0.001 | 0.8125 | 1.00 | 0.67 | 0.73 | 0.98 | 0.48 | 0.91 | 0.51 | 0.96 | 0.65 | 0.77 | 0.60 | 0.97
VGG16-CNN | 15.7M | 0.26 | 0.0001 | 0.83 | 0.73 | 0.66 | 0.94 | 0.99 | 0.73 | 0.66 | 0.41 | 1.00 | 0.60 | 0.78 | 0.57 | 1.00
EfficientNetB0-CNN | 4.05M | 0.27 | 0.001 | 0.37 | 0.00 | 0.17 | 0.00 | 0.39 | 0.04 | 0.04 | 0.00 | 0.86 | 0.00 | 0.07 | 0.00 | 0.54
EfficientNetB3-CNN | 10.8M | 0.47 | 1.6e-5 | 0.60 | 0.50 | 0.60 | 0.38 | 0.65 | 0.06 | 0.46 | 0.27 | 0.96 | 0.11 | 0.52 | 0.31 | 0.78
ResNet50-CNN | 23.5M | 0.31 | 0.001 | 0.78 | 1.00 | 0.87 | 0.60 | 0.83 | 0.39 | 0.59 | 0.82 | 0.97 | 0.56 | 0.70 | 0.69 | 0.90
DenseNet201-CNN | 51.8M | 0.57 | 0.001 | 0.47 | 0.33 | 0.43 | 0.67 | 0.49 | 0.23 | 0.32 | 0.05 | 0.82 | 0.27 | 0.37 | 0.10 | 0.61
MobileNet-CNN | 3M | 0.126 | 0.001 | 0.30 | 0.00 | 0.30 | 0.00 | 0.00 | 0.00 | 1.00 | 0.00 | 0.00 | 0.00 | 0.46 | 0.00 | 0.00
GoogleNet-CNN | 6M | 0.14 | 0.0001 | 0.77 | 0.71 | 0.64 | 0.58 | 0.97 | 0.65 | 0.67 | 0.56 | 0.96 | 0.68 | 0.65 | 0.57 | 0.97
XceptionNet-CNN | 25M | 0.15 | 0.0001 | 0.83 | 0.90 | 0.66 | 0.89 | 0.99 | 0.61 | 0.95 | 0.42 | 0.99 | 0.73 | 0.78 | 0.57 | 0.99

Table 3: Evaluation of Integrated frameworks performance on imbalanced dataset
when trained with Focal Loss

Model |  Space Complexity |  Time (hrs) Complexity | LR | Acc. | Precision | Recall | F1 Score
---|---|---|---|---|---|---|---
P | M | G | NT | P | M | G | NT | P | M | G | NT
CNN | 440k | 0.084 | 0.001 | 0.80 | 0.88 | 0.68 | 0.61 | 0.99 | 0.45 | 0..81 | 0.56 | 0.98 | 0.60 | 0.74 | 0.59 | 0.98
VGG16-CNN | 15.7M | 0.267 | 0.0001 | 0.85 | 0.93 | 0.68 | 0.88 | 1.00 | 0.87 | 0.98 | 0.37 | 0.98 | 0.90 | 0.81 | 0.52 | 0.99
EfficientNetB0-CNN | 4.05M | 0.27 | 0.001 | 0.42 | 0.00 | 0.00 | 0.00 | 0.42 | 0.00 | 0.00 | 0.00 | 1.00 | 0.00 | 0.00 | 0.00 | 0.59
EfficientNetB3-CNN | 10.8M | 0.47 | 1.6e-5 | 0.32 | 0.00 | 0.22 | 0.18 | 0.37 | 0.00 | 0.14 | 0.09 | 0.62 | 0.00 | 0.17 | 0.12 | 0.46
ResNet50-CNN | 23.5M | 0.31 | 0.001 | 0.62 | 0.69 | 0.45 | 1.00 | 0.97 | 0.71 | 0.95 | 0.17 | 0.59 | 0.70 | 0.61 | 0.29 | 0.73
DenseNet201-CNN | 51.8M | 0.57 | 0.001 | 0.66 | 0.57 | 0.50 | 0.24 | 0.96 | 0.39 | 0.68 | 0.15 | 0.94 | 0.46 | 0.58 | 0.19 | 0.95
MobileNet-CNN | 3M | 0.13 | 0.001 | 0.42 | 0.00 | 0.00 | 0.00 | 0.42 | 0.00 | 0.00 | 0.00 | 1.00 | 0.00 | 0.00 | 0.00 | 0.59
GoogleNet-CNN | 6M | 0.14 | 0.0001 | 0.78 | 0.73 | 0.62 | 0.74 | 0.97 | 0.61 | 0.88 | 0.33 | 0.97 | 0.67 | 0.72 | 0.46 | 0.97
XceptionNet-CNN | 25M | 0.20 | 0.0001 | 0.81 | 0.83 | 0.62 | 0.84 | 1.00 | 0.61 | 0.92 | 0.33 | 1.00 | 0.70 | 0.74 | 0.48 | 1.00

Table 4: Evaluation of Integrated frameworks performance on imbalanced dataset
when trained with SMOTE oversampling

Model |  Space Complexity |  Time (hrs) Complexity | LR | Acc. | Precision | Recall | F1 Score
---|---|---|---|---|---|---|---
P | M | G | NT | P | M | G | NT | P | M | G | NT
CNN | 440k | 0.11 | 0.001 | 0.82 | 0.86 | 0.71 | 0.71 | 0.93 | 0.77 | 0.83 | 0.45 | 0.99 | 0.81 | 0.76 | 0.55 | 0.96
VGG16-CNN | 15.7M | 0.52 | 0.0001 | 0.86 | 0.96 | 0.71 | 0.83 | 0.99 | 0.81 | 0.96 | 0.50 | 0.97 | 0.88 | 0.81 | 0.62 | 0.98
EfficientNetB0-CNN | 4.05M | 0.45 | 0.001 | 0.47 | 0.00 | 0.38 | 0.00 | 0.51 | 0.00 | 0.41 | 0.00 | 0.82 | 0.00 | 0.39 | 0.00 | 0.63
EfficientNetB3-CNN | 10.8M | 1.04 | 1.6e-5 | 0.50 | 0.00 | 0.39 | 0.00 | 0.88 | 0.00 | 0.97 | 0.00 | 0.51 | 0.00 | 0.55 | 0.00 | 0.64
ResNet50-CNN | 23.5M | 0.77 | 0.001 | 0.66 | 1.00 | 0.97 | 0.37 | 0.96 | 0.19 | 0.27 | 0.95 | 0.89 | 0.32 | 0.42 | 0.54 | 0.92
DenseNet201-CNN | 51.8M | 2.08 | 0.001 | 0.86 | 0.96 | 0.70 | 0.97 | 0.99 | 0.77 | 0.98 | 0.45 | 0.99 | 0.86 | 0.82 | 0.61 | 0.99
MobileNet-CNN | 3M | 0.27 | 0.001 | 0.42 | 0.00 | 0.00 | 0.00 | 0.42 | 0.00 | 0.00 | 0.00 | 1.00 | 0.00 | 0.00 | 0.00 | 0.59
GoogleNet-CNN | 6M | 0.14 | 0.0001 | 0.79 | 0.88 | 0.68 | 0.57 | 1.00 | 0.74 | 0.66 | 0.72 | 0.93 | 0.81 | 0.67 | 0.64 | 0.96
XceptionNet-CNN | 25M | 0.27 | 0.0001 | 0.83 | 0.91 | 0.66 | 0.82 | 1.00 | 0.68 | 0.94 | 0.41 | 0.99 | 0.78 | 0.78 | 0.55 | 1.00

Table 5: Evaluation of Integrated frameworks performance on imbalanced dataset
when trained with ADASYN oversampling

Model |  Space Complexity |  Time (hrs) Complexity | LR | Acc. | Precision | Recall | F1 Score
---|---|---|---|---|---|---|---
P | M | G | NT | P | M | G | NT | P | M | G | NT
CNN | 440k | 0.12 | 0.001 | 0.83 | 0.83 | 0.75 | 0.64 | 0.97 | 0.77 | 0.77 | 0.59 | 1.00 | 0.80 | 0.76 | 0.61 | 0.98
VGG16-CNN | 15.7M | 0.60 | 0.0001 | 0.87 | 0.91 | 0.75 | 0.81 | 1.00 | 0.94 | 0.92 | 0.54 | 0.99 | 0.92 | 0.82 | 0.65 | 0.99
EfficientNetB0-CNN | 4.05M | 0.63 | 0.001 | 0.42 | 0.00 | 0.00 | 0.00 | 0.42 | 0.00 | 0.00 | 0.00 | 1.00 | 0.00 | 0.00 | 0.00 | 0.59
EfficientNetB3-CNN | 10.8M | 1.11 | 1.6e-5 | 0.52 | 1.00 | 0.36 | 0.33 | 0.86 | 0.13 | 0.15 | 0.87 | 0.69 | 0.23 | 0.21 | 0.48 | 0.69
ResNet50-CNN | 23.5M | 0.83 | 0.001 | 0.42 | 0.00 | 0.00 | 0.00 | 0.42 | 0.00 | 0.00 | 0.00 | 1.00 | 0.00 | 0.00 | 0.00 | 0.59
DenseNet201-CNN | 51.8M | 1.67 | 0.001 | 0.90 | 1.00 | 0.85 | 0.94 | 0.85 | 0.77 | 0.90 | 0.76 | 1.00 | 0.87 | 0.87 | 0.84 | 0.95
MobileNet-CNN | 3M | 0.30 | 0.001 | 0.42 | 0.00 | 0.00 | 0.00 | 0.42 | 0.00 | 0.00 | 0.00 | 1.00 | 0.00 | 0.00 | 0.00 | 0.59
GoogleNet-CNN | 6M | 0.15 | 0.0001 | 0.77 | 0.82 | 0.66 | 0.49 | 1.00 | 0.74 | 0.58 | 0.63 | 0.97 | 0.78 | 0.62 | 0.55 | 0.98
XceptionNet-CNN | 25M | 0.31 | 0.0001 | 0.86 | 0.96 | 0.70 | 0.90 | 0.99 | 0.81 | 0.97 | 0.45 | 0.99 | 0.88 | 0.81 | 0.60 | 0.99

Precision, recall (or sensitivity), and F1-Score are crucial indicators, and
they are determined using the relations described above. For each class, the
harmonic mean of recall and precision yields the F1-score, another significant
statistical classification metric. If we obtain high F1-Score values across
all classes, this suggests that we have successfully identified samples free
of any class of brain tumor. Due to the existence of imbalance among the 4
classes, we also showed individual precision values for each class that have
been implemented using different approach. We also intend to compare several
analytical factors, such as how well each model performs with less training
samples from practicality aspect and how overfitting with lower training
samples affects performance of the classifier.

##  4 Results and Discussion

This study aimed to appropriately categorize four distinct kind of brain tumor
MRI scans. We compared 5 different methodologies—augmentation, focal loss,
SMOTE, and ADASYN — while using 8 transfer learning + CNN models. We used a
variety of evaluation criteria, including accuracy, precision, recall, and
F1-score, to evaluate the effectiveness of our suggested model. This
unbalanced dataset was categorized using deep learning models that have
already been trained as well as the suggested transfer learning + CNN model,
with 90% of the dataset being used for training and 10% being utilized for
testing. The performance of deep learning networks for classification can be
measured using various methods. In the CNN process, classification tasks are
frequently carried out using a confusion matrix. The accuracy comparison on
imbalanced dataset utilized in this research is shown in the Table 6.

Figure 9: Performance comparison using various transfer learning models + CNN
on different techniques for handling data imbalance Table 6: Accuracy
Comparison on Unbalanced Dataset using several imbalance solution approaches.
The highlighted row indicates the best performing version of our proposed
Transfer Learning-CNN approach

Model | CE | Focal Loss | Augment CE | SMOTE CE | ADASYN CE
---|---|---|---|---|---
CNN | 0.8125 | 0.80 | 0.91 | 0.82 | 0.83
VGG16-CNN | 0.83 | 0.85 | 0.96 | 0.86 | 0.87
EfficientNetB0-CNN | 0.37 | 0.42 | 0.31 | 0.47 | 0.42
EfficientNetB3-CNN | 0.60 | 0.32 | 0.32 | 0.50 | 0.52
ResNet50-CNN | 0.78 | 0.62 | 0.80 | 0.66 | 0.42
DenseNet201-CNN | 0.47 | 0.66 | 0.94 | 0.86 | 0.90
MobileNet-CNN | 0.30 | 0.42 | 0.23 | 0.42 | 0.42
GoogleNet-CNN | 0.77 | 0.78 | 0.80 | 0.79 | 0.77
XceptionNet-CNN | 0.83 | 0.81 | 0.94 | 0.83 | 0.86

The result of the above experiments showed that VGG-16 performed best overall
across all techniques for the accuracy evaluation metric. With augmentation,
0.96 was the highest accuracy ever attained overall. In this instance,
MobileNet and EfficientNets had the worst results with an overall accuracy of
0.23 and 0.31, respectively. For the relevant unbalanced dataset on VGG16,
EfficientNets, DenseNet201, MobileNet, and GoogleNet, it was found that focal
loss obtained higher accuracies than cross entropy loss when the loss
functions were compared having values 0.85, 0.66, 0.42 and 0.78, respectively.
With the cross-entropy loss, the remaining models fared better. Data
augmentation produced the overall best accuracy of the techniques used to
address the data imbalance issue.

In all classes where VGG-16 + CNN with data augmentation outperformed other
approaches, we computed the average precision, recall, and F1 scores. To
summarize, despite having far less layers than other models, transfer learning
using the VGG-16 model performed the best overall. The performance measures
employing various transfer learning + CNN models on our unbalanced dataset are
compared in Figure 9. Here, we note that performance on an augmented dataset
with cross-entropy loss outperforms alternative methods across a range of
evaluation parameters.

##  5 Conclusion and Future Work

Most computer-aided modeling techniques for the investigation of medical image
data heavily rely on the CNN model for precise and reliable medical image
classification Imam et al. [2023]. We have proposed a deep learning-based
detection and identification technique for classifying brain tumors in our
research study. With the help of brain tumor MRI imaging data, we have
examined various transfer learning models for the classification of tumor
kinds, including meningioma, glioma, and pituitary. We have added various
transfer learning techniques, followed by a CNN model for each method, to
improve the CNN model’s predictive power. According to the experimental
findings, the proposed strategy, which combines VGG-16 and CNN, has a 96%
accuracy rate which is significantly higher when compared with other
approaches.

The proposed method’s notable prediction accuracy was mostly due to data
augmentation; further contributing aspects included altering the number of
layers, using optimizers, and experimenting with various activation functions.
In the future, we will concentrate on optimizing the VGG16 and DenseNet models
through rigorous GridSearchCV and hyperparameter tuning, which requires a lot
of effort and computing power. We also want to focus on MobileNet due to its
simplicity and increased use in real-world settings.

## Acknowledgement

We sincerely thank Dr. Rao Anwer, our course advisor at Mohamed Bin Zayed
University of Artificial Intelligence, for his invaluable guidance and support
throughout this work. Additionally, we extend our special appreciation to
Mohamed Bin Zayed University of Artificial Intelligence for providing us with
the required computational resources to conduct the experiments featured in
this study. Their support has been pivotal in the successful completion of our
research.

## References

  * Abd El Kader et al. [2021] Isselmou Abd El Kader, Guizhi Xu, Zhang Shuai, Sani Saminu, Imran Javaid, and Isah Salim Ahmad.  Differential deep convolutional neural network model for brain tumor classification.  _Brain Sciences_ , 11(3):352, 2021.
  * Alam et al. [2021] Mohammed Talha Alam, Syed Ubaid, Shakil, Shahab Saquib Sohail, Maryam Nadeem, Shiraz Hussain, and Jamshed Siddiqui.  Comparative analysis of machine learning based filtering techniques using movielens dataset.  _Procedia Computer Science_ , 194:210–217, 2021.  https://doi.org/10.1016/j.procs.2021.10.075.
  * Alam et al. [2022] Mohammed Talha Alam, Shahab Saquib Sohail, Syed Ubaid, Shakil, Zafar Ali, Mohammad Hijji, Abdul Khader Jilani Saudagar, and Khan Muhammad.  Its your turn, are you ready to get vaccinated? towards an exploration of vaccine hesitancy using sentiment analysis of instagram posts.  _Mathematics_ , 10(22), 2022.  10.3390/math10224165.
  * Badža and Barjaktarović [2020] Milica M Badža and Marko Č Barjaktarović.  Classification of brain tumors from mri images using a convolutional neural network.  _Applied Sciences_ , 10(6):1999, 2020.
  * Chatterjee et al. [2022] Soumick Chatterjee, Faraz Ahmed Nizamani, Andreas Nürnberger, and Oliver Speck.  Classification of brain tumours in mr images using deep spatiospatial models.  _Scientific Reports_ , 12(1):1–11, 2022.
  * Deepak and Ameer [2019] S Deepak and PM Ameer.  Brain tumor classification using deep cnn features via transfer learning.  _Computers in biology and medicine_ , 111:103345, 2019.
  * Deepak and Ameer [2022] S Deepak and PM Ameer.  Brain tumor categorization from imbalanced mri dataset using weighted loss and deep feature fusion.  _Neurocomputing_ , 2022.
  * Imam et al. [2022] Raza Imam, Kaushal Kumar, Syed Mehran Raza, Rumi Sadaf, Faisal Anwer, Noor Fatima, Mohammad Nadeem, Mohamed Abbas, and Obaidur Rahman.  A systematic literature review of attribute based encryption in health services.  _Journal of King Saud University - Computer and Information Sciences_ , 34(9):6743–6774, 2022.  https://doi.org/10.1016/j.jksuci.2022.06.018.  URL https://www.sciencedirect.com/science/article/pii/S1319157822002269.
  * Imam et al. [2023] Raza Imam, Muhammad Huzaifa, and Mohammed El-Amine Azz.  On enhancing the robustness of vision transformers: Defensive diffusion, 2023.
  * Irmak [2021] Emrah Irmak.  Multi-classification of brain tumor mri images using deep convolutional neural network with fully optimized framework.  _Iranian Journal of Science and Technology, Transactions of Electrical Engineering_ , 45(3):1015–1036, 2021.
  * Kibriya et al. [2022] Hareem Kibriya, Rashid Amin, Asma Hassan Alshehri, Momina Masood, Sultan S Alshamrani, and Abdullah Alshehri.  A novel and effective brain tumor classification model using deep feature fusion and famous machine learning classifiers.  _Computational Intelligence and Neuroscience_ , 2022, 2022.
  * Li et al. [2022] Jian Ping Li, Shakir Khan, Mohammed Ali Alshara, Reemiah Muneer Alotaibi, CobbinahBernard Mawuli, et al.  Dacbt: deep learning approach for classification of brain tumors using mri data in iot healthcare environment.  _Scientific Reports_ , 12(1):1–14, 2022.
  * Petmezas et al. [2022] Georgios Petmezas, Grigorios-Aris Cheimariotis, Leandros Stefanopoulos, Bruno Rocha, Rui Pedro Paiva, Aggelos K Katsaggelos, and Nicos Maglaveras.  Automated lung sound classification using a hybrid cnn-lstm network and focal loss function.  _Sensors_ , 22(3):1232, 2022.
  * Raza et al. [2022] Asaf Raza, Huma Ayub, Javed Ali Khan, Ijaz Ahmad, Ahmed S. Salama, Yousef Ibrahim Daradkeh, Danish Javeed, Ateeq Ur Rehman, and Habib Hamam.  A hybrid deep learning-based approach for brain tumor classification.  _Electronics_ , 11(7):1146, 2022.
  * Veeramuthu et al. [2022] A Veeramuthu, S Meenakshi, G Mathivanan, Ketan Kotecha, Jatinderkumar R Saini, V Vijayakumar, and V Subramaniyaswamy.  Mri brain tumor image classification using a combined feature and image-based classifier.  _Frontiers in Psychology_ , 13, 2022.
  * Wang et al. [2020] Tao Wang, Changhua Lu, Mei Yang, Feng Hong, and Chun Liu.  A hybrid method for heartbeat classification via convolutional neural networks, multilayer perceptrons and focal loss.  _PeerJ Computer Science_ , 6:e324, 2020.
  * Xie et al. [2022] Yuting Xie, Fulvio Zaccagna, Leonardo Rundo, Claudia Testa, Raffaele Agati, Raffaele Lodi, David Neil Manners, and Caterina Tonon.  Convolutional neural network techniques for brain tumor classification (from 2015 to 2022): Review, challenges, and future perspectives.  _Diagnostics_ , 12(8):1850, 2022.


# Towards Foundation Models for Materials Science: The Open MatSci ML Toolkit

Kin Long Kelvin Lee<EMAIL_ADDRESS>0000-0002-1903-9242
Accelerated Graphics & Computing Systems, Intel Corporation2111 25th NE
Ave.HillsboroORUSA97124 ,  Carmelo Gonzales<EMAIL_ADDRESS>Intel
Labs, Intel Corporation2200 Mission College Blvd.Santa ClaraCAUSA95054 ,
Matthew Spellings<EMAIL_ADDRESS>Vector InstituteTorontoONCanada
,  Mikhail Galkin<EMAIL_ADDRESS>Intel Labs, Intel Corporation2200
Mission College Blvd.Santa ClaraCAUSA95054 ,  Santiago Miret
<EMAIL_ADDRESS>Intel Labs, Intel Corporation2200 Mission College
Blvd.Santa ClaraCAUSA95054 and  Nalini Kumar<EMAIL_ADDRESS>Accelerated Graphics & Computing Systems, Intel Corporation2200 Mission
College Blvd.Santa ClaraCAUSA95054

###### Abstract.

Artificial intelligence and machine learning have shown great promise in their
ability to accelerate novel materials discovery. As researchers and domain
scientists seek to unify and consolidate chemical knowledge, the case for
models with potential to generalize across different tasks within materials
science—so-called “foundation models”—grows with ambitions. This manuscript
reviews our recent progress with development of Open MatSci ML Toolkit, and
details experiments that lay the groundwork for foundation model research and
development with our framework. First, we describe and characterize a new
pretraining task that uses synthetic data generated from symmetry operations,
and reveal complex training dynamics at large scales. Using the pretrained
model, we discuss a number of use cases relevant to foundation model
development: semantic architecture of datasets, and fine-tuning for property
prediction and classification. Our key results show that for simple
applications, pretraining appears to provide worse modeling performance than
training models from random initialization. However, for more complex
instances, such as when a model is required to learn across multiple datasets
and types of targets simultaneously, the inductive bias from pretraining
provides significantly better performance. This insight will hopefully inform
subsequent efforts into creating foundation models for materials science
applications.

datasets, neural networks, materials science, chemistry, graph neural
networks, point clouds

††ccs: Applied computing Chemistry††ccs: Applied computing Computer-aided
design††ccs: Software and its engineering Designing software††ccs: Software
and its engineering Software design engineering††ccs: Computing methodologies
Artificial intelligence††ccs: Computing methodologies Machine learning

##  1\. Introduction

Materials science—despite its relative youth as a scientific discipline—has
become a dominant driver for modern day technology, with applications ranging
from semiconductors in electronics, to battery materials and light harvesting
in energy, to catalysts for faster, safer, and more efficient chemical
manufacturing. Key to all of this is the discovery of new materials: as each
of the aforementioned areas grow, the materials we use need to grow and adapt
to meet demand. Until quite recently, the design of materials has relied
exclusively on a feedback loop between experiment and theory: the former
attempts to synthesis and characterize novel materials with desirable
properties, while the latter seeks to understand underlying electronic
structure to guide design.

The shortcomings of this conventional loop, however, generally pertain to
scalability and cost. Experimental facilities are costly to start-up and run,
and are typically limited in throughput: chemical synthesis procedures can
span weeks and months, and the expertise to realize these procedures can be
hard to cultivate. For this reason, computational modeling has traditionally
complemented experimental work by enabling higher throughput screening,
whereby properties of promising materials can be calculated either from first-
principles (ab initio methods such as coupled-cluster theory (Coester and
Kümmel, 1960; Matthews et al., 2020)) and/or semi-empirical/approximate
methods such as density functional theory (DFT) (Kohn and Sham, 1965). While
in some cases, computer hardware and algorithmic improvements (Nakata et al.,
2020) have rendered computation faster than experiment, the design of space of
chemistry is so vast (Reymond, 2015) and the scaling of compute requirements
is so costly that computational modeling is still reserved only for the most
promising materials.

The advent of machine learning (ML) and artificial intelligence (AI) provides
an avenue for significantly changing—hopefully improving—the way materials can
be and are discovered. Advancements in machine-learned representations of
atoms, learning algorithms (in particular neural network architectures and
optimization), and high-performance computing jointly enable new routes to
approach materials discovery. For example, various groups (Gasteiger et al.,
2021; Chen et al., 2019; Spellings, 2022) have seen success in parameterizing
deep neural networks as surrogate models for property prediction, which once
trained, can be used for prediction at a fraction of the computational cost of
conventional electronic structure methods.

AI/ML workflows have their own requirements that need to be realized: high-
quality and diverse datasets are needed to train _useful_ models, and new
algorithms and architectures central to materials science should be developed.
To do so requires uniting high performance computing with intuitive
abstractions for materials scientists, chemists, and AI/ML practitioners and
engineers to cooperate seamlessly. These are the guiding principles behind the
development of Open MatSci ML Toolkit—a framework that seeks to unite large-
scale and disparate materials datasets, state-of-the-art model
implementations, and an end-to-end, modular pipeline that scales with use
cases from developing and testing locally on laptops to many-node distributed
computing. An outcome of this is to facilitate general purpose AI/ML models
for materials science—so-called “foundation models” (Bommasani et al.,
2022)—trained on large amounts of diverse data to produce unified
representations (Rebuffi et al., 2017) that can be tuned to specific tasks
with less data.

The present paper provides a brief introduction to these concepts, and reviews
some of the recent discoveries and challenges we have found towards this goal.
Our key contributions include:

  1. (1)

A new pretraining task based on synthetic point clouds based off sets of
symmetry operations,

  2. (2)

Demonstration of a distributed development, training, and experimentation
workflow with up to 32 nodes of 4th Generation Intel® Xeon® Scalable
Processors (codenamed Sapphire Rapids),

  3. (3)

Characterization of complex training dynamics during pretraining, which we
ascribe to recent investigations into instabilities of the Adam optimizer,

  4. (4)

Investigation of qualitative and quantitative benefits of pretraining on
downstream tasks, including under a multi-task, multi-dataset setting.

##  2\. Background and related work

###  2.1. Atomistic representations

There have been a host of methods for representing chemical structures in
machine-readable formats, which are subsequently ingested by some AI/ML
algorithm to transform the data into vector representations. Graph structures
have been the workhorse of chemists for many years (Cayley, 1874), with atoms
represented as nodes/vertices and connectivity between atoms can be
arbitrarily structured as edges. The development of neural networks for graphs
has grown rapidly, with chemical applications being at the forefront: graph
neural networks (GNN) have frequently used molecular data (Ramakrishnan et
al., 2014) as proving grounds for learning symmetry-adapted properties that
are equivariant (Satorras et al., 2022) or invariant (Schütt et al., 2017) to
3D geometry (i.e. translation, rotation, reflection).

A recent line of research extends these desirable properties to point clouds
with geometric algebra (Spellings, 2022; Brehmer et al., 2023), which work in
the space of multivector products. Some advantages of this approach over
graphs is bypassing the need to impose structure, i.e. construction of the
graph, on the data, which in turn has downstream implications such as
computational performance; graph structures can exhibit poor cache reuse
without reordering (Gong et al., 2022), in addition to the need for
specialized sparse kernels (Heinecke et al., 2016). In contrast, geometric
algebra networks rely on generally optimized dense compute in combination with
the well-studied attention mechanism, provides comparable modeling
capabilities for small graphs.

###  2.2. Related Work

####  2.2.1. Open Catalyst Project

The Open Catalyst Project (OCP) (Chanussot et al., 2021; Tran et al., 2022)
was developed in collaboration by researchers at Fundamental AI Research
(FAIR) at Meta AI and Carnegie Mellon University’s Department of Chemical
Engineering. It comprises a set of benchmark datasets and challenges
associated with catalyst activity, whereby a molecular adsorbate interacts
with a catalytic surface. At a high level, the tasks and datasets were
constructed to parameterize graph neural networks as learned potentials for
catalyst simulation. OCP has since motivated a number of related projects
including state-of-the-art models (Gasteiger et al., 2021), data-centric
transformations (Duval et al., 2023), and indeed, our own project Open MatSci
ML Toolkit. Open MatSci ML Toolkit differentiates from OCP by expanding the
scope in data representations (i.e. use of point clouds), tasks (property
prediction and classification), and datasets.

####  2.2.2. MatBench

MatBench (Dunn et al., 2020) provides a suite of evaluation benchmarks using
structures and targets provided by the related Materials Project (Jain et al.,
2013), with various tasks including property prediction and material
classification. Similar to OCP, communities are able to submit to leaderboard
run by the authors, allowing new architectures and strategies to be
demonstrated; to date, most submissions have involved GNN-based architectures.
Open MatSci ML Toolkit uses the same data source (Materials Project) as
MatBench, but differs in approach by allowing free composition of tasks and
datasets as opposed to independent benchmarks, allowing knowledge to be pooled
across areas in materials science.

####  2.2.3. MAST-ML

MAST-ML (Jacobs et al., 2020) provides a platform for materials science
research mainly with classical ML techniques in familiar notebook
environments. The advantages of this approach is a low barrier to entry for
practitioners, as well as rapid reproducible research and development by
providing interfaces to several community data repositories such as the
Materials Data Foundry (Blaiszik et al., 2016, 2019) and abstractions for
scikit-learn models (Pedregosa et al., 2011). Open MatSci ML Toolkit differs
by focusing mainly on deep learning techniques, which generally possess a
higher degree of expressivity and capacity to learn than classical ML models
like gradient boosting (Friedman, 2001). Perhaps the most striking similarity
between the two projects is the ability to interact with different datasets,
which contrasts with the other two related benchmarks.

##  3\. Open MatSci ML Toolkit

As we have mentioned in the introduction, the Open MatSci ML Toolkit was
designed to be an open source platform for AI-accelerated materials discovery,
applying concepts ranging from MLOps best practices to supporting various
representations of material structures, as described in the previous section.
We direct interested readers to an earlier introduction to the toolkit (Miret
et al., 2023) and for brevity we will focus on some newer and relevant aspects
to results we will discuss in this paper. The original concept for Open MatSci
ML Toolkit was to improve performance portability of the Open Catalyst Project
by refactoring significant amounts of functionality to PyTorch Lightning
(Falcon and The PyTorch Lightning team, 2019) and the Deep Graph Library
(DGL): the former provides pipeline abstractions that facilitate accelerator
offloading, scale-up and scale-out, and the latter offers GPU _and_ CPU-
optimized sparse kernels for efficient graph neural network compute (Wang et
al., 2019). Since the original release, we have significantly reworked the
scope of the framework to move beyond Open Catalyst tasks and datasets, and
provide a generalized abstraction to incorporate and maintain interfaces to a
variety of data sources and modeling tasks.

Figure 1. Modular workflow and abstraction for Open MatSci ML Toolkit, with
some examples provided in each block. A number of widely used materials
science _datasets_ have been incorporated, providing the foundations for
experimentation. As each data sample is retrieved, a chain of
_transformations_ can be applied to freely convert between representations,
and/or modified to introduce inductive biases. A _task_ adapts the emitted
data for a particular set of learning targets, and comprises an _encoder_ that
processes the material to generate a vector embedding that is subsequently
used by individual output heads for task predictions.

A diagram depicting how Open MatSci ML Toolkit was designed and can be used.

###  3.1. Datasets

As of writing, Open MatSci ML Toolkit supports datasets that can be used for
both property prediction and pretraining. For the former, we provide
interfaces for a diverse set of datasets, including: the Open Catalyst Project
(Chanussot et al., 2021; Tran et al., 2022); the Materials Project (Jain et
al., 2013); the Carolina Materials Database (Zhao et al., 2021; Dan et al.,
2020); and the LiPS dataset (Batzner et al., 2022). These datasets range from
material properties such as band gap and formation energy, as well as time-
dependent dynamics with energy/force labels for trajectory samples.

For pretraining, we include a pipeline for generating synthetic point clouds
based off symmetry groups relevant to crystal structures: at a high-level,
this involves creating a collection of particles in 3D space, followed by
repetition based on randomly chosen symmetry operations. While these data do
not necessarily reflect real material structures nor contain chemical
information, they share the same symmetry basis, and by learning to identify
point groups, we believe should instill some structural concepts in
embeddings. Further, as the positions of particles are randomly generated, we
are able to sample configurations at arbitrary scales and uniformly across
point groups; this is especially important as materials datasets typically
focus on particular types of structures and symmetries, thereby inducing a
selection bias.

From Figure 1, the transformation interface and pipeline allows datasets to be
freely converted between representations, as well as include additional
variables and manipulations as needed by downstream modeling.

###  3.2. Tasks

Given a dataset and some neural network model, we define a “task” as a
learning objective such as property regression (e.g. the semiconductor band
gap), or material classification (e.g. is the material stable at 298 K).
Concretely, a task comprises an encoding model which ingests a graph or point
cloud structure and emits an embedding that is subsequently used by one or
several output heads (see Appendix A) to independently predict targets, in
which the encoder and each output head is trained jointly through an
optimizer. Each task is implemented as Python classes and can be freely paired
with datasets in a modular way (Fig. 1); for instance in this work, the
Materials Project dataset is used for both regression and classification
tasks.

For more complex learning required for foundation models, Open MatSci ML
Toolkit is able to seemlessly compose multiple datasets and tasks together,
using an encoder model that learns generalized representations across them. On
the one hand, multi-task learning makes full use of all labels present in a
dataset, for example performing regression and classification jointly. On the
other hand, using multiple datasets to perform a similar task creates the
opportunity for models to specialize by scaling up available training data.
Combining both multiple datasets and tasks in theory facilitates highly
general models, providing the underlying encoder has the capacity and
expressivity. In this case, the joint encoder is updated separately to each
task output head.

##  4\. Experimental setup

###  4.1. Platform and software

All experiments were carried out on the Endeavour cluster at Intel. Each
compute node consists of two CPU sockets—both populated with Intel® Xeon®
Platinum 8480+ processors comprising 56 physical cores partitioned into four
non-uniform memory access (NUMA) nodes with combined 256 GB of DDR5 (4800
MT/s) memory.111Using Rocky Linux 8.8, kernel version 4.18 Each compute node
is connected via Mellanox HDR200 interconnect, and datasets stored on a Lustre
filesystem. Scale-out experiments for measured throughput were conducted on
8/2/2023.

We used an internal build of Open MatSci ML Toolkit which we are preparing for
open source release. The main dependencies include DGL (`0.9.1`), PyTorch
(`1.13.1`), and PyTorch Lightning (`1.8.6`). PyTorch was obtained as part of
the Intel® AI Kit `2023.1` via the `intel` `conda` channel, which is prebuilt
against optimized Intel® oneAPI libraries including oneMKL and oneDNN, along
with framework optimized libraries such as oneCCL Bindings for PyTorch. For
intraop parallelism, we used Intel® OpenMP threading, and for distributed
training and inference, Intel® MPI. In all instances of training and
inference, processes were bound to their respective NUMA domains (`-map-by
numa`) with threads pinned to physical cores (`I_MPI_PIN_CELL=core`).

###  4.2. Training strategy

To focus the scope of work, we only consider a single GNN architecture, the
equivariant graph neural network [E(n)-GNN] developed by Satorras et al.
(2022), with a nominal, untuned set of hyperparameters. Details regarding
neural network architectures can be found in Appendix A.

In all cases, we use distributed data parallelism (DDP) to accelerate
training: the full dataset is divided across MPI processes ($N$) receiving the
same number of samples ($B$) per batch to yield an effective batch size
$B_{\mathrm{eff}}=NB$. For a single node, the number of parallel workers—16
for our machine configuration—was chosen to balance compute per worker (i.e.
the available FLOP/s determined by `OMP_NUM_THREADS`) and memory bandwidth per
socket, and interprocess communication only occurs at the gradient averaging
step. In this work, we report scaling up to 32 nodes, yielding $N=512$ and
$B_{\mathrm{eff}}=16,384$.

For all training, we use the Adam optimizer with decoupled weight decay
(`AdamW`) (Loshchilov and Hutter, 2019), using default momentum values
($\beta_{1}=0.9$, $\beta_{2}=0.999$). For DDP training, the base learning rate
($\eta_{\mathrm{base}}$) is scaled by the total number of DDP workers to
ensure constant gradient variance (Goyal et al., 2018). Over the course of
training, the learning rate $\eta$ is modified based on learning rate
schedulers: a warmup period of eight epochs linearly ramps $\eta$ to its
nominal value, followed by an exponential decay with $\gamma=0.8$. For fine-
tuning, we scale $\eta_{\mathrm{base}}$ down by a factor of ten, which is
typical practice to mitigate forgetting (Li and Hoiem, 2017).

##  5\. Results & discussion

###  5.1. Training scale-out

To facilitate high-throughput experimentation, we relied on both scaling-up
and scaling-out training and inference to many MPI processes. Figure 2 shows
the combined training throughput for the symmetry pretraining task as a
function of DDP workers, as we scale out from one to a total of 32 nodes. We
see a linear growth in training throughput, implying that the communication
overhead associated with gradient synchronization for every training step is
negligible compared to the compute of individual workers. For additional
context, the right ordinate of Figure 2 converts each tick for throughput into
time per epoch; in most configurations considered, the full dataset is
traversed within minutes, enabling quick feedback for experimental
configurations.

Figure 2. Pretraining throughput scaling as a function of the number of DDP
workers, shown as both samples per second, and based on a dataset size of
2,000,000 samples, the time to complete an epoch on the right ordinate. The
dashed line represents a linear fit to the measurements.

###  5.2. Symmetry group pretraining

To ensure quality of down stream modeling, it is important to characterize the
convergence dynamics of pretraining, as well as the time to solution in a
wall-time basis: higher training throughput can be achieved by scaling out to
a larger number of DDP workers, however as the gradients are averaged over
large $B_{\mathrm{eff}}$, fewer optimizer updates are performed with less
information content. To partially compensate, we follow the recommendation of
Goyal et al. (2018) to scale the learning rate for a single batch proportional
to the number of DDP workers (i.e. the world size).

Figure 3 shows the validation error as a function of the number of DDP workers
for a fixed number of training steps. The top plot shows that, with
$\eta_{\mathrm{base}}=10^{-3}$—a typical value for the learning rate—learning
stagnates early for all scale-out configurations at large validation errors,
indicating a combination of fairly broad local minima, and insufficiently
informative gradients to escape them. An intuitive solution is to
substantially lower the base learning rate, shown in the bottom plot of Figure
3. Here, while we see the validation error does decrease over time, we draw
two general observations: first, the single node (16 DDP ranks) converges
albeit slowly; second, the early rate of convergence seems to increase with
the number of workers, however the prevalence of spikes in the validation loss
also increases. In the largest scale-out case (512 ranks), the model
experiences a large spike around step ${\sim}$550 and does not recover to
former levels.

Recently, Molybog et al. (2023) highlighted previously reported instabilities
in the training dynamics of large language models (Chowdhery et al., 2022),
and provided a thorough investigation into the convergence properties of Adam-
based optimizers. The authors observe similar optimizer divergence, correlated
with large gradient norms and attribute these occurrences to inherent
instabilities in the Adam algorithm. In particular, one of assumptions
required for convergent properties of Adam is zero time-correlation between
update steps, i.e. Markovian dynamics, which are violated for large batch
training, and result in a disconnect in the learning dynamics between layers
whereby gradients decay to the order of $\varepsilon$ used to prevent division
by zero. While Molybog et al. (2023) discuss this in the context of large
language models, the instabilities are general to “Adam-like” optimizers and
likely to apply here, particularly as the use of GNN encoders can lead to
over-smoothing (see (Rusch et al., 2023) for a review) and therefore decrease
signal contained within gradients; a symptom described by Molybog et al.
(2023) in early layers of language models.

A thorough analysis into this phenomenon is outside our scope, and in the
interim, we found a balance between convergence characteristics and training
throughput with $N=256$, or 16 nodes. While this kind of optimization is
conventionally offloaded to hyperparameter optimization, further work is
required to assess a more principled approach for large scale pretraining in
Open MatSci ML Toolkit. Subsequent experiments involve using a model trained
for 20 epochs, taking approximately ${\sim}80$ min.—an adequate turnover rate
for experimentation.

Figure 3. Early training dynamics for the pretraining task as a function of
the number of DDP workers, which increases the effective batch size
$B_{\mathrm{eff}}$ proportionally. The top and bottom frames correspond to
higher and lower base learning rates. In both instances, the
ordinate—representing the cross-entropy error—has been truncated to show
scales relevant to the discussion.

###  5.3. Dataset exploration with pretrained models

Despite ongoing challenges with training, in this section we show that the
model with learned representations of symmetry groups still possesses valuable
inductive bias. One application of pretraining is the potential to architect
datasets semantically based on structure, chemical composition, or both. Both
elements reflect in the material design process, and thus the corresponding
datasets used to train foundation models for materials science should include
this aspect for composition. Here, we use the pretrained E(n)-GNN model from
the previous section to embed samples contained across the datasets supported
in Open MatSci ML Toolkit, and using UMAP (McInnes et al., 2018) as a
structure-preserving projection method, visualize the distribution and
diversity of material structures holistically.

The result is shown in Figure 4, from which we can draw three qualitative
observations: first, datasets tend to comprise the same kinds of structural
motifs, which may be a reflection of inherent selection bias when it comes to
composing datasets; second, somewhat expectedly, there is significant overlap
between the OCP datasets; third, Materials Project data share some structural
motifs with OCP, but perhaps offers the broadest variety of structures. For
context, these observation can be calibrated by the fact that the LiPS
dataset, which comprises trajectories of the same material composition, forms
a clear independent cluster. The conclusions we can draw from these
observations align well with our design philosophy: the inclusion of a variety
of datasets provides the diversity of materials needed to complement one
another, and this holistic view can provide valuable insight when considering
what data to include, and how much is needed. The same analysis could be done
using an encoder trained with chemical information, for example Materials
Project, to find dataset gaps in _chemical space_.

Figure 4. UMAP visualization of each dataset implemented in Open MatSci ML
Toolkit. Each point represents a structure contained within its dataset,
embedded with an E(n)-GNN encoder pretrained on synthetic constructs. For each
dataset, 10,000 structures were randomly chosen as a representative sample.
The UMAP parameters used were: 200 n_neighbors, 0.05 min_dist, using euclidean
distance as the metric.

UMAP visualization of datasets implemented in Open MatSci ML Toolkit—this
illustrates the space of structures contained.

###  5.4. Fine-tuning on material properties

Table 1. Comparison of task metrics for multi-task, multi-data training starting from a pretrained model, and from random initialization. With the exception of stability, which corresponds to the binary cross-entropy error, regression errors are provided as mean-absolute error in their respective units. Values in boldface correspond to the lowest error. $\zeta$ corresponds to the Fermi energy. | Materials Project | Carolina Materials Database
---|---|---
Configuration | Band gap (eV) |  $\zeta$ (eV) |  $E_{\mathrm{form}}$ (eV/atom) | Stability classification |  $E_{\mathrm{form}}$ (eV/atom)
Pretrained | 1.27 | 0.76 | 0.83 | 0.42 | 0.14
From scratch | 4.80 | 3.86 | 3.54 | 0.40 | 0.10

With a pretrained model, this section assesses whether there are quantiative
benefits to pretraining for tasks relevant to materials discovery. We compare
the validation performance of models that have been pretrained, versus from
random initialization or “from scratch” for a given task.

The simplest case, shown in Figure 5, corresponds to a band gap regression
with the Materials Project Dataset. For both initializations, we observe
smooth, convergent behavior. In the case of the pretrained model, we see that
the early training stages converge to a lower error more quickly than the
randomly initialized model, however appears to fall into a local minimum. In
contrast, the model trained from scratch sees slower convergence, but
eventually to a better performing model by the end of training. From this
result, it appears that the pretraining may introduce a strong bias that
allows for faster convergence that may see benefits with early stopping
algorithms with a fixed compute budget, but for longer regimes, training from
random initialization may provide sufficient momentum to escape local minima.

Figure 5. Validation curves for E(n)-GNN models trained to perform band gap
prediction with the Materials Project dataset, comparing trajectories starting
from a model pretrained on synthetic data (red) and from random initialization
(gray blue).

In a more complex procedure, and one that resembles more of foundation model
training, we perform the same experiment but in a multi-task, multi-dataset
setting. The joint task comprises band gap, Fermi energy ($\zeta$), and
formation energy ($E_{\mathrm{form}}$) regression along with stability
classification for the Materials Project, in addition to formation energy
prediction for the Carolina Materials Database. The results shown in Table 1
provide a different perspective on pretraining from the simple task: when
trained jointly against multiple types of targets, the general performance for
the pretrained model is significantly better than training from scratch, with
lower errors for three of the five metrics, with the two remaining being
relatively comparable in magnitude. While a more in-depth investigation into
the learning dynamics is likely needed, this result does indicate that there
is merit in pretraining for foundation models in materials science.

##  6\. Conclusions and Future Work

In this work, we have provided a brief introduction into the capabilities of
Open MatSci ML Toolkit, focusing on areas relevant to foundation models for
materials science. Throughout, we have demonstrated a new pretraining pipeline
at scale, using up to 32 compute nodes for fast experimentation turnover. As
part of our pretraining experiments, we reported significant training
instabilities correlating with large batch training, which we discuss in the
context of work by (Molybog et al., 2023) underscoring shortcomings in the
Adam algorithm. Our results indicate that, when considering pretraining,
additional attention may be required balancing training throughput and model
convergence when using Adam-like optimizers, and more work is needed into
finding optimal learning schedules.

With the pretrained model, we briefly discuss a proof-of-concept application
for semantic dataset design. Using a UMAP visualization based on embedded
structures across datasets supported in Open MatSci ML Toolkit, we highlight
gaps in data space that could be filled to improve the uniformity of
datasets—in both chemical and structural spaces—used to train foundation
models. Finally, we investigate whether there are tangible benefits with
pretraining for downstream tasks relevant to materials science. We find that
in the simpler case of single target regression, pretraining can improve early
convergence rates but does not result in a more general model. However, for
more complex instances such as multi-task, multi-dataset cases relevant to
foundation model training, pretraining results in a far more generalized model
across targets and datasets than training from random initialization. We hope
our early contributions should hopefully guide future experimentation and
development of foundation models for materials science, both in terms of data
composition and training strategies.

## Notices & disclaimers

Performance varies by use, configuration and other factors. Learn more on the
Performance Index site. Performance results are based on testing as of dates
shown in configurations and may not reflect all publicly available updates.
See backup for configuration details. No product or component can be
absolutely secure. Your costs and results may vary. Intel technologies may
require enabled hardware, software or service activation. © Intel Corporation.
Intel, the Intel logo, and other Intel marks are trademarks of Intel
Corporation or its subsidiaries. Other names and brands may be claimed as the
property of others.

## References

  * (1)
  * Batzner et al. (2022) Simon Batzner, Albert Musaelian, Lixin Sun, Mario Geiger, Jonathan P. Mailoa, Mordechai Kornbluth, Nicola Molinari, Tess E. Smidt, and Boris Kozinsky. 2022.  E(3)-Equivariant Graph Neural Networks for Data-Efficient and Accurate Interatomic Potentials.  https://doi.org/10.24435/MATERIALSCLOUD:S0-5N
  * Blaiszik et al. (2016) B. Blaiszik, K. Chard, J. Pruyne, R. Ananthakrishnan, S. Tuecke, and I. Foster. 2016\.  The Materials Data Facility: Data Services to Advance Materials Science Research.  _JOM_ 68, 8 (Aug. 2016), 2045–2052.  https://doi.org/10.1007/s11837-016-2001-3
  * Blaiszik et al. (2019) Ben Blaiszik, Logan Ward, Marcus Schwarting, Jonathon Gaff, Ryan Chard, Daniel Pike, Kyle Chard, and Ian Foster. 2019\.  A Data Ecosystem to Support Machine Learning in Materials Science.  _MRS Communications_ 9, 4 (Dec. 2019), 1125–1133.  https://doi.org/10.1557/mrc.2019.118
  * Bommasani et al. (2022) 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 Nie, 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. 2022.  On the Opportunities and Risks of Foundation Models.  arXiv:2108.07258 [cs]
  * Brehmer et al. (2023) Johann Brehmer, Pim de Haan, Sönke Behrends, and Taco Cohen. 2023.  Geometric Algebra Transformers.  arXiv:2305.18415 [cs, stat]
  * Cayley (1874) Cayley. 1874.  LVII. _On_ _the Mathematical Theory of Isomers_.  _The London, Edinburgh, and Dublin Philosophical Magazine and Journal of Science_ 47, 314 (June 1874), 444–447.  https://doi.org/10.1080/14786447408641058
  * Chanussot et al. (2021) Lowik Chanussot, Abhishek Das, Siddharth Goyal, Thibaut Lavril, Muhammed Shuaibi, Morgane 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. 2021.  The Open Catalyst 2020 (OC20) Dataset and Community Challenges.  _ACS Catalysis_ 11, 10 (May 2021), 6059–6072.  https://doi.org/10.1021/acscatal.0c04525 arXiv:2010.09990 [cond-mat]
  * Chen et al. (2019) Chi Chen, Weike Ye, Yunxing Zuo, Chen Zheng, and Shyue Ping Ong. 2019.  Graph Networks as a Universal Machine Learning Framework for Molecules and Crystals.  _Chemistry of Materials_ 31, 9 (May 2019), 3564–3572.  https://doi.org/10.1021/acs.chemmater.9b01294
  * Chowdhery et al. (2022) Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, Parker Schuh, Kensen Shi, Sasha Tsvyashchenko, Joshua Maynez, Abhishek Rao, Parker Barnes, Yi Tay, Noam Shazeer, Vinodkumar Prabhakaran, Emily Reif, Nan Du, Ben Hutchinson, Reiner Pope, James Bradbury, Jacob Austin, Michael Isard, Guy Gur-Ari, Pengcheng Yin, Toju Duke, Anselm Levskaya, Sanjay Ghemawat, Sunipa Dev, Henryk Michalewski, Xavier Garcia, Vedant Misra, Kevin Robinson, Liam Fedus, Denny Zhou, Daphne Ippolito, David Luan, Hyeontaek Lim, Barret Zoph, Alexander Spiridonov, Ryan Sepassi, David Dohan, Shivani Agrawal, Mark Omernick, Andrew M. Dai, Thanumalayan Sankaranarayana Pillai, Marie Pellat, Aitor Lewkowycz, Erica Moreira, Rewon Child, Oleksandr Polozov, Katherine Lee, Zongwei Zhou, Xuezhi Wang, Brennan Saeta, Mark Diaz, Orhan Firat, Michele Catasta, Jason Wei, Kathy Meier-Hellstern, Douglas Eck, Jeff Dean, Slav Petrov, and Noah Fiedel. 2022.  PaLM: Scaling Language Modeling with Pathways.  https://doi.org/10.48550/arXiv.2204.02311 arXiv:2204.02311 [cs]
  * Coester and Kümmel (1960) F. Coester and H. Kümmel. 1960.  Short-Range Correlations in Nuclear Wave Functions.  _Nuclear Physics_ 17 (June 1960), 477–485.  https://doi.org/10.1016/0029-5582(60)90140-1
  * Dan et al. (2020) Yabo Dan, Yong Zhao, Xiang Li, Shaobo Li, Ming Hu, and Jianjun Hu. 2020\.  Generative Adversarial Networks (GAN) Based Efficient Sampling of Chemical Composition Space for Inverse Design of Inorganic Materials.  _npj Computational Materials_ 6, 1 (June 2020), 1–7.  https://doi.org/10.1038/s41524-020-00352-0
  * Dunn et al. (2020) Alexander Dunn, Qi Wang, Alex Ganose, Daniel Dopp, and Anubhav Jain. 2020.  Benchmarking Materials Property Prediction Methods: The Matbench Test Set and Automatminer Reference Algorithm.  _npj Computational Materials_ 6, 1 (Sept. 2020), 1–10.  https://doi.org/10.1038/s41524-020-00406-3
  * Duval et al. (2023) Alexandre Duval, Victor Schmidt, Santiago Miret, Yoshua Bengio, Alex Hernández-García, and David Rolnick. 2023.  PhAST: Physics-Aware, Scalable, and Task-specific GNNs for Accelerated Catalyst Design.  https://doi.org/10.48550/arXiv.2211.12020 arXiv:2211.12020 [physics]
  * Falcon and The PyTorch Lightning team (2019) William Falcon and The PyTorch Lightning team. 2019.  _PyTorch Lightning_.  https://doi.org/10.5281/zenodo.3828935
  * Friedman (2001) Jerome H. Friedman. 2001\.  Greedy Function Approximation: A Gradient Boosting Machine.  _The Annals of Statistics_ 29, 5 (Oct. 2001), 1189–1232.  https://doi.org/10.1214/aos/1013203451
  * Gasteiger et al. (2021) Johannes Gasteiger, Florian Becker, and Stephan Günnemann. 2021.  GemNet: Universal Directional Graph Neural Networks for Molecules. In _Advances in Neural Information Processing Systems_ , Vol. 34. Curran Associates, Inc., 6790–6802.
  * Gong et al. (2022) Zhangxiaowen Gong, Houxiang Ji, Yao Yao, Christopher W. Fletcher, Christopher J. Hughes, and Josep Torrellas. 2022\.  Graphite: Optimizing Graph Neural Networks on CPUs through Cooperative Software-Hardware Techniques. In _Proceedings of the 49th Annual International Symposium on Computer Architecture_. ACM, New York New York, 916–931.  https://doi.org/10.1145/3470496.3527403
  * Goyal et al. (2018) Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He. 2018.  Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour.  arXiv:1706.02677 [cs]
  * Heinecke et al. (2016) Alexander Heinecke, Greg Henry, Maxwell Hutchinson, and Hans Pabst. 2016.  LIBXSMM: Accelerating Small Matrix Multiplications by Runtime Code Generation. In _Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis_ _(SC ’16)_. IEEE Press, Salt Lake City, Utah, 1–11.
  * Jacobs et al. (2020) Ryan Jacobs, Tam Mayeshiba, Ben Afflerbach, Luke Miles, Max Williams, Matthew Turner, Raphael Finkel, and Dane Morgan. 2020\.  The Materials Simulation Toolkit for Machine Learning (MAST-ML): An Automated Open Source Toolkit to Accelerate Data-Driven Materials Research.  _Computational Materials Science_ 176 (April 2020), 109544\.  https://doi.org/10.1016/j.commatsci.2020.109544
  * Jain et al. (2013) Anubhav Jain, Shyue Ping Ong, Geoffroy Hautier, Wei Chen, William Davidson Richards, Stephen Dacek, Shreyas Cholia, Dan Gunter, David Skinner, Gerbrand Ceder, and Kristin A. Persson. 2013.  Commentary: The Materials Project: A Materials Genome Approach to Accelerating Materials Innovation.  _APL Materials_ 1, 1 (July 2013), 011002\.  https://doi.org/10.1063/1.4812323
  * Klambauer et al. (2017) Günter Klambauer, Thomas Unterthiner, Andreas Mayr, and Sepp Hochreiter. 2017.  Self-Normalizing Neural Networks.  arXiv:1706.02515 [cs, stat]
  * Kohn and Sham (1965) W. Kohn and L. J. Sham. 1965.  Self-Consistent Equations Including Exchange and Correlation Effects.  _Physical Review_ 140, 4A (Nov. 1965), A1133–A1138.  https://doi.org/10.1103/PhysRev.140.A1133
  * Li and Hoiem (2017) Zhizhong Li and Derek Hoiem. 2017.  Learning without Forgetting.  arXiv:1606.09282 [cs, stat]
  * Loshchilov and Hutter (2019) Ilya Loshchilov and Frank Hutter. 2019.  Decoupled Weight Decay Regularization.  https://doi.org/10.48550/arXiv.1711.05101 arXiv:1711.05101 [cs, math]
  * Matthews et al. (2020) Devin A. Matthews, Lan Cheng, Michael E. Harding, Filippo Lipparini, Stella Stopkowicz, Thomas-C. Jagau, Péter G. Szalay, Jürgen Gauss, and John F. Stanton. 2020\.  Coupled-Cluster Techniques for Computational Chemistry: The CFOUR Program Package.  _The Journal of Chemical Physics_ 152, 21 (June 2020), 214108.  https://doi.org/10.1063/5.0004837
  * McInnes et al. (2018) Leland McInnes, John Healy, Nathaniel Saul, and Lukas Großberger. 2018.  UMAP: Uniform Manifold Approximation and Projection.  _Journal of Open Source Software_ 3, 29 (Sept. 2018), 861.  https://doi.org/10.21105/joss.00861
  * Miret et al. (2023) Santiago Miret, Kin Long Kelvin Lee, Carmelo Gonzales, Marcel Nassar, and Matthew Spellings. 2023\.  The Open MatSci ML Toolkit: A Flexible Framework for Machine Learning in Materials Science.  _Transactions on Machine Learning Research_ (2023).
  * Molybog et al. (2023) Igor Molybog, Peter Albert, Moya Chen, Zachary DeVito, David Esiobu, Naman Goyal, Punit Singh Koura, Sharan Narang, Andrew Poulton, Ruan Silva, Binh Tang, Diana Liskovich, Puxin Xu, Yuchen Zhang, Melanie Kambadur, Stephen Roller, and Susan Zhang. 2023.  A Theory on Adam Instability in Large-Scale Machine Learning.  https://doi.org/10.48550/arXiv.2304.09871 arXiv:2304.09871 [cs, math]
  * Nakata et al. (2020) Ayako Nakata, Jack S. Baker, Shereif Y. Mujahed, Jack T. L. Poulton, Sergiu Arapan, Jianbo Lin, Zamaan Raza, Sushma Yadav, Lionel Truflandier, Tsuyoshi Miyazaki, and David R. Bowler. 2020.  Large Scale and Linear Scaling DFT with the CONQUEST Code.  _The Journal of Chemical Physics_ 152, 16 (April 2020), 164112.  https://doi.org/10.1063/5.0005074
  * Pedregosa et al. (2011) F. Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B. Thirion, O. Grisel, M. Blondel, P. Prettenhofer, R. Weiss, V. Dubourg, J. Vanderplas, A. Passos, D. Cournapeau, M. Brucher, M. Perrot, and E. Duchesnay. 2011.  Scikit-Learn: Machine Learning in Python.  _Journal of Machine Learning Research_ 12 (2011), 2825–2830.
  * Ramakrishnan et al. (2014) Raghunathan Ramakrishnan, Pavlo O. Dral, Matthias Rupp, and O. Anatole von Lilienfeld. 2014\.  Quantum Chemistry Structures and Properties of 134 Kilo Molecules.  _Scientific Data_ 1, 1 (Aug. 2014), 140022\.  https://doi.org/10.1038/sdata.2014.22
  * Rebuffi et al. (2017) Sylvestre-Alvise Rebuffi, Hakan Bilen, and Andrea Vedaldi. 2017\.  Learning Multiple Visual Domains with Residual Adapters. In _Advances in Neural Information Processing Systems_ , Vol. 30. Curran Associates, Inc.
  * Reymond (2015) Jean-Louis Reymond. 2015\.  The Chemical Space Project.  _Accounts of Chemical Research_ 48, 3 (March 2015), 722–730.  https://doi.org/10.1021/ar500432k
  * Rusch et al. (2023) T. Konstantin Rusch, Michael M. Bronstein, and Siddhartha Mishra. 2023.  A Survey on Oversmoothing in Graph Neural Networks.  arXiv:2303.10993 [cs]
  * Satorras et al. (2022) Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. 2022\.  E(n) Equivariant Graph Neural Networks.  arXiv:2102.09844 [cs, stat]
  * Schütt et al. (2017) Kristof T. Schütt, Pieter-Jan Kindermans, Huziel E. Sauceda, Stefan Chmiela, Alexandre Tkatchenko, and Klaus-Robert Müller. 2017\.  SchNet: A Continuous-Filter Convolutional Neural Network for Modeling Quantum Interactions.  _arXiv:1706.08566 [physics, stat]_ (Dec. 2017).  arXiv:1706.08566 [physics, stat]
  * Spellings (2022) Matthew Spellings. 2022\.  Geometric Algebra Attention Networks for Small Point Clouds.  arXiv:2110.02393 [cs]
  * Tran et al. (2022) Richard Tran, Janice Lan, Muhammed Shuaibi, Brandon M. Wood, Siddharth Goyal, Abhishek Das, Javier Heras-Domingo, Adeesh Kolluru, Ammar Rizvi, Nima Shoghi, Anuroop Sriram, Felix Therrien, Jehad Abed, Oleksandr Voznyy, Edward H. Sargent, Zachary Ulissi, and C. Lawrence Zitnick. 2022.  The Open Catalyst 2022 (OC22) Dataset and Challenges for Oxide Electrocatalysts.  arXiv:2206.08917 [cond-mat, physics:physics]
  * Wang et al. (2019) Minjie Wang, Da Zheng, Zihao Ye, Quan Gan, Mufei Li, Xiang Song, Jinjing Zhou, Chao Ma, Lingfan Yu, Yu Gai, Tianjun Xiao, Tong He, George Karypis, Jinyang Li, and Zheng Zhang. 2019.  Deep Graph Library: A Graph-Centric, Highly-Performant Package for Graph Neural Networks.  _arXiv preprint arXiv:1909.01315_ (2019).  arXiv:1909.01315
  * Zhang and Sennrich (2019) Biao Zhang and Rico Sennrich. 2019.  Root Mean Square Layer Normalization.  https://doi.org/10.48550/arXiv.1910.07467 arXiv:1910.07467 [cs, stat]
  * Zhao et al. (2021) Yong Zhao, Mohammed Al-Fahdi, Ming Hu, Edirisuriya M. D. Siriwardane, Yuqi Song, Alireza Nasiri, and Jianjun Hu. 2021.  High-Throughput Discovery of Novel Cubic Crystal Materials Using Deep Generative Neural Networks.  _Advanced Science_ 8, 20 (2021), 2100566.  https://doi.org/10.1002/advs.202100566

##  Appendix A Model descriptions

The encoder backbone used throughout this work is the E(n)-GNN architecture
(Satorras et al., 2022; Batzner et al., 2022), which is designed to model
equivariance in $n$-dimensions; for chemical applications, equivariance in
three-dimensions allows parameterizing functions that commute in euclidean
space. Within this framework, model outputs are equivalent for a given set of
translation, rotation, and reflection operations, regardless of their sequence
ordering. We defer interested readers to Satorras et al. (2022), and only
briefly explain the architecture at a high level. Given a graph $\mathcal{G}$
comprising nodes (atoms) $\mathcal{V}$ and edges (bonds) $\mathcal{E}$, the
Equivariant Graph Convolutional Layer (Satorras et al., 2022) transforms node
$\mathbf{h}$ and coordinate $\mathbf{x}$ embeddings through message function
$\mathbf{m}$ between two nodes $i,j$:

(1)
$\mathbf{m}_{ij}=\phi_{e}(\mathbf{h}_{i},\mathbf{h}_{j},||x_{i}-x_{j}||^{2},a_{ij})$

The coordinate embeddings for node $i$ in a subsequent layer ($l+1$) are then
given by:

(2) $\mathbf{x}^{l+1}_{i}=x^{l}_{i}+C\sum_{j\neq
i}(x^{l}_{i}-x^{l}_{j})\phi_{x}(\mathbf{m}_{i}j)$

where $C$ normalizes the sum aggregation by the dimensionality of the
embeddings. E(n)-GNN differentiates from conventional graph convolution by
accounting for $x^{l}_{ij}$ in node and coordinate embeddings given by the
last two equations—by taking capturing absolute and relative distances between
nodes, thereby preserving translation and rotational equivariance. For our
work, the input features comprise atom embeddings from learnable embedding
tables, and a vector of $x,y,z$ coordinates, which are passed sequentially
through three E(n)-GNN layers, yielding equivariant representations based off
three-hop neighborhoods of nodes. Each step combines modeling capabilities
from three fully-connected layers with respective nonlinearities to transform
each type of embedding: globally, we use `SiLU` activations; for node and
message layers, a width of 256 neurons; for positional embeddings, widths of
64. Residual connections are included between graph layers. Finally, at a
graph-level, we chose to perform size-extensive modeling, whereby the graph
embedding is obtained by summation over the individual node embeddings.

The system-level embeddings are subsequently passed into blocks of multilayer
perceptrons (MLP) to transform graph (or point cloud) embeddings into each
individual target within a prescribed task. Their design was to strike a
balance in the capacity of each output head as to be expressive enough to map
onto their targets, while constrained to not ignore the embedding completely.
Each output head comprises a sequence of MLP blocks with residual connection,
which in turn correspond to MLP→non-linearity→normalization→dropout. The
output of the block is subsequently added with the inputs. In the present
work, MLPs typically have a hidden dimension of 256, using `SELU` activation
(Klambauer et al., 2017) with `RMSNorm` (Zhang and Sennrich, 2019), and a
dropout probability of 0.2. We did not perform extensive hyperparameter tuning
throughout this work; the activation was chosen for slightly better
performance over `SiLU`, and `RMSNorm` provided more reliable behavior when
operating in the multi-task, multi-data setting where irregular batches may
cause issues with traditional `BatchNorm`. For pretraining and fine-tuning in
the single task setting, three output blocks were used. For the multi-task,
multi-dataset case, each output head comprised six output blocks.

##  Appendix B Training curves

Figure 6. Learning curves for the pretrained E(n)-GNN model used for
subsequent fine-tuning experiments, measured by the multiclass cross-entropy
error. The yellow dashed line corresponds to the monitored learning rate
$\eta$ plot on a logarithmic scale, based on $\eta_{\mathrm{base}}=10^{-5}$
and $N=512$. One epoch comprises approximately 500 training steps. Figure 7.
Validation (lower is better) curves for each metric for the multi-task, multi-
dataset experiments comparing against models pretrained and from random
initialization. For regression targets, units are provided, while “stability”
is the binary cross-entropy error. For the formation energy,
$E_{\mathrm{form}}$, we differentiate between data from the Carolina Materials
Database (CMDataset) and the Materials Project. The $E_{\mathrm{form}}$ panel
for Carolina Materials Database truncates the ordinate, where as the loss
spikes to abnormal levels before recovering.

Figure 6 shows the training curve for the final pretrained E(n)-GNN model used
for downstream tasks. In the early stages of training, we observe the training
spikes described in Section 5.2, which we attributed to the combination of
large effective batch sizes, and instabilities owing to the Adam optimizer.
The learning rate curve shows the initial linear ramp up to
$\eta_{\mathrm{base}}N$ over five epochs, followed by exponential decay. As
the learning rate is decreased, the optimizer stabilizes albeit learning
gradually plateaus.

Figure 7 presents the validation error across training, where the final
validation errors are shown in Table 1. We see that for the three cases where
pretraining significantly outperforms training from scratch, the model
generally struggles to learn throughout training but the pretraining model
contains sufficient inductive bias to provide a better baseline.