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 Existing approach: Knowledge from selfsupervised tasks could not be aligned or even conflict with what predictions need. Parameter initialization could hardly interpret how unlabeled graphs were or would be able to improve the models, leading to high prediction errors.
Labeled graphs in task 𝑘 G 𝑖 ∈ G 𝑘 , 𝑦 𝑖 ∈ Y 𝑘 𝑖=1 𝑁 [𝑘] ≪𝑁 Generated graph for task 𝑘 (G′ ∈ G 𝑘 , 𝑦 ′ ∈ Y 𝑘 )
Tasks 𝑘 = 1 … 𝐾 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 0 to sufficiently preserve the label of G with f . The second objective is to optimize G 0 to be very different from (i.e., minimally similar to) G. These two objectives ensure that G 0 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
i=1 , where G [k] is the graph space and Y [k] is the label space of the task. The prediction model with parameters ✓ is defined as f
✓ 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 = (V, E) 2 G [k] in the task k, GIN updates the representation vector of node v 2 V at l-layer:
where ✏ is a learnable scalar and u 2 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:
✓ 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 2 G [U ] , j = 1, . . . , N} as pre-trained model parameters ✓. The uninterpretable parameters are transferred to warm up the prediction models {f
✓ } K k=1 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] 2 G [k] , y [k] 2 Y [k] ) in the task k, the framework is expected to output a new example G 0 [k] with the label y 0 [k] such that (1) y 0[k] = y [k] and (2) G 0 [k] and G [k] are from the same graph space G [k] . However, if the graph structures of G 0 [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.
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
d ෨ 𝐺 .
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 2 [0, T ] to index multiple diffusion steps {G (t) } T t=1 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:
where w is the standard Wiener process, f (•, t) : G ! G is the drift coefficient and g(t) : R ! R is the diffusion coefficient. 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 r G (t) log p t (G (t) ) for denoising and graph generation from T to 0 (Song et al., 2021):
where p t (G (t) ) is the marginal distribution at time t in forward diffusion. w is a reverse time standard Wiener process. dt here is an infinitesimal negative time step. The score r G (t) log p t (G (t) ) is unknown in practice and it is approximated by the score function 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 s(G (t) , t) to de-noise both node features and graph structures and details are in appendix B.3.
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 directlygenerated 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 0 , y 0 ) 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 0 2 G and y 0 2 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 0 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 0 = y). It ensures that the knowledge from unlabeled graphs is task-related; Minimality refers to the minimization of graph similarity between G 0 and G to ensure that the augmentation introduces diversity. Both optimizations can be formulated using mutual information I(• ; •) to generate task-specific labeled data (G 0 , y 0 ):
Definition 3.1 (Sufficiency for Data Augmentation). The augmented graph G 0 sufficiently preserves the label of the original graph G if and only if I(G 0 ; y) = I(G; y).
Definition 3.2 (Minimal Sufficiency for Data Augmentation). The Sufficiency is minimal for data augmentation if and only if I(G 0 ; G) I( Ḡ; G), 8 Ḡ 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:
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 ),
where G 0 i is the augmented graph. When G 0 i is optimized, G i makes a positive pair; {G j } (j 6 = 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
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 0 by f ✓ (G 0 ). We maximize the log likelihood log p (y|f ✓ (G 0 )) to maximize I 2 (G 0 ; y). Specifically, after the predictor f ✓ 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:
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 ⌧ T ) steps. The perturbed noisy graph, denoted by 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 taskspecific augmented example G 0 , DCT integrates the loss function in Eq. ( 7) into the score function s(•, t) for minimal sufficient knowledge transfer:
where ↵ is a scalar for score alignment between s and rL to avoid the dominance of any of them: . Finally, an outer-loop sampling takes one step to guide denoising using Eq. ( 8).
DCT iteratively creates the augmented graphs (G 0 , y 0 ), updates the training dataset {(G i , y i )}, and trains the graph property predictor f ✓ . In each iteration, for task k, n ⌧ 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.
Baselines and implementation: Besides GIN, there are three lines of baseline methods:
(1) selfsupervised learning methods including EDGEPRED, ATTRMASK, CONTEXTPRED in (Hu et al., 2019), INFOMAX (Velickovic et al., 2019), JOAO (You et al., 2021), GRAPHLOG (Xu et al., 2021), MGSSL Zhang et al. (2021) and D-SLA (Kim et al., 2022), (2) semi-supervised learning methods including self-training with selected unlabeled graphs (ST-REAL) and generated graphs (ST-GEN) and INFOGRAPH (Sun et al., 2020), and (3) graph data augmentation (GDA) methods including FLAG (Kong et al., 2022), GREA (Liu et al., 2022a), and G-MIXUP (Han et al., 2022). For self-supervised pre-training, we follow their own settings and directly use their pre-trained models if available. For semi-supervised learning methods and DCT, we use 113K QM9 (Ramakrishnan et al., 2014) and 306K PPI graphs (Hu et al., 2019) as unlabeled data sources for the tasks on molecules/polymers and proteins, respectively. For DCT, we tune three major hyper-parameters: the number of perturbation steps D 2 [1, 10], the number of negative samples M 2 [1, 10], and top-n % labeled graphs of lowest property prediction loss selected for data augmentation.
We report the model performance using mean and standard deviation over 10 runs Table 2. DCT is the best on all 15 tasks compared to the state-of-the-art baselines. Our observations are: (1) GIN is the most competitive baseline and outperforms self-supervised learning methods. On 7 of 15 tasks, GIN outperforms all the 7 self-supervised learning methods. Because self-supervised pre-training imposes constraints on the model architecture, it undermines the true power of GNNs and under-performs the GNNs that are properly used.
(2) Self-training and GDA methods perform better than GIN but cannot effectively learn from unlabeled data. Self-training (ST-REAL and ST-GEN) is often the best baseline in regression tasks. GDA (GREA and G-MIXUP) methods outperform self-training in most classification tasks except ogbg-SIDER, because they are often designed to exploit categorical labeled data and remain underexplored for regression. Although self-training benefits from selecting unlabeled examples in some graph regression tasks, they are negatively affected by the unlabeled graphs in the classification tasks such as ogbg-ToxCast and ogbg-ClinTox. As indicated in Figure 3, it is inappropriate to pseudo-label unlabeled graphs in self-training due to the huge gap between the unlabeled data and target task.
(3) DCT transfers useful knowledge from unlabeled data by data augmentation. DCT outperforms the best baseline relatively by +3.9%, +13.4%, and +10.2% when there are only 1,210, 513, and 4,303 training graphs on ogbg-BACE, ogbg-FreeSolv, and GlassTemp, respectively. Compared to the self-supervised baselines, the improvement from DCT is more significant, so the knowledge transfer is more effective. For example, on ogbg-FreeSolv and O 2 Perm, DCT performs better than the best self-supervised baselines relatively by +45.8% and +8.0%, respectively. On regression tasks that involve knowledge transfer across domains (e.g., from molecules to polymers), DCT reduces MAE relatively by 1.9% ⇠ 10.2% compared to the best baseline. All these results demonstrate the outstanding performance of task-specific data augmentation in DCT.
Comprehensive ablation studies: In Table 3, we investigate how the task-related objectives in Eq. ( 5) impact the performance of DCT. First, DCT outperforms the top baseline even if the two task-related optimization objectives are disabled. This is because DCT generates new training examples based on original labeled graphs: the data augmentation has already improved the diversity of the training dataset a little bit. Second, adding the objective I 1 further improves the performance by encouraging the generation of diverse examples, because it minimizes the similarity between the original graph and augmented graph in the graph space. Third, we receive the best performance of DCT when it combines I 1 and I 2 objectives to generate task-related and diverse augmented graphs. When we change the unlabeled data source from QM9 to the ZINC dataset from (Jo et al., 2022), similar observations confirm the necessity of the task-related objectives.
The impacts of three hyper-parameters of DCT are studied: the number of perturbation steps D, the number of negative samples M in Eq. ( 6), and the number of augmented graphs in each iteration (i.e., top-n % selected graph for augmentation). Results from Figure 4 show that DCT is robust to a wide range of D and M valued from 0 to 10. They suggest that D and M can be set as 5 in most cases. As for the number of the augmented graphs in each iteration, results show that noisy graphs are often created when n is higher than 30%, because the predictor cannot effectively guide the data augmentation for those labeled graphs whose labels are hard to predict. So, 10% is suggested as the default of top-n%.
Iterative process: Figure 5 investigates the relationship between the quality of augmented graphs and the accuracy of property prediction models. We save a predictor checkpoint every 20 epochs We visualize a few augmented graphs in DCT using ogbg-BACE and O 2 Perm. We adapt top-k pooling (Knyazev et al., 2019) to select the subgraphs that GIN used for prediction. The selected subgraphs are highlighted in green in Figure 6. The three examples show that the augmented graphs can identify and preserve the core structures that GIN uses to predict property values. These augmented graphs are chemically valid, showing that concepts such as some chemical rules from the unlabeled graphs are successfully transferred to downstream tasks. More results are in appendix D.2. Regarding task-specific knowledge, it is known that the fluorine atom and the methyl group are usually negatively and positively correlated to the permeability, respectively (Park et al., 2003;Corrado and Guo, 2020). The augmented examples show that DCT captures this domain knowledge with the task-related objectives. In example (b), DCT replaces most of the fluorine atoms with the methyl groups. It encourages GIN to learn the positive relationship between the methyl group and the permeability so that GIN predicts a high label value. In example (c), DCT replaces the methyl groups with fluorine atoms. It encourages GIN to learn the negative relationship between the fluorine atom and the permeability so that GIN predicts a low label value.
5 Related Work
Graph neural networks (GNNs) (Kipf and Welling, 2017;Xu et al., 2019) are commonly used for graph property prediction in chemistry and polymer informatics tasks (Otsuka et al., 2011;Hu et al., 2020;Zhou et al., 2022). However, it is hard to annotate enough labels in these domains. Recent work used self-supervised tasks such as node attribute prediction and graph structure prediction (Hu et al., 2019;You et al., 2021;Kim et al., 2022) to pre-train architecture-fixed GNNs. Sun et al. (2022) observed that existing methods might fail to transfer knowledge from unlabeled graph data. Flexible GNN architectures for downstream tasks would be desirable.
Graph data augmentation (GDA) methods do not restrict GNN architecture choices to improve prediction accuracy (Trivedi et al., 2022;Zhao et al., 2022Zhao et al., , 2023;;Ding et al., 2022). They learn to create new examples that preserve the properties of original graphs (Liu et al., 2022b,a;Kong et al., 2022;Han et al., 2022;Luo et al., 2022). However, they purely manipulate labeled examples and thus cannot utilize the knowledge in unlabeled graphs. Our DCT combines the knowledge from the unlabeled dataset and the labeled task dataset. It creates label-preserved graph examples with the knowledge transferred from the unlabeled data. It allows the GNN models to have flexible architectures.
Pre-training on self-supervised tasks such as masked image modeling and autoregressive text generation is effective for large language and vision models (Brown et al., 2020;He et al., 2022). However, the hand-crafted self-supervised tasks could hardly help models learn useful knowledge from unlabeled graphs due to the gap between these label-agnostic tasks and the downstream prediction tasks towards drug discovery and material discovery (Sun et al., 2021;Kim et al., 2022;Inae et al., 2023).
A universal self-supervised task to learn from the unlabeled graphs remains under-explored (Sun et al., 2022;Trivedi et al., 2022).
Semi-supervised learning assumes that unlabeled and labeled data are from the same source (Liu et al., 2023). The learning objective in the latent space is usually mutual information maximization that encourages similarity between the representations of unlabeled and labeled graphs (Sun et al., 2020). However, the distributions of the unlabeled and labeled data could be very different due to the different types of sources (Hu et al., 2019), leading to negative impacts on the property prediction on the labeled graphs. Self-training, as a specific type of semi-supervised learning method, selects the unlabeled graphs of confidently predictable labels and assigns pseudo-labels for them (Lee et al., 2013;Iscen et al., 2019). Many studies have explored improving uncertainty estimation (Gal and Ghahramani, 2016;Tagasovska and Lopez-Paz, 2019;Amini et al., 2020) to help the model filter out noise for reliable pseudo-labels. Recently, pseudo-labels have been applied in imbalanced learning (Liu et al., 2023) and representation learning (Ghiasi et al., 2021). However, self-training is restricted to confidently predictable labels and may ignore the huge number of any other unlabeled graphs (Huang et al., 2022). Therefore, it cannot fully utilize the knowledge in the unlabeled graphs.
In contrast, our DCT employs a diffusion model to extract knowledge (as the diffusion and reverse processes) from all the unlabeled graphs. DCT represents the knowledge as task-specific labeled examples to augment the target dataset, instead of uninterpretable pre-trained model parameters.
We note that self-or semi-supervised learning does not conflict with DCT, and we leave their combinations for future work.
Recent works have improved the diffusion models on graphs (Niu et al., 2020;Jo et al., 2022;Vignac et al., 2022;Kong et al., 2023;Chen et al., 2023). EDP-GNN (Niu et al., 2020) employed score matching for permutation-invariant graph data distribution. GDSS (Jo et al., 2022) extended the continuous-time framework [6] to model node-edge joint distribution. DiGress (Vignac et al., 2022) used the transition matrix to preserve the discrete natures of the graph structure. GraphARM (Kong et al., 2023) introduced a node-absorbing autoregressive diffusion process. EDGE (Chen et al., 2023) focused on efficiently generating larger graphs. Instead of improving the generation performance of the diffusion model, our model builds on score-based diffusion models (Jo et al., 2022;Song et al., 2021) for predictive tasks, i.e., , graph classification and graph regression.
In this work, we made the first attempt to transfer minimal sufficient knowledge from unlabeled graphs by data augmentation. We proposed a data-centric framework to use the diffusion model trained on the unlabeled graphs and use two task-related objectives to generate task-specific augmented graphs.
Experiments demonstrated the performance of the proposed framework through visible augmented examples. It is better than self-supervised learning, self-training, and graph data augmentation methods on as many as 15 tasks.
Data Augmentation Data augmentation creates new examples with preserved labels but uses no unlabeled data (Shorten and Khoshgoftaar, 2019;Kashefi and Hwa, 2020;Balestriero et al., 2022). Examples of heuristic data augmentation techniques include flipping, distorting, and rotating images (Shorten and Khoshgoftaar, 2019), using lexical substitution, inserting words, and shuffling sentences in texts (Kashefi and Hwa, 2020), and deleting nodes and dropping edges in graphs (Zhao et al., 2021(Zhao et al., , 2023)). While human knowledge can be used to improve data diversity and reduce over-fitting in heuristic methods, it is difficult to use a single heuristic method to preserve the different labels for different tasks (Balestriero et al., 2022;Cubuk et al., 2019). So, automated augmentation (Cubuk et al., 2019) learned from data to search for the best policy to combine a bunch of predefined heuristic augmentations. Generation models (Antoniou et al., 2017;Bowles et al., 2018;Han et al., 2022) create in-class examples. Other learning ideas such as FATTEN (Liu et al., 2018) and GREA (Liu et al., 2022a) learned to split the latent space for data augmentation. However, learning and augmenting from insufficient labels at the same time may limit the diversity of new examples and cause over-fitting. DCT leverages unlabeled data to avoid them. Learn to Augment (Kong et al.,Liu et al.,etc.) Heuristic Augmentation (Rong et al., etc.) Task w/ Labels Unlabeled Data (Self-Training) Contained Knowledge
Relationship between Data-Centric Approaches As presented in Figure 7, perturb edges, delete nodes and mask attributes (Rong et al., 2019;Trivedi et al., 2022) for graphs are some heuristic ways for data augmentation. The augmented knowledge from them is mainly controlled by human prior knowledge on the perturbations and it often fails to be close to the task, i.e., , random perturbations hardly preserve labels for the augmented graphs. The learning to augment approaches learn from labeled graphs to perturb graph structures (Luo et al., 2022), to estimate graphons for different classes (Han et al., 2022), or to split the latent space for augmentation (Liu et al., 2022a). Although these approaches could preserve labels for the augmented graphs, they introduce less extra knowledge to improve the model prediction. In summary, graph data augmentation is effective in expanding knowledge for limited labels, but it makes no use of unlabeled graphs. Besides, the diversity and richness of the domain knowledge from augmented graphs are far from that contained in a large number of unlabeled graphs. To learn from unlabeled graphs, data-centric approaches like the self-training is assumed to be useful when the unlabeled and labeled data are from the same source. It is less studied when we have a single unified unlabeled source for different tasks.
In Eq. ( 6), we use a leave-one-out variant of InfoNCE (I bound ) to derive the upper bound of mutual information. We summarize the derivation (Poole et al., 2019) here.
The intractable upper bound is minimized when the variational approximation q(G 0 ) matches the true marginal p(G 0 ) (Poole et al., 2019). For each G i , its augmented output G 0 i , and M 1 negative examples with different labels, we could approximate q(G 0
For each molecule and polymer graph, we concatenate the following vectors or values for statistical feature extraction.
• the sum of the degree in the graph;
• the vector indicating the distribution of atom types;
• the vector containing the maximum, minimum and mean values of atoms weights in a molecule or polymer; • the vector containing the maximum, minimum, and mean values of bond valence.
For each protein-protein interaction ego-graph in the biology field, we use the sorted vector of node degree distribution in the graph as the statistical features.
The Lookup Table from Atom Type to Node Embedding Space Given a graph G, we assume the node feature matrix on the graph is X 2 R n⇥Fn , where n is the number of nodes. The edge feature matrix is E 2 R m⇥Fe , where m is the number of edges. There are two ways for G to represent the graph structure in practice. We can use either the dense adjacency matrix A 2 R n⇥n or sparse edge index I e 2 R 2⇥m . The diffusion model (Jo et al., on graphs prefers the former, which is more straightforward for graph generations. The prediction model prefers the latter because of its flexibility, and less computational cost time. The transformation between two types of graph structure representation takes additional time. Particularly for molecular graphs, the node features used for generation (one-hot encoding of the atom type) and for prediction (see the official package of OGBGfoot_0 for details) are different, which introduces extra time to process the graph data. For details, we (1) first need to extract discrete node attributes given the atom type and its neighborhoods; (2) we then need to use an embedding table to embed node attributes in a continuous embedding space;
(3) the embedding features of nodes with their graph structure are inputted into the graph neural networks to get the latent representation for nodes. The reverse process for data augmentation in DCT may need to repeatedly process graph data with steps (1) and (2). It introduces additional time. To address these technical problems, we build up a lookup table to directly map the atom type to the node embedding. We average the node attributes for the same type of node within the batch. We then use the continuous node attributes as weights to average the corresponding node embedding according to the table.
Instantiations of SDE on Graphs According to Song et al. (2021), we use the Variance Exploding (VE) SDE for the diffusion process. Given the minimal noise min and the maximal noise max , the VE SDE is:
The perturbation kernel is derived (Song et al., 2021) as:
On graphs, we follow Jo et al. (2022) to separate the perturbation of adjacency matrix and node features:
(7) having drug side effects of 27 system organ classes in ogbg-SIDER, (8) predicting the property of lipophilicity in ogbg-Lipo, (9) predicting the water solubility (log solubility in mols per litre) from chemical structures in ogbg-ESOL, (10) predicting the hydration free energy of molecules in water in ogbg-FreeSolv. For all molecule datasets, we use the scaffold splitting procedure as the open graph benchmark adopted (Hu et al., 2020). It attempts to separate structurally different molecules into different subsets, which provides a more realistic estimate of model performance in experiments (Wu et al., 2018).
Four polymer regression tasks include GlassTemp, MeltingTemp, ThermCond, and O 2 Perm. They are used to predict different polymer properties such as glass transition temperature ( C), melting temperature ( C), thermal conductivity (W/mK) and oxygen permeability (Barrer). GlassTemp and MeltingTemp are collected from PolyInfo, which is the largest web-based polymer database (Otsuka et al., 2011). The ThermCond dataset is from molecular dynamics simulation and is an extension from the dataset used in (Ma et al., 2022). The O 2 Perm dataset is created from the Membrane Society of Australasia portal, consisting of a variety of gas permeability data (Thornton et al., 2012). Since a polymer is built from repeated units, researchers often use a single unit graph with polymerization points as polymer graphs to predict properties. Different from molecular graphs, two polymerization points are two special nodes (see "⇤" in Figure 2), indicating the polymerization of monomers (Cormack and Elorza, 2004). For all the polymer tasks, we randomly split by 60%/10%/30% for training, validation, and test.
An additional task is protein function prediction using protein-protein interaction graphs (Hu et al., 2019). A node is a protein without attributes, an edge is a relation type between two proteins such as co-expression and co-occurrence. In our DCT, we treat all the relations as the undirected edge without attributes.
When implementing GIN (Xu et al., 2019), we tune its hyper-parameters for different tasks with an early stop on the validation set. We generally implement pre-training baselines following their own setting. For molecule and polymer property prediction and protein function prediction, the pre-trained GIN models with self-supervised tasks such as EDGEPRED, ATTRMASK, CONTEXTPRED in (Hu et al., 2019), INFOMAX (Velickovic et al., 2019) are available. So we directly use them. For other self-supervised methods, we implement their codes with default hyper-parameters. Following their settings, we use 2M ZINC15 (Sterling and Irwin, 2015) to pre-train GIN models for molecule and polymer property prediction. We use 306K unlabeled protein-protein interaction ego-networks (Hu et al., 2019) to pre-train the GIN for the downstream protein function property prediction. For self-training with real unlabeled graphs and INFOGRAPH (Sun et al., 2020), we use 113K QM9 (Ramakrishnan et al., 2014). For self-training with generated unlabeled graphs, we train the diffusion model (Jo et al., 2022) on the real QM9 dataset and then produce the same number of generated unlabeled graphs. To train the diffusion model in our DCT, we also use QM9 (Ramakrishnan et al., 2014).
In Table 3, when we replace the 133K QM9 with the 249K ZINC (Jo et al., 2022) to train the diffusion model, which nearly doubles the size of the unlabeled graphs and includes more atom types, we do not observe any additional improvement, and in some cases, even worse performance. It is possible because of the constraint of the current diffusion model's capacity to model the data distribution for a much larger number of more complex graphs. It encourages powerful generative models in the future, which could be directly used to benefit predictive models under the proposed framework.
In Figure 6, we show through some examples that concepts, such as certain chemical rules from the unlabeled graphs, are successfully transferred to downstream tasks. To further validate this point, we gathered 1,000 task-specific graphs generated in the intermediate steps on the tasks of ogbg-BACE, ogbg-BBBP, ogbg-FreeSolv, and O 2 Perm. We then assessed the chemical validity of these graphs and observed that the validity is 92.8%, 87.9%, 97.4%, and 62.1%, respectively. Results show that transferring knowledge from pre-trained molecular data to target molecules yields relatively high chemical validity. However, the validity drops to 62% when transferring knowledge from pre-trained molecular data to target polymer data. This finding indicates that the transferability of chemical rules becomes more challenging when the distribution gap between the pre-training data and downstream task data is larger.
https://github.com/snap-stanford/ogb/blob/master/ogb/utils/features.py