Is Self-Supervised Pretraining Good for Extrapolation in Molecular Property Prediction?

Extrapolation generally means predicting data sampled from unseen distributions or domains. The general definition of extrapolation is described in the following way according to [22]. Firstly, this paper lets $\mathcal{X}$ be the domain of interests such as the structural data of materials. A target function $g\colon\mathcal{X}\to\mathbb{R}$ generates labels from the inputs. A label $y\in\mathbb{R}$ and an input $\boldsymbol{x}\in\mathcal{X}$ have a relationship of $y=g(\boldsymbol{x})$. Given that $\mathcal{D}$ is the support of certain training distribution, training data are expressed as $\{(\boldsymbol{x}_{i},y_{i})\}^{n}_{i=1}\subset\mathcal{D}$. Note that as an input $\boldsymbol{x}\in\mathcal{X}$ can be an arbitrary format such as a graph, various definitions can be considered for the data region and distribution. Extrapolation aims to learn the distribution $\mathcal{X}\setminus\mathcal{D}$ by minimizing the extrapolation error $\mathbb{E}_{\boldsymbol{x}\sim\mathcal{X}\setminus\mathcal{D}}\left[\ell\left(f(\boldsymbol{x}),g(\boldsymbol{x})\right)\right]$. Note that $f\colon\mathcal{X}\to\mathbb{R},\ell\colon\mathbb{R}\times\mathbb{R}\to\mathbb{R}$ is a model and a loss function.

Also, it is recognized as a difficult task as models are poor at learning non-linearity outside the training distributions. This issue was examined by Xu et al.(2021) [22] proving that MLP with ReLU as an activation function could not converge nonlinear function outside training distribution.

In molecular property prediction, the input is a molecular graph and the output is a continuous label. Since a set of graphs can be represented as $\mathcal{G}=(V,E)$ using a set of nodes $V$ representing atoms, and edges $E$ representing bonds, the problem is to learn the relationships $g\colon\mathcal{G}\to\mathbb{R}$. To model it, GNNs and transformer-based models are extensively developed [13, 14, 25].

Although they mainly work in i.i.d. assumption, some research tackled extrapolation, or out-of-distribution generalization in molecular property prediction. Then, they are specified according to the formation of data distributions: one based on input features or based on out labels.

Most previous studies [23, 12] defined extrapolation by dividing data distributions based on input features. For example, Yang et al. (2022) [23] use scaffolds and sizes of graphs to decide distributions. Although Chan et al. (2021)[2] formed distributions by output labels, its aim is generating data rather than prediction. Whereas the definition based on input features has the advantage of making models robust to a wide variety of inputs, it doesn’t meet the needs of finding material extraordinary property value. To realize it, a new definition of extrapolation based on output labels in molecular property prediction is expected.

Self-supervised learning is a method of training a model using only input data without labels. Since self-supervised learning does not use labels, it has been widely studied for utilizing unlabeled data. Considering the fact that there are plenty of molecules whose target property values are unknown, it is a powerful tool for our research. The following three types of pretext tasks are related to self-supervised learning on graphs.

Attribute Level Task. A popular method for node-level pretext tasks is a node and edge masking problem. This task is a classification problem, and the model predicts true attributes from the hidden inputs by masking. These masking tasks were first developed in the field of natural language processing, such as BERT [3], and have been applied to graphs as well. Through this task, encoders can gain effective information related to graph attributes. [19, 5]

Structure Level Task. The task handles structures in graphs and extracts important information regarding the connectivity between nodes without considering the attributes. Structures strongly influence the characteristics of molecules because electrons involved in the bonds have an important role to determine the property values. For example, the number of rings in a graph can be a useful indicator for grasping molecules.

Attribute + Structure Level Task. The task is mainly to predict the properties of subgraphs with attributes and structures. This often needs to prepare a set of specific graphs in advance for training objectives, e.g., whether the specific graphs are included in a given input graph. Most of them perform graph matching between subgraphs of inputs and specific graphs in a set of graphs called motif vocabulary [19]. The motif vocabularies are often created using domain knowledge in chemistry.

Instead of the traditional definition, we define extrapolation by dividing data distributions based on output labels. We formulate this definition using the same notations in section 2.1, and it can be described as follows.

In this research, two hypotheses regarding the effects of self-supervised pretraining on extrapolation will be verified. We give an overview of the hypotheses and the technical reasons why they are made.

Firstly, we give an overview of the flow when using self-supervised pretraining in our experiment in Figure 2. Training models in this experiment consists of two steps: self-supervised pretraining, and fine-tuning. Also, there are three important elements: 1) the model for predicting property from the molecular graph; 2) the way to train the model; 3) the method of data split; We will introduce all of them in this section.

To validate the two hypotheses, we train models by the following 7 training methods.

No pretraining (baseline). As a baseline, models are trained by only fine-tuning with training data. Then their interpolation and extrapolation performance is evaluated by holdout validation and forward holdout validation.

Pretraining without validation data (task1, task2, task3). To check on whether utilizing unlabeled data in self-supervised pretraining is effective or not, the models are firstly pretrained with only training data. As self-supervised pretraining, task1, task2, and task3 are used. After that, they are fine-tuned by training data. Their extrapolation performance is evaluated by forward holdout validation.

Pretraining with validation data (task1, task2, task3). As opposed to the above methods, in this training method, we use both training and validation data when pretraining. The rest procedures are the same as above.

We evaluated all of the performances by Mean Absolute Error. Firstly, we compare the result of interpolation and extrapolation in the baseline. The average of minimum MAEs in interpolation reaches 0.091, whereas that in extrapolation does 0.692. Additionally, the minimums in extrapolation are recorded from the 1st epoch to the 10th epoch. It means that, in extrapolation, the model is considered to be poorly learned.

To investigate the issue in more detail, we explain Figure 6 which shows the relationship between the predicted value and the label. First of all, in interpolation, the validation data is plotted along the black dotted line. On the other hand, in extrapolation, the validation data are far apart, and the predicted value hardly exceeds a certain value. One of the reasons for this result may be that there are no data with labels greater than about 8 in the training data. However, since a large number of predicted values exist around 8, the model seems to be able to predict the tendency that the validation data is relatively high in the dataset.

Figure 6 shows the mean absolute error in extrapolation with 7 training methods. Considering baseline and pretraining with validation data, we can find that the MAEs of task1, 2, and 3 are lower than that of the baseline method When compared to the baseline, the values decreased by 4.91% for method1, 3.47% for method2, and 9.68% for method3. However, these MAEs are much larger than the interpolation MAEs, and the extrapolation performances are far behind the interpolation.

When focusing on the result of pretraining with and without unlabeled data, tasks 1 and 2 with unlabeled data are almost 3% better than without them. In contrast, the MAE of task 3 increased by 0.8% when unlabeled data is used in pretraining.

In contrast, all of the MAEs are significantly worse than that in interpolation. It means that self-supervised pretraining can’t help models extrapolate exact values. Therefore, in terms of predicting exactly, the hypothesis h1 and h2 are false.

Also, we illustrate Figure 7 to show the relationship between labels and predicted values.In these four figures, the models could not almost predict values higher than 8. In contrast, the difference between them is that the validation data in the baseline figure is slightly tilted to the lower right, while in task1 and 2 is almost parallel, and in task3 slightly tilted to the upper right. This fact suggests that pretraining may help the model to learn relative trends in the validation data. From this point of view, we considered that MAE alone is insufficient as an evaluation metric for extrapolation in this study, and decide to conduct further analysis.

To analyze the relative tendency of the predictions and find out the difference between pretraining and non-pretraining methods more clearly, we introduce the idea of ranking.

This is an index that indicates whether the ranking in descending order of physical property value is correctly predicted in the validation data or not. In addition, in the field of materials chemistry, understanding this ranking also has a positive effect on improving the efficiency of material development, so we decided to adopt this index. Furthermore, ranking is a discrete value, and it is expected that the plot will be more sparse than Figure 7, making it easier to see the difference from the figure.

To evaluate it, the rankings of the labels and the predicted values are calculated in advance, and the correlation coefficient is used as the evaluation metrics. The reason why we choose the metric is that since the number of validation data is about 70,000, the metric is suitable for grasping the macro trends related to rankings.

Figure 9 illustrates the ranking correlation coefficients in extrapolation with 8 training methods. What we can see from this figure is that the correlation coefficients for task1, 2, and 3 when using validation data are much closer to 1.0 than the baseline. Specifically, the figure of task3 is 0.516, which is the best result. On the other hand, the baseline is 0.255, which is the lowest value. Correlation coefficients in the range from 0.4 to 0.7 are considered to have some degree of positive correlation, while if it is between -0.2 and 0.2 there is almost no correlation. It can be said that self-supervised pretraining enhanced the ability to judge whether certain data is greater than other data in the validation data.

Then, for checking hypothesis h2, we contrast pretraining with and without validation data. The result is that the metrics in all of the tasks decrease from almost 0.5 to 0.4 when removing validation data from self-supervised pretraining. Therefore, it can be concluded that validation data has a good influence on extrapolation performance in terms of predicting the relative tendency.

Also, we calculated the correlation coefficients for each epoch and gain Figure 9 showing the change through training. From this figure, we can see the fact that by 10 epochs the correlation coefficient values fall into negative for most methods. Especially in the first 5 epochs, the correlation coefficient is very high, and the method using pretraining has a value of 0.4 or higher, but it can be seen that it drops sharply after the epoch. This phenomenon might be because, through training, a model that is too well-fitted to the training data fails to predict label values so high that the training data doesn’t include it. The issue is a significant limitation in our study and should be well analyzed.