SlowGuess/NeurIPS_2023_OCR huggingface dataset:
The Lookup Table from Atom Type to Node Embedding Space Given a graph  $G$ , we assume the node feature matrix on the graph is  $\mathbf{X} \in \mathbb{R}^{n \times F_{\mathrm{n}}}$ , where  $n$  is the number of nodes. The edge feature matrix is  $\mathbf{E} \in \mathbb{R}^{m \times F_{\mathrm{e}}}$ , 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  $\mathbf{A} \in \mathbb{R}^{n \times n}$  or sparse edge index  $\mathbf{I}_{e} \in \mathbb{R}^{2 \times m}$ . The diffusion model (Jo et al., 2022) 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 and 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 OGBG<sup>1</sup> 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  $\sigma_{\mathrm{min}}$  and the maximal noise  $\sigma_{\mathrm{max}}$ , the VE SDE is:
$$
\mathrm {d} G = \sigma_ {\min } \left(\frac {\sigma_ {\max }}{\sigma_ {\min }}\right) ^ {t} \sqrt {2 \log \frac {\sigma_ {\max }}{\sigma_ {\min }}} \mathrm {d} \mathbf {w}, \quad t \in (0, 1 ] \tag {11}
$$
The perturbation kernel is derived (Song et al., 2021) as:
$$
p _ {0 t} \left(G ^ {(t)} \mid G ^ {(0)}\right) = \mathcal {N} \left(G ^ {(t)}; G ^ {(0)}, \sigma_ {\min } ^ {2} \left(\frac {\sigma_ {\max }}{\sigma_ {\min }}\right) ^ {2 t} \mathbf {I}\right), \quad t \in (0, 1 ] \tag {12}
$$
On graphs, we follow Jo et al. (2022) to separate the perturbation of adjacency matrix and node features:
$$
p _ {0 t} \left(G ^ {(t)} \mid G ^ {(0)}\right) = p _ {0 t} \left(\mathbf {A} ^ {(t)} \mid \mathbf {A} ^ {(0)}\right) p _ {0 t} \left(\mathbf {X} ^ {(t)} \mid \mathbf {X} ^ {(0)}\right). \tag {13}
$$
<https://github.com/snap-stanford/ogb/blob/master/ogb/util/features.py>
The Sampling Algorithm in the Reverse Process for Graph Data Augmentation We adapt the Predictor-Corrector (PC) samplers for the graph data augmentation in the reverse process. The algorithm is shown in Algorithm 1.
Algorithm 1 Diffusion-Based Graph Augmentation with PC Sampling
Input: Graph  $G$  with node feature  $\mathbf{X}$  and adjacency matrix  $\mathbf{A}$ , the denoising function for node feature  $\mathbf{s}_{\mathbf{X}}$  and adjacency matrix  $\mathbf{s}_{\mathbf{A}}$ , the fine-tune loss  $\mathcal{L}_{\mathrm{aug}}$ , Lagevin MCMC step size  $\beta$ , scaling coefficient  $\epsilon_{1}$ $\mathbf{A}^{(D)} \leftarrow \mathbf{A} + \mathbf{z}_{A}$ ;  $\mathbf{z}_{A} \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$ $\mathbf{X}^{(D)} \leftarrow \mathbf{X} + \mathbf{z}_{X}$ ;  $\mathbf{z}_{X} \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$   
for  $t = D - 1$  to 0 do  
 $\hat{G}_{(t+1)} \sim p_{0t+1}(\hat{G}_{(t+1)}|G^{(t+1)})$  {inner-loop sampling with another PC sampler}  
 $\mathbf{S}_{A} = \frac{1}{2}\mathbf{s}_{\mathbf{A}}(G^{(t+1)}, t+1) - \frac{1}{2}\alpha\nabla_{\mathbf{A}^{(t)}}\mathcal{L}_{\mathrm{aug}}(\hat{G}_{(t+1)})$ $\mathbf{S}_{X} = \frac{1}{2}\mathbf{s}_{\mathbf{X}}(G^{(t+1)}, t+1) - \frac{1}{2}\alpha\nabla_{\mathbf{X}^{(t)}}\mathcal{L}_{\mathrm{aug}}(\hat{G}_{(t+1)})$ $\tilde{\mathbf{A}}^{(t)} \leftarrow \mathbf{A}^{(t+1)} + g(t)^2\mathbf{S}_{A} + g(t)\mathbf{z}_{A}$ ;  $\mathbf{z}_{A} \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$  {Predictor for adjacency matrix}  
 $\tilde{\mathbf{X}}^{(t)} \leftarrow \mathbf{X}^{(t+1)} + g(t)^2\mathbf{S}_{X} + g(t)\mathbf{z}_{X}$ ;  $\mathbf{z}_{X} \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$  {Predictor for node features}  
 $\mathbf{A}^{(t)} \leftarrow \tilde{\mathbf{A}}^{(t)} + \frac{\beta}{2}\mathbf{S}_{A} + \epsilon_{1}\sqrt{\beta}\mathbf{z}_{A}$ ;  $\mathbf{z}_{A} \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$  {Corrector for adjacency matrix}  
 $\mathbf{X}^{(t)} \leftarrow \tilde{\mathbf{X}}^{(t)} + \frac{\beta}{2}\mathbf{S}_{X} + \epsilon_{1}\sqrt{\beta}\mathbf{z}_{X}$ ;  $\mathbf{z}_{X} \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$  {Corrector for node features}  
end for  
return  $G' = (\mathbf{A}^{(0)}, \mathbf{X}^{(0)})$
The Algorithm of the Framework The algorithm of the proposed data-centric knowledge transfer framework is presented in Algorithm 2 and Algorithm 3. In detail, Algorithm 2 corresponds to Section 3.2 and Algorithm 3 corresponds to Section 3.3.
Algorithm 2 The Data-Centric Knowledge Trans-Algorithm 3 The Data-Centric Knowledge Transfer Framework: Learning from Unlabeled Graphs Framework: Generating Task-specific Labeled Graphs
Input: Given unlabeled graphs from the space  $\mathcal{G}^{[U]}$ , randomly initialized score models  $\mathbf{s}_{\mathbf{X}}$  and  $\mathbf{s}_{\mathbf{A}}$  for node feature and graph adjacency matrix, respectively, the total diffusion time step  $T$ .  
while  $\mathbf{s}_{\mathbf{X}}$  and  $\mathbf{s}_{\mathbf{A}}$  not converged do  
Sample  $G = (\mathbf{X},\mathbf{A})\in \mathcal{G}^{[U]}$   
Sample  $t\in \mathrm{Uniform}(1,2,\dots,T)$   
Sample  $\mathbf{z}_A\sim \mathcal{N}(\mathbf{0},\mathbf{I})$   
Sample  $\mathbf{z}_X\sim \mathcal{N}(\mathbf{0},\mathbf{I})$   
Sample  $\hat{G}$  with  $t,\mathbf{z}_A,\mathbf{z}_X$  and Eq. (13)  
Optimize  $\mathbf{s}_{\mathbf{A}}$  with the gradient:  
 $\nabla \| \mathbf{z}_A - \mathbf{s}_{\mathbf{A}}(\hat{G},t)\|^2$   
Optimize  $\mathbf{s}_{\mathbf{X}}$  with the gradient:  
 $\nabla \| \mathbf{z}_X - \mathbf{s}_{\mathbf{X}}(\hat{G},t)\|^2$   
end while
C Additional Experiments Set-ups
We perform experiments on 15 datasets, including eight classification and seven regression tasks from chemistry, material science, and biology. We use Area under the ROC curve (AUC) for classification performance and mean absolute error (MAE) for regression.
C.1 Molecule Classification and Regression Tasks
Seven molecule classification and three molecule regression tasks are from open graph benchmark (Hu et al., 2020). They were originally collected by MoleculeNet (Wu et al., 2018) and used to predict molecule properties. They include (1) inhibition to HIV virus replication in ogbg-HIV, (2) toxicological properties of 617 types in ogbg-ToxCast, (3) toxicity measurements such as nuclear receptors and stress response in ogbg-Tox21, (4) blood-brain barrier permeability in ogbg-BBBP, (5) inhibition to human  $\beta$ -secretase 1 in ogbg-BACE, (6) FDA approval status or failed clinical trial in ogbg-ClinTox,
(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).
C.2 Polymer Regression Tasks
Four polymer regression tasks include GlassTemp, MeltingTemp, ThermCond, and  $\mathrm{O}_2\mathrm{Perm}$ . They are used to predict different polymer properties such as glass transition temperature  $(^{\circ}\mathrm{C})$ , melting temperature  $(^{\circ}\mathrm{C})$ , thermal conductivity  $(\mathrm{W} / \mathrm{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  $\mathrm{O}_2\mathrm{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.
C.3 Protein Classification Task
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.
C.4 Baselines and Implementation
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).
D Additional Experiment Analysis
D.1 The Power of Diffusion Model to Learn from Unlabeled Graphs
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.
D.2 Chemical Validity of the Generated Graphs in Downstream Tasks
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  $\mathrm{O}_2\mathrm{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.