SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Gang Liu
University of Notre Dame gliu7@nd.edu
Eric Inae
University of Notre Dame einae@nd.edu
Tong Zhao
Snap Inc. tzhao@snap.com
Jiaxin Xu
University of Notre Dame jxu24@nd.edu
Tengfei Luo
University of Notre Dame tluo@nd.edu
Meng Jiang
University of Notre Dame mjiang2@nd.edu
Graph property prediction tasks are important and numerous. While each task offers a small size of labeled examples, unlabeled graphs have been collected from various sources and at a large scale. A conventional approach is training a model with the unlabeled graphs on self-supervised tasks and then fine-tuning the model on the prediction tasks. However, the self-supervised task knowledge could not be aligned or sometimes conflicted with what the predictions needed. In this paper, we propose to extract the knowledge underlying the large set of unlabeled graphs as a specific set of useful data points to augment each property prediction model. We use a diffusion model to fully utilize the unlabeled graphs and design two new objectives to guide the model's denoising process with each task's labeled data to generate task-specific graph examples and their labels. Experiments demonstrate that our data-centric approach performs significantly better than fifteen existing various methods on fifteen tasks. The performance improvement brought by unlabeled data is visible as the generated labeled examples unlike the self-supervised learning.
Graph data such as molecules and polymers are found to have attractive properties in drug and material discovery (Böhm et al., 2004; Huang et al., 2021), but annotating them requires specialized knowledge, as well as lengthy and costly experiments in wet labs (Cormack and Elorza, 2004). So, it is important for graph property predictors to learn useful knowledge from unlabeled graphs.
Self-supervised learning (Hu et al., 2019; Rong et al., 2020; You et al., 2021; Kim et al., 2022) utilizes unlabeled graphs to learn through predictive tasks or contrastive tasks to represent and transfer the knowledge as model parameters. Despite the empirical success in language and vision (Brown et al., 2020; He et al., 2022), their performance on graph data applications remains unsatisfactory because of the significant gap between the graph self-supervised task and the graph label prediction task. Models trained on node attribute prediction (Hu et al., 2019) as a simple predictive self-supervised task extract too limited knowledge from the graph structure, which has been observed after too fast convergence (Sun et al., 2022). More complex tasks like graph context prediction (Hu et al., 2019; Zhang et al., 2021) may transfer knowledge that conflicts with downstream tasks. Aromatic rings, for instance, are a prevalent structure in molecules (Maziarka et al., 2020) and are considered valuable in context prediction tasks (Zhang et al., 2021). However, graph properties such as oxygen permeability can be more related to non-aromatic rings in some cases (Liu et al., 2022a), which is overlooked if not using tailored predictive tasks specifically for downstream tasks. As predictive tasks strive for
Figure 1: Comparing the diagrams of the existing approach and the proposed approach to learning from unlabeled graphs for a variety of graph property prediction tasks.
universality, the transferred knowledge may force models to focus more on aromatic rings, leading to poor prediction.
On the other line, contrastive tasks (You et al., 2021; Kim et al., 2022) aim to learn the similarity between original and perturbed graphs. However, the learned similarity can hardly generalize across tasks (Kim et al., 2022). First, perturbations without domain knowledge, e.g., bioisosteres, do not preserve broad biological properties (Sun et al., 2021). Second, it is difficult, if not impossible, to find universal perturbations that generalize to diverse property prediction tasks. For example, bioisosteric (subgraph) replacements produce similar biological properties for molecules. And they may reduce toxicity (Brown, 2014). So, contrastive tasks with bioisosteric replacement enforce the similarity between toxic and non-toxic molecules. However, models pre-trained on such contrastive tasks hurt the performance on downstream tasks, e.g., toxicity prediction.
Our data-centric idea avoids the use of self-supervised tasks that are not appropriate. We use a diffusion probabilistic model (known as diffusion model) to capture the data distribution of unlabeled graphs, leveraging its capability of distribution coverage, stationarity, and scalability (Dhariwal and Nichol, 2021). At the stage of performing a particular property prediction task, the reverse process, guided by novel task-related optimization objectives, generates new task-specific labeled examples. Minimal sufficient knowledge from the unlabeled data is transferred into these examples, instead of uninterpretable model parameters, and then to enhance the training of prediction models.
To implement our idea, we propose a Data-Centric Transfer framework (DCT) based on a diffusion model for graph data, as shown in Figure 1b. It aims to transfer minimal sufficient knowledge from unlabeled graphs to property predictors by data augmentation. The diffusion model gradually adds Gaussian noise to a graph from which a score function (i.e., the gradient of the log probability density) is then learned to estimate the noise step by step to reverse the process. DCT trains the diffusion model on the unlabeled graphs to get ready to augment any labeled dataset. Given a labeled graph from a particular task (i.e., type of property), the diffusion model adds noise to perturb it by a few steps and then generates a new graph through the score function. The new graph could be close to the distribution of the unlabeled graphs for diversity, however, it would lose the relatedness to the target task. So, we add two task-related objectives into the score function to guide the reverse process. When a predictor model  $f$  has been trained on the task, given an original labeled graph  $G$ , the first objective is to optimize the new graph  $G'$  to sufficiently preserve the label of  $G$  with  $f$ . The second objective is to optimize  $G'$  to be very different from (i.e., minimally similar to)  $G$ . These two objectives ensure that  $G'$  carries minimal sufficient knowledge from the unlabeled graphs to be an augmentation of  $G$ . DCT iteratively generates new examples to augment the labeled dataset and progressively trains the prediction model with it.
We test DCT on fifteen graph property prediction datasets from three fields: chemistry (molecules), material science (polymers), and biology (protein-protein interaction graphs). DCT achieves the best performance over all these tasks. We find that the state-of-the-art self-supervised methods often struggle to transfer knowledge to regression tasks, etc. DCT reduces the mean absolute error relatively by  $13.4\%$  and  $10.2\%$  compared to the best baseline on the molecule and polymer graph regression tasks, respectively.
Given  $K$  property prediction tasks, there are  $N^{[k]}$  labeled graph examples for the  $k$ -th task. They are  $\{(G_i,y_i)\mid G_i\in \mathcal{G}^{[k]},y_i\in \mathcal{Y}^{[k]}\}_{i = 1}^{N^{[k]}}$ , where  $\mathcal{G}^{[k]}$  is the graph space and  $\mathcal{V}^{[k]}$  is the label space of the task. The prediction model with parameters  $\theta$  is defined as  $f_{\theta}^{[k]}\colon \mathcal{G}^{[k]}\to \mathcal{V}^{[k]}$ .  $f_{\theta}^{[k]}$  consists of a GNN and a multi-layer perceptron (MLP). Without the loss of generality, we consider Graph Isomorphism Networks (GIN) (Xu et al., 2019) to encode graph structures. Given a graph  $G = (\mathcal{V},\mathcal{E})\in \mathcal{G}^{[k]}$  in the task  $k$ , GIN updates the representation vector of node  $v\in \mathcal{V}$  at  $l$ -layer:
$$
\mathbf {h} _ {v} ^ {l} = \operatorname {M L P} ^ {l} \left((1 + \epsilon) \cdot \mathbf {h} _ {v} ^ {l - 1} + \sum_ {u \in \mathcal {N} (v)} \mathbf {h} _ {u} ^ {l - 1}\right), \tag {1}
$$
where  $\epsilon$  is a learnable scalar and  $u\in \mathcal{N}(v)$  is one of node  $v$  's neighbor nodes. After stacking  $L$  layers, the READOUT function (e.g., summation) gets the graph representation across all the nodes. The predicted label is:
$$
\hat {y} = \operatorname {M L P} \left(\operatorname {R E A D O U T} \left(\left\{\mathbf {h} _ {v} ^ {L} \mid v \in G \right\}\right)\right). \tag {2}
$$
$f_{\theta}^{[k]}$  is hard to be well-trained because it is hard to collect graph labels at a large scale ( $N^{[k]}$  is small). Fortunately, regardless of the tasks, a large number of unlabeled graphs are usually available from the same or similar domains. Self-supervised learning methods (Hu et al., 2019) rely on hand-crafted tasks to extract knowledge from the unlabeled examples  $\{G_j \in \mathcal{G}^{[U]}, j = 1, \dots, N\}$  as pre-trained model parameters  $\theta$ . The uninterpretable parameters are transferred to warm up the prediction models  $\{f_{\theta}^{[k]}\}_{k=1}^{K}$  on the  $K$  downstream graph property prediction tasks. However, the gap and even conflict between the self-supervised tasks and the property prediction tasks lead to suboptimal performance of the prediction models. In the next section, we present the DCT framework that transfers knowledge from the unlabeled graphs with a data-centric approach.
3 The Data-Centric Transfer Framework
The goal of data-centric approaches is to augment training datasets by generating useful labeled data examples. Under that, the goal of our data-centric transfer (DCT) framework is to transfer the knowledge from unlabeled data into the data augmentation. Specifically, for each graph-label pair  $(G^{[k]} \in \mathcal{G}^{[k]}, y^{[k]} \in \mathcal{Y}^{[k]})$  in the task  $k$ , the framework is expected to output a new example  $G^{\prime [k]}$  with the label  $y^{\prime [k]}$  such that (1)  $y^{\prime [k]} = y^{[k]}$  and (2)  $G^{\prime [k]}$  and  $G^{[k]}$  are from the same graph space  $\mathcal{G}^{[k]}$ . However, if the graph structures of  $G^{\prime [k]}$  and  $G^{[k]}$  were too similar, the augmentation would duplicate the original data examples, become useless, and even cause over-fitting. So, the optimal graph data augmentation should enrich the training data with good diversity as well as preserve the labels of the original graphs. To achieve this, DCT utilizes a diffusion probabilistic model to first learn the data distribution from unlabeled graphs (Section 3.2). Then DCT adapts the reverse process in the diffusion model to generate task-specific labeled graphs for data augmentation (Section 3.3). Thus, the augmented graphs will be derived from the distribution of a huge collection of unlabeled data for diversity. To preserve the labels, DCT controls the reverse process with two task-related optimization objectives to transfer minimal sufficient knowledge from the unlabeled data. The first objective minimizes an upper bound of mutual information between the augmented and the original graphs in the graph space. The second objective maximizes the probability of the predicted label of augmented graphs being the same as the label of original graphs. The first is for minimal knowledge transfer, and the second is for sufficient knowledge transfer. DCT integrates the two objectives into the reverse process of the diffusion model to guide the generation of new labeled graphs. DCT iteratively trains the graph property predictor (used in the second objective) and creates the augmented training data. To simplify notations, we remove the task superscript  $[k]$  in the following sections.
3.2 Learning Data Distribution from Unlabeled Graphs
The diffusion process for graphs in Figure 2 applies to both graph structure and node features. The diffusion model slowly corrupts unlabeled graphs to a standard normal distribution with noise. For
Figure 2: Diffusion model in DCT: It performs task-specific data augmentation using objectives  $\mathcal{I}_1$  and  $\mathcal{I}_2$  in the reverse process. The model was trained on unlabeled graphs to learn the general data distribution. Then it generates  $(G', y' = y)$  based on  $(G, y)$  in the reverse process. It perturbs  $G$  with  $D$  steps and optimizes  $G'$  to be minimally similar to  $G$  (Objective  $\mathcal{I}_1$ ) and sufficiently preserve the label of  $G$  (Objective  $\mathcal{I}_2$ ).
Figure 3: Relative improvement (increased AUC or reduced MAE) from three data-centric methods (over ten runs), compared to the basic GIN: Blue is for self-training with selected real unlabeled graphs. Green is for self-training with graphs directly generated by a standard diffusion model. Red is for DCT that generates task-specific labeled graphs. The first two often make little or negative impact. Our DCT has consistent and significant improvement shown as the percentages in red.
graph generation, the model samples noise from the normal distribution and learns a score function to reverse the perturbed noise. Given an unlabeled graph  $G$ , we use continuous time  $t \in [0,T]$  to index multiple diffusion steps  $\{G^{(t)}\}_{t=1}^T$  on the graph, such that  $G^{(0)}$  follows the original data distribution and  $G^{(T)}$  follows a prior distribution like the normal distribution. The forward diffusion is a stochastic differential equation (SDE) from the graph to the noise:
$$
\mathrm {d} G ^ {(t)} = \mathbf {f} (G ^ {(t)}, t) \mathrm {d} t + g (t) \mathrm {d} \mathbf {w}, \tag {3}
$$
where  $\mathbf{w}$  is the standard Wiener process,  $\mathbf{f}(\cdot ,t):\mathcal{G}\to \mathcal{G}$  is the drift coefficient and  $g(t):\mathbb{R}\rightarrow \mathbb{R}$  is the diffusion coefficient.  $\mathbf{f}(G^{(t)},t)$  and  $g(t)$  relate to the amount of noise added to the graph at each infinitesimal step  $t$ . The reverse-time SDE uses gradient fields or scores of the perturbed graphs  $\nabla_{G^{(t)}}\log p_t(G^{(t)})$  for denoising and graph generation from  $T$  to 0 (Song et al., 2021):
$$
\mathrm {d} G ^ {(t)} = \left[ \mathbf {f} \left(G ^ {(t)}, t\right) - g (t) ^ {2} \nabla_ {G ^ {(t)}} \log p _ {t} \left(G ^ {(t)}\right) \right] \mathrm {d} t + g (t) \mathrm {d} \overline {{\mathbf {w}}}, \tag {4}
$$
where  $p_t(G^{(t)})$  is the marginal distribution at time  $t$  in forward diffusion.  $\overline{\mathbf{w}}$  is a reverse time standard Wiener process. dt here is an infinitesimal negative time step. The score  $\nabla_{G^{(t)}}\log p_t(G^{(t)})$  is unknown in practice and it is approximated by the score function  $\mathbf{s}(G^{(t)},t)$  with score matching techniques (Song et al., 2021). On graphs, Jo et al. (2022) used two GNNs to develop the score function  $\mathbf{s}(G^{(t)},t)$  to de-noise both node features and graph structures and details are in appendix B.3.
3.3 Generating Task-specific Labeled Graphs
Self-training approaches would propose to either (1) select unlabeled graphs by a graph property predictor or (2) generate graphs directly from the standard diffusion model, and then use the predictor to assign them labels so that the training data could be enriched. However, we have observed that neither of them can guarantee positive impact on the prediction performance. In fact, as shown in Figure 3, they make very little or even negative impact. That is because the selected or directly-generated graphs are too different from the labeled graph space of the target tasks. Task details of ten datasets are in appendix C.
Given a labeled graph  $(G, y)$  from the original dataset of a specific task, the new labeled graph  $(G', y')$  is expected to provide useful knowledge to augment the training set. We name it the augmented graph throughout this section. The augmented graph is desired to have the following two properties, as in Section 3.1: Task relatedness: As an effective training data point,  $G' \in \mathcal{G}$  and  $y' \in \mathcal{Y}$  are from the graph/label spaces of the specific task where  $(G, y)$  come from and thus transfer sufficient task knowledge into the training set; Diversity: If  $G'$  was too similar to  $G$ , the new data point would cause severe over-fitting on the property prediction model. The augmentation aims to learn from unlabeled graph to create diverse data points, which should contain minimal task knowledge about  $G$ .
The selected unlabeled graphs used in self-training have little task relatedness because the unlabeled data distribution might be too far from the one of the specific task. Existing graph data augmentation methods could not create diverse graph examples because they manipulated labeled graphs and did not learn from the unlabeled graphs. Our novel data-centric approach DCT works towards both desired properties by transferring minimally sufficient knowledge from the unlabeled graphs: Sufficiency is achieved by maximizing the possibility for label preservation (i.e.,  $y' = y$ ). It ensures that the knowledge from unlabeled graphs is task-related; Minimality refers to the minimization of graph similarity between  $G'$  and  $G$  to ensure that the augmentation introduces diversity. Both optimizations can be formulated using mutual information  $\mathcal{I}(\cdot ;\cdot)$  to generate task-specific labeled data  $(G', y')$ :
Definition 3.1 (Sufficiency for Data Augmentation). The augmented graph  $G'$  sufficiently preserves the label of the original graph  $G$  if and only if  $\mathcal{I}(G'; y) = \mathcal{I}(G; y)$ .
Definition 3.2 (Minimal Sufficiency for Data Augmentation). The Sufficiency is minimal for data augmentation if and only if  $\mathcal{I}(G'; G) \leq \mathcal{I}(\bar{G}; G)$ ,  $\forall \bar{G}$  represents any augmented graph that sufficiently preserves the original graph's label.
Self-supervised tasks applied a similar philosophy in pre-training (Soatto and Chiuso, 2016), however, they did not use labeled data from any specific tasks. So the optimizations were unable to extract useful knowledge and transfer it to the downstream (Tian et al., 2020). In our DCT that performs task-specific data augmentation, the augmented graphs can be optimized toward the objectives using any labeled graph  $G$  and its label  $y$ :
$$
\min  _ {\mathcal {I} _ {1}} \max  _ {\mathcal {I} _ {2}} \mathbb {E} _ {G} \left[ \mathcal {I} _ {1} \left(G ^ {\prime}; G\right) + \mathcal {I} _ {2} \left(G ^ {\prime}; y\right) \right]. \tag {5}
$$
For the first objective, we use the leave-one-out variant of InfoNCE (Poole et al., 2019; Oord et al., 2018) as the upper bound estimation. For the  $i$ -th labeled graph  $(G_i, y_i)$ ,
$$
\mathcal {I} _ {1} \leq \mathcal {I} _ {\text {b o u n d}} \left(G _ {i} ^ {\prime}; G _ {i}\right) = \log \frac {p \left(G _ {i} ^ {\prime} \mid G _ {i}\right)}{\sum_ {j = 1 , j \neq i} ^ {M} p \left(G _ {i} ^ {\prime} \mid G _ {j}\right)}, \tag {6}
$$
where  $G_{i}^{\prime}$  is the augmented graph. When  $G_{i}^{\prime}$  is optimized,  $G_{i}$  makes a positive pair;  $\{G_{j}\} (j\neq i)$  are  $M - 1$  negative samples of labels that do not equal  $y_{i}$ . ( $M$  is a hyperparameter.) We use cosine similarity and a softmax function to calculate  $p(G_i'|G_j) = \frac{\exp(\sin(G_i',G_j))}{\sum_{j = 1}^{M}\exp(\sin(G_i',G_j))}$ . In practice, we extract statistical features of graphs to calculate their similarity. Details are in appendix B.2.
For the second objective, we denote the predicted label of the augmented graph  $G'$  by  $f_{\theta}(G')$ . We maximize the log likelihood  $\log p(y|f_{\theta}(G'))$  to maximize  $\mathcal{I}_2(G';y)$ . Specifically, after the predictor  $f_{\theta}$  is trained for several epochs on the labeled data, we freeze its parameters and use it to optimize the augmented graphs so they are task-related:
$$
\mathcal {L} \left(G ^ {\prime}\right) = \mathcal {I} _ {\text {b o u n d}} \left(G ^ {\prime}; G\right) - \log p (y | f _ {\theta} \left(G ^ {\prime}\right)). \tag {7}
$$
**Framework details:** As shown in Figure 2, after the diffusion model learns the data distribution from unlabeled graphs, given a labeled graph  $G$  from a specific task, DCT perturbs it for  $D$  ( $D \ll T$ ) steps. The perturbed noisy graph, denoted by  $\tilde{G}^{(D)}$ , stays inside the task-specific graph and label space, rather than the noise distribution (at step  $T$ ). To reverse the noise in it and generate a task-specific augmented example  $G'$ , DCT integrates the loss function in Eq. (7) into the score function  $\mathbf{s}(\cdot, t)$  for minimal sufficient knowledge transfer:
$$
\mathrm {d} \tilde {G} ^ {(t)} = \left[ \mathbf {f} \left(\tilde {G} ^ {(t)}, t\right) - g (t) ^ {2} \left(\mathbf {s} \left(\tilde {G} ^ {(t)}, t\right) - \alpha \nabla_ {\tilde {G} ^ {(t)}} \mathcal {L} \left(\tilde {G} ^ {(t)}\right)\right) \right] \mathrm {d} t + g (t) \mathrm {d} \overline {{\mathbf {w}}}, \tag {8}
$$
where  $\alpha$  is a scalar for score alignment between s and  $\nabla \mathcal{L}$  to avoid the dominance of any of them:  $\alpha = \frac{\|\mathbf{s}(\tilde{G}^{(t)},t)\|_2}{\|\nabla_{\tilde{G}^{(t)}}\mathcal{L}(\tilde{G}^{(t)})\|_2}$ . Because  $\tilde{G}^{(t)}$  is an intermediate state in the reverse process, the noise in it
Table 1: Statistics of datasets for graph property prediction in different domains.
Data Type	Dataset	# Graphs	Prediction Task	# Task	Avg./Max # Nodes	Avg./Max # Edges
Molecules	ogbg-HIV	41,127	Classification	1	25.5 / 222	54.9 / 502
ogbg-ToxCast	8,576	Classification	617	18.8 / 124	38.5 / 268
ogbg-Tox21	7,831	Classification	12	18.6 / 132	38.6 / 290
ogbg-BBBP	2,039	Classification	1	24.1 / 132	51.9 / 290
ogbg-BACE	1,513	Classification	1	34.1 / 97	73.7 / 202
ogbg-ClinTox	1,477	Classification	2	26.2 / 136	55.8 / 286
ogbg-SIDER	1,427	Classification	27	33.6 / 492	70.7 / 1010
ogbg-Lipo	4200	Regression	1	27 / 115	59 / 236
ogbg-ESOL	1128	Regression	1	13.3 / 55	27.4 / 124
ogbg-FreeSolv	642	Regression	1	8.7 / 24	16.8 / 50
Polymers	GlassTemp	7,174	Regression	1	36.7 / 166	79.3 / 362
MeltingTemp	3,651	Regression	1	26.9 / 102	55.4 / 212
ThermCond	759	Regression	1	21.3 / 71	42.3 / 162
O2Perm	595	Regression	1	37.3 / 103	82.1 / 234
Proteins	PPI	88000	Classification	40	49.4 / 111	890.8 / 11556
may fail the optimizations. So, we design a new sampling method named double-loop sampling for accurate loss calculation. It has an inner-loop sampling using Eq. (4) to sample  $\hat{G}_{(t)}$ , as the denoised version of  $\tilde{G}^{(t)}$  at the reverse time  $t$ . Then  $\nabla_{\hat{G}}\mathcal{L}(\hat{G}_{(t)})$  is calculated as an alternative for  $\nabla_{\tilde{G}^{(t)}}\mathcal{L}(\tilde{G}^{(t)})$ . Finally, an outer-loop sampling takes one step to guide denoising using Eq. (8).
DCT iteratively creates the augmented graphs  $(G', y')$ , updates the training dataset  $\{(G_i, y_i)\}$ , and trains the graph property predictor  $f_\theta$ . In each iteration, for task  $k$ ,  $n \ll N^{[k]}$  labeled graphs of the lowest property prediction loss are selected to create the augmented graphs. The predictor is better fitted to these graphs for more accurate sufficiency estimation of the augmentation.
In this section, we present and analyze experimental results to demonstrate the outstanding performance of DCT, the usefulness of new optimization objectives, the effect of hyperparameters and iterative process, and the interpretability of "visible" knowledge transfer from unlabeled graphs.
Tasks and metrics: Experiments are conducted on 15 graph property prediction tasks in chemistry, material science, and biology, including seven molecule classification, three molecule regression tasks from open graph benchmarks (Hu et al., 2020), four polymer regression tasks, and protein function prediction (PPI) (Hu et al., 2019). Dataset statistics is presented in Table 1. We use the area under the ROC curve (AUC) to evaluate classifiers and mean absolute error (MAE) for regressors.