SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Hyosoon Jang $^{1}$ , Seonghyun Park $^{1}$ , Sangwoo Mo $^{2}$ , Sungsoo Ahn $^{1}$ $^{1}$ POSTECH  $^{2}$ University of Michigan  
{hsjang1205,shpark26,sungsoo.ahn}@postech.ac.kr, swmo@umich.edu
This paper studies structured node classification on graphs, where the predictions should consider dependencies between the node labels. In particular, we focus on solving the problem for partially labeled graphs where it is essential to incorporate the information in the known label for predicting the unknown labels. To address this issue, we propose a novel framework leveraging the diffusion probabilistic model for structured node classification (DPM-SNC). At the heart of our framework is the extraordinary capability of DPM-SNC to (a) learn a joint distribution over the labels with an expressive reverse diffusion process and (b) make predictions conditioned on the known labels utilizing manifold-constrained sampling. Since the DPMs lack training algorithms for partially labeled data, we design a novel training algorithm to apply DPMs, maximizing a new variational lower bound. We also theoretically analyze how DPMs benefit node classification by enhancing the expressive power of GNNs based on proposing AGG-WL, which is strictly more powerful than the classic 1-WL test. We extensively verify the superiority of our DPM-SNC in diverse scenarios, which include not only the transductive setting on partially labeled graphs but also the inductive setting and unlabeled graphs.
In this paper, we address the node classification problem, which is a fundamental problem in machine learning on graphs with various applications, such as social networks [1] and citation networks [2]. One representative example is a transductive problem to classify documents from a partially labeled citation graph. Recently, graph neural networks (GNNs) [3, 4] have shown great success in this problem over their predecessors [5, 6]. Their success stems from the high capacity of non-linear neural architectures combined with message passing.
However, the conventional GNNs are incapable of structured node classification: predicting node labels while considering the node-wise label dependencies [7-9]. That is, given a graph  $G$  with vertices  $\mathcal{V}$  and node labels  $\{y_i:i\in \mathcal{V}\}$ , a GNN with parameter  $\theta$  outputs an unstructured prediction, i.e.,  $p_{\theta}(y_i,y_j|G) = p_{\theta}(y_i|G)p_{\theta}(y_j|G)$  for  $i,j\in \mathcal{V}$ . Especially, this limitation becomes problematic in a transductive setting where the prediction can be improved by incorporating the known labels, e.g., output  $p_{\theta}(y_i|G,y_j)$  for known  $y_{j}$ . In Figures 1(a) and 1(b), we elaborate on this issue with an example where the conventional GNN fails to make the correct prediction.
To resolve this issue, recent studies have investigated combining GNNs with classical structured prediction algorithms, i.e., schemes that consider dependencies between the node labels [8-13]. They combine GNNs with conditional random fields [14], label propagation [15], or iterative classification algorithm [2]. Despite the promising outcomes demonstrated by these studies, their approach relies on the classical algorithms to express joint dependencies between node labels and may lack sufficient expressive power or consistency in incorporating known labels.
(b) GCN [3]
(c) GMNN [7]
Figure 1: Results of various methods to solve node classification on a non-attributed and partially labeled cyclic grid. (a) The task is to estimate  $p(\pmb{y}_U|\pmb{y}_L, G)$  where the known labels  $\pmb{y}_L$  are highlighted in green and the unknown labels  $\pmb{y}_U$  are colored in gray. (b)-(e) Red highlights the incorrect predictions made by the corresponding algorithm. Conventional GNN (GCN) fails to make an informative prediction while our DPM-SNC is the only method to perfectly predict all the labels.
(d) CLGNN [13]
(e) DPM-SNC (ours)
Contribution. We propose a novel framework for structured node classification called DPM-SNC. Our key idea is to leverage the diffusion probabilistic model (DPM) [16], motivated by two significant advantages of DPM for structured node classification: (a) it can effectively learn a joint distribution over both the unknown and the known labels, and (b) it can easily incorporate conditions during inference via posterior sampling. Figure 1 highlights that DPM significantly outperforms previous methods when the prediction problem requires a complete understanding of label dependencies.
However, DPM cannot be directly applied to the transductive scenario, where the model needs to maximize its log-likelihood for partially labeled graphs. We propose a novel training algorithm to address this challenge. Our method maximizes a new variational lower bound of the marginal likelihood of the graph over the unlabeled nodes, involving the alternative optimization of DPM and a variational distribution. In this process, we estimate the variational distribution from the current DPM using a similar approach to fixed-point iteration [17]. After training, DPM can predict the remaining node labels by constraining its predictions with the given labeled nodes, utilizing the manifold-constrained sampling technique [18].
In addition, we provide a theoretical analysis that supports the benefits of DPM-SNC for node classification by enhancing the expressive power of GNNs. To this end, we derive an analog of our DPM-SNC based on the Weisfeiler-Lehman (WL) [19] test, evaluating the ability to distinguish non-isomorphic graphs. Specifically, we prove that DPM-SNC is as powerful as our newly derived aggregated Weisfeiler-Lehman (AGG-WL) test, which is strictly more powerful than the 1-WL test, a known analog to standard GNN architectures [20] such as GCN [3] and GAT [21].
We demonstrate the effectiveness of DPM-SNC on various datasets, covering both transductive and inductive settings. In the transductive setting, we conduct experiments on seven benchmarks: Pubmed, Cora, and Citeseer [22]; Photo and Computer [23]; and Empire and Ratings [24]. Additionally, we introduce a synthetic benchmark to evaluate the ability to capture both short and long-range dependencies. DPM-SNC outperforms the baselines, including both homophilic and heterophilic graphs. In the inductive setting, we conduct experiments on four benchmarks: Pubmed, Cora, Citeseer, and PPI [25]. Furthermore, we evaluate DPM-SNC on the algorithmic reasoning benchmark [26], which requires a higher-level understanding of node relations. DPM-SNC also outperforms previous algorithms for these tasks.
To summarize, our contribution can be listed as follows:
- We explore DPM as an effective solution for structured node classification due to its inherent capability in (a) learning the joint node-wise label dependency in data and (b) making predictions conditioned on partially known data (Section 3).
- We propose a novel method for training a DPM on partially labeled graphs that leverages the probabilistic formulation of DPM to maximize the variational lower bound (Section 4).
- We provide a theoretical analysis of how DPM enhances the expressive power of graph neural networks, strictly improving its capability over the 1-WL test (Section 5).
- Our experiments demonstrate the superiority over the baselines in both transductive and inductive settings. We additionally consider heterophilic graphs and algorithmic reasoning benchmarks to extensively evaluate the capability of DPMs in understanding the label dependencies (Section 6).
Graph neural networks (GNNs). Graph neural networks (GNNs) are neural networks specifically designed to handle graph-structured data. Over the past years, GNNs have shown promising outcomes in various tasks related to graph representation learning [3, 4, 27, 28]. In the case of node classification, GNNs aggregate neighboring node information and generate node representations for label predictions. Their expressive power has been analyzed through the lens of the Weisfeiler-Lehman (WL) test [19], a classical algorithm to distinguish non-isomorphic graphs. For example, Xu et al. [20] proved how the expressive power of conventional GNNs is bounded by the 1-dimensional WL (1-WL) test.
Structured node classification with GNNs. Apart from investigating the power of GNNs to discriminate isomorphic graphs, several studies focused on improving the ability of GNNs to learn the joint dependencies [7, 9, 13, 26, 29-31]. In particular, researchers have considered GNN-based structured node classification [7, 9, 13]. First, Qu et al., [7] parameterize the potential functions of conditional random fields using a GNN and train it using pseudo-likelihood maximization [32]. Next, Qu et al., [9] proposed a training scheme with a GNN-based proxy for the conditional random field. Hang et al., [13] proposed an iterative classification algorithm-based learning framework of GNNs that updates each node-wise prediction by aggregating previous predictions on neighboring nodes.
Diffusion probabilistic models (DPMs). Inspired by non-equilibrium thermodynamics, DPMs are latent variable models (LVMs) that learn to reverse a diffusion process and construct data from a noise distribution. Recent studies demonstrated great success of DPMs in various domains, e.g., generating images [33], text [34], and molecules [35]. An important advantage of diffusion models is their ability to incorporate additional constraints. This allows DPM to solve the posterior inference problems, e.g., conditioning on the partially observed data to recover the original data [18, 36, 37].
3 Diffusion Probabilistic Models for Structured Node Classification
In this section, we introduce the problem of structured node classification and discuss the limitations of graph neural networks (GNNs) in addressing this problem (Section 3.1). We then explain how the diffusion probabilistic models (DPMs) offer a promising solution for this problem (Section 3.2).
3.1 Structured node classification
We address structured node classification, which involves predicting node labels while considering their joint dependencies. Specifically, we focus on predicting labels in partially labeled graphs, where some true labels are known in advance. Our goal is to devise an algorithm that considers the dependencies with known labels to enhance predictions for unknown labels.
To be specific, our problem involves a graph  $G = (\mathcal{V},\mathcal{E},\boldsymbol {x})$  consisting of nodes  $\mathcal{V}$  edges  $\mathcal{E}$  and node attributes  $\pmb {x} = \{x_{i}:i\in \mathcal{V}\}$ . We also denote the node labels by  $\pmb {y} = \{y_{i}:i\in \mathcal{V}\}$ . We let  $\nu_{L}$  and  $\nu_{U}$  denote the set of labeled and unlabeled nodes, while  $\pmb {y}_L$  and  $\pmb {y}_U$  denote the corresponding labels for each set, e.g.,  $\pmb {y}_L = \{y_i:i\in \mathcal{V}_L\}$ . Our objective is to predict the unknown labels  $\pmb {y}_U$  by training on the partially labeled graph, aiming to infer the true conditional distribution  $p(\pmb {y}_U|G,\pmb {y}_L)$ .
Figure 2: The reverse process makes predictions with label dependencies. The node color indicates the label value, and opacity reflects the likelihood.
We point out that the conventional GNNs [3, 4, 27, 28] are suboptimal for estimating  $p(\pmb{y}_U|G,\pmb{y}_L)$  since their predictions on the node labels are independently conditioned on the input, i.e., their predictions are factorized by  $p_{\theta}(\pmb{y}_U|G,\pmb{y}_L) = \prod_{i\in \mathcal{V}_U}p_{\theta}(y_i|G)$ . Namely, conventional GNNs cannot incorporate the information of the known labels  $\pmb{y}_L$  into a prediction of the unknown labels  $\pmb{y}_U$ .
<sup>1</sup>We also show structured node classification to be important for unlabeled graphs in Section 5.
unlabeled node
$(G,y_{\mathbf{L}})$
$(G,y)$
(a) Update the buffer with samples from  $p_{\theta}(\mathbf{y}_U|G,\mathbf{y}_L)$
(b) Update the DPM parameter to maximize  $\mathcal{L}_{\mathrm{VLB}}$
Figure 3: Illustration of DPM-SNC training on partially labeled graphs, repeating steps (a) and (b).
maximize lower bound
$(G,\mathbf{y}^{(t)})$
Algorithm 1 DPM-SNC
1: Train a mean-field GNN  $p_{\phi}(\boldsymbol{y}|G)$ .
2: Initialize the buffer  $\mathcal{B}$  by  $p_{\phi}(\boldsymbol{y}_U|G)$ .
3: repeat
4: for  $i = 1,\dots ,N_{1}$  do
5: Get  $\mathbf{y}_U\sim p_\theta (\mathbf{y}_U|G,\mathbf{y}_L)$  using manifold-constrained sampling.
6: Update  $\mathcal{B} \gets \mathcal{B} \cup \{\pmb{y}_U\}$ .
7: If  $|\mathcal{B}| > K$ , remove oldest one in  $\mathcal{B}$ .
8: end for
9: for  $i = 1,\dots ,N_{2}$  do
10: Sample  $\pmb{y}_U\sim \mathcal{B}$
11: Update  $\theta$  to maximize  $\mathcal{L}_{\mathrm{VLB}}$  with  $G$  and  $\pmb {y} = \pmb {y}_L\cup \pmb {y}_U$
12: end for
13: until converged
Structured prediction algorithms [2, 14] overcome this issue by solving two tasks: (a) learning a joint distribution  $p_{\theta}(\pmb{y}_U, \pmb{y}_L|G)$  to maximize the likelihood of labeled data  $p_{\theta}(\pmb{y}_L|G)$  and (b) inferring from the distribution  $p_{\theta}(\pmb{y}_U|G, \pmb{y}_L)$  conditioned on known labels.
3.2 Diffusion probabilistic models for structured node classification
In this work, we consider diffusion probabilistic models for structured node classification (DPM-SNC) motivated by how their strengths align with the essential tasks for solving structured node classification, outlined in Section 3.1. In particular, DPMs are equipped with (a) high expressivity in learning a joint distribution over data and (b) the ability to easily infer from a posterior distribution conditioned on partially observed data.
To this end, we formally describe DPMs that generate node-wise labels  $\pmb{y}$  associated with a graph  $G$ . At a high level, our DPM consists of two parts: the forward process, which injects noise into the labels, and the reverse process, which denoises to predict the true labels  $\pmb{y}$  [16]. Given the number of diffusion steps  $T$ , the forward process constructs a sequence of noisy labels  $\pmb{y}^{(1:T)} = [\pmb{y}^{(1)}, \dots, \pmb{y}^{(T)}]$  using a fixed distribution  $q(\pmb{y}^{(1:T)}|\pmb{y}^{(0)})$  starting from the true label  $\pmb{y}^{(0)} = \pmb{y}$ . Next, given an initial noises  $\pmb{y}^{(T)}$  sampled from  $p(\pmb{y}^{(T)})$ , the reverse process  $p_{\theta}(\pmb{y}^{(0:T-1)}|\pmb{y}^{(T)}, G)$  is trained to recover the forward process conditioned on the graph  $G$ . We defer the detailed parameterization to Section 4. To be specific, the forward and the reverse process are factorized as follows:
$$
q (\boldsymbol {y} ^ {(1: T)} | \boldsymbol {y} ^ {(0)}) = \prod_ {t = 1} ^ {T} q (\boldsymbol {y} ^ {(t)} | \boldsymbol {y} ^ {(t - 1)}), \quad p _ {\boldsymbol {\theta}} (\boldsymbol {y} ^ {(0: T - 1)} | \boldsymbol {y} ^ {(T)}, G) = \prod_ {t = 1} ^ {T} p _ {\boldsymbol {\theta}} (\boldsymbol {y} ^ {(t - 1)} | \boldsymbol {y} ^ {(t)}, G),
$$
where the training of the reverse process is trivial on a fully-labeled graph  $(\pmb{y}, G)$  by maximizing the variational lower bound of the marginal log-likelihood  $\log p_{\pmb{\theta}}(\pmb{y} | G) = \log \sum_{\pmb{y}^{(1:T)}} p_{\pmb{\theta}}(\pmb{y}^{(0:T)} | G)$  with respect to the noisy labels  $\pmb{y}^{(1:T)}$  sampled from the forward process [16].
By leveraging shared latent variables  $\pmb{y}^{(t)}$  across multiple steps in the reverse process, the reverse process effectively considers the dependencies in the output. Next, DPMs can easily infer from a posterior distribution conditioned on partially observed data. Specifically, the incremental updating of  $\pmb{y}^{(t)}$  for  $t = 1,\dots ,T$  allows the DPM to incorporate the known label  $\pmb{y}_L$ , e.g., applying iterative projection with manifold-based correction [18]. This enables inference from the distribution  $p_{\theta}(\pmb{y}_U|G,\pmb{y}_L)$  conditioned on the known labels  $\pmb{y}_L$ . See Figure 2 for an illustration of DPM for structured node classification with known labels.
2We relax the discrete labels to one-hot vectors to apply the forward process.
4 Training Diffusion Probabilistic Models on Partially Labeled Graphs
Despite the potential of DPMs for structured node classification, they lack the algorithms to learn from partially labeled graphs, i.e., maximize the likelihood of known labels. To resolve this issue, we introduce a novel training algorithm for DPM-SNC, based on maximizing a variational lower bound for the log-likelihood of known labels.
Variational lower bound. At a high-level, our algorithm trains the DPM to maximize the log-likelihood of training data  $\log p_{\theta}(\pmb{y}_L|G)$ , which is defined as follows:
$$
\mathcal {L} = \log p _ {\theta} (\boldsymbol {y} _ {L} | G) = \log \sum_ {\boldsymbol {y} _ {U}} \sum_ {\boldsymbol {y} ^ {(1: T)}} p _ {\boldsymbol {\theta}} (\boldsymbol {y} _ {L}, \boldsymbol {y} _ {U}, \boldsymbol {y} ^ {(1: T)} | G).
$$
However, this likelihood is intractable due to the exponentially large number of possible combinations for the unknown labels  $\pmb{y}_U$  and the noisy sequence of labels  $\pmb{y}^{(1:T)}$ . To address this issue, we train the DPM based on a new variational lower bound  $\mathcal{L} \geq \mathcal{L}_{\mathrm{VLB}}$ , expressed as follows:
$$
\mathcal {L} _ {\mathrm {V L B}} = \mathbb {E} _ {q (\boldsymbol {y} _ {U} | \boldsymbol {y} _ {L})} \left[ \mathbb {E} _ {q (\boldsymbol {y} ^ {(1: T)} | \boldsymbol {y})} \left[ \log p _ {\boldsymbol {\theta}} (\boldsymbol {y}, \boldsymbol {y} ^ {(1: T)} | G) - \log q (\boldsymbol {y} ^ {(1: T)} | \boldsymbol {y}) \right] - \log q (\boldsymbol {y} _ {U} | \boldsymbol {y} _ {L}) \right]. \tag {1}
$$
Here,  $q(\cdot)$  is a variational distribution factorized by  $q(\pmb{y}_U, \pmb{y}^{(1:T)}|\pmb{y}_L) = q(\pmb{y}^{(1:T)}|\pmb{y})q(\pmb{y}_U|\pmb{y}_L)$ , where  $\pmb{y} = \pmb{y}_U \cup \pmb{y}_L$ . We provide detailed derivation in Appendix A.1.
Parameterization. We define the reverse process  $p_{\theta}(\pmb{y}^{(t-1)}|\pmb{y}^{(t)}, G)$  with a Gaussian distribution, where the mean value parameterized by a graph neural network (GNN) and the variance set to be a hyper-parameter. Following the prior work [16], we also employ a Gaussian diffusion to parameterize the forward process  $q(\pmb{y}^{(1:T)}|\pmb{y})$  as follows:
$$
q \left(\boldsymbol {y} ^ {(1)}, \dots , \boldsymbol {y} ^ {(T)} \mid \boldsymbol {y} ^ {(0)}\right) = \prod_ {t = 1} ^ {T} \mathcal {N} \left(\boldsymbol {y} ^ {(t)}; \sqrt {1 - \beta_ {t}} \boldsymbol {y} ^ {(t - 1)}, \beta_ {t} \mathbf {I}\right),
$$
where  $\mathbf{I}$  is an identity matrix,  $\beta_{1},\ldots ,\beta_{T}$  are fixed variance schedules. We set the variance schedule to promote  $q(\pmb{y}^{(T)}|\pmb{y}^{(0)})\approx \mathcal{N}(\pmb{y}^{(T)};\pmb {0},\pmb {I})$  by setting  $\beta_{t} < \beta_{t + 1}$  for  $t = 0,\dots ,T - 1$  and  $\beta_T = 1$ . Finally, we describe the variational distribution  $q(\pmb{y}_U|\pmb{y}_L)$  as an empirical distribution over a fixed-size buffer  $\mathcal{B}$  containing multiple estimates of the unknown labels  $\pmb{y}_U$ . This buffer is updated throughout the training. We provide more details on our parameterization in Appendix A.2.
Training algorithm. To maximize the variational lower bound, we alternatively update the GNN-based reverse process  $p_{\theta}(\pmb{y}, \pmb{y}^{(1:T)}|G)$  and the buffer-based variational distribution  $q(\pmb{y}_U|\pmb{y}_L)$ . In particular, we update the parameters  $\theta$  of the reverse process  $p_{\theta}(\pmb{y}, \pmb{y}^{(1:T)}|G)$  to maximize the Monte Carlo approximation of  $\mathcal{L}_{\mathrm{LVB}}$  by applying ancestral sampling to the variational distribution  $q(\pmb{y}^{(1:T)}|\pmb{y})q(\pmb{y}_U|\pmb{y}_L)$ , i.e., sampling  $\pmb{y}_U$  from the buffer  $\mathcal{B}$  and applying the diffusion process to  $\pmb{y}$ . This is the same as the original denoising-based training of DPM [16] with estimated unknown labels, where the detailed training objective is described in Appendix A.3.
Next, we update the variational distribution  $q(\pmb{y}_U|\pmb{y}_L)$  by inserting samples from the distribution  $p_{\theta}(\pmb{y}_U|G,\pmb{y}_L)$  into the buffer  $\mathcal{B}$ . This update is derived from the condition  $q(\pmb{y}_U|\pmb{y}_L) = p_{\theta}(\pmb{y}_U|G,\pmb{y}_L)$  being necessary to maximize  $\mathcal{L}_{\mathrm{VLB}}$ , similar to the derivation of fixed-point iteration for optimization [17]. We describe the overall optimization procedure in Figure 3 and Algorithm 1.
Finally, we emphasize that our training algorithm is specialized for DPMs. Previous studies introducing variational lower bounds for structured node classifications [7, 38] face intractability in maximizing  $\log p_{\theta}(\pmb{y}|G)$  or sampling from  $p_{\theta}(\pmb{y}_U|G,\pmb{y}_L)$ . They require defining the pseudolikelihood for the former [32], or parameterizing variational distribution for the latter. However, in our case, the formal simply requires maximizing the lower bound  $\mathcal{L}_{\mathrm{VLB}}$ , and the latter is easily solved by manifold-constrained sampling [18].
<sup>3</sup>We use manifold-constrained sampling of DPM to infer from the distribution  $p_{\theta}(\pmb{y}_U|G, \pmb{y}_L)$  [18]. The detailed sampling procedure is described in Appendix A.4.
4In practice, we initialize the buffer  $\mathcal{B}$  with samples from a mean-field GNN  $p_{\phi}(\pmb {y}|G)$ , which outputs an independent joint distribution over node labels, i.e.,  $p_{\theta}(\pmb{y}_U|G,\pmb{y}_L) = \prod_{i\in \mathcal{V}_U}p_{\theta}(\pmb{y}_i|G)$ .
In this section, we present a secondary motivation for using DPMs in node classification, distinct from the one given in Section 3 and Section 4. Specifically, we demonstrate that DPMs provably enhance the expressive power of conventional GNNs for solving the graph isomorphism test, implying improved expressive power for node classification problems as well.
To this end, we assess the expressive power of our DPM-SNC by its analog for discriminating isomorphic graphs, which we call the aggregated Weisfeiler-Lehman (AGG-WL) test. Then, we show that AGG-WL is strictly more powerful than the 1-dimensional WL (1-WL) test [19], which is an analog of GNNs [20]. We formalize our results in the following theorem.
Theorem 1. Let 1-WL-GNN be a GNN as powerful as the 1-WL test. Then, DPM-SNC using a 1-WL-GNN is strictly more powerful than the 1-WL-GNN in distinguishing non-isomorphic graphs.
We provide the formal proof in Appendix B. At a high level, both the 1-WL test and the AGG-WL test assign colors to nodes and iteratively refine them based on the colors of neighboring nodes, enabling the comparison of graph structures. The key difference is that AGG-WL applies perturbations into the initial graph to create multiple augmented graphs, and aggregates the results of the 1-WL algorithm on each augmented graph. The perturbations and aggregation can be done by node-wise concatenations with binary random features and hashing a multiset of the refined graph colors, respectively. Hence, any graph pair distinguishable by the 1-WL test is also distinguishable by our AGG-WL test, and there exist pairs of graphs indistinguishable by the WL test but distinguishable by our variant.
We remark that our analysis can be easily extended to applying other latent variable models, e.g., variational auto-encoders [39] or normalizing flows [40], for node classification. In this analysis, sampling a latent variable corresponds to fixing a perturbation for color assignments at initialization, and aggregating over the perturbed graphs corresponds to computing the prediction over the true labels with marginalization over the latent variables. Our theory is based on that of Bevilaqua et al. [41] which was originally developed for analyzing a particular GNN architecture; our key contribution lies in extending this analysis to the latent variable models.
6.1 Transductive setting
We first evaluate the performance of our algorithm in the transductive setting. Specifically, we conduct an evaluation of DPM-SNC on synthetic data, as well as real-world node classification problems. We consider different types of label dependencies by considering both homophilic and heterophilic graphs. We provide the details of our implementation in Appendix C.1.