SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yejiang Wang $^{1,2}$  Yuhai Zhao $^{1,2,\dagger}$  Zhengkui Wang $^{3}$  Ling Li $^{1,2}$
<sup>1</sup> School of Computer Science and Engineering, Northeastern University, China
$^{2}$  Key Laboratory of Intelligent Computing in Medical Image of Ministry of Education, Northeastern University, China
<sup>3</sup> InfoComm Technology Cluster, Singapore Institute of Technology, Singapore wangyejiang@stumail.neu.edu.cn, zhaoyuhai@mail.neu.edu.cn, zhengkui.wang@singaporetech.edu.sg, lilingneu@gmail.com
Self-supervised learning on graphs aims to learn graph representations in an unsupervised manner. While graph contrastive learning (GCL - relying on graph augmentation for creating perturbation views of anchor graphs and maximizing/minimizing similarity for positive/negative pairs) is a popular self-supervised method, it faces challenges in finding label-invariant augmented graphs and determining the exact extent of similarity between sample pairs to be achieved. In this work, we propose an alternative self-supervised solution that (i) goes beyond the label invariance assumption without distinguishing between positive/negative samples, (ii) can calibrate the encoder for preserving not only the structural information inside the graph, but the matching information between different graphs, (iii) learns isometric embeddings that preserve the distance between graphs, a by-product of our objective. Motivated by optimal transport theory, this scheme relies on an observation that the optimal transport plans between node representations at the output space, which measure the matching probability between two distributions, should be consistent with the plans between the corresponding graphs at the input space. The experimental findings include: (i) The plan alignment strategy significantly outperforms the counterpart using the transport distance; (ii) The proposed model shows superior performance using only node attributes as calibration signals, without relying on edge information; (iii) Our model maintains robust results even under high perturbation rates; (iv) Extensive experiments on various benchmarks validate the effectiveness of the proposed method.
Self-supervised graph learning involves learning representations of real-world graph data without the need for human supervision. Graph contrastive learning (GCL) [6, 16, 71] has been identified as one of the most successful graph self-supervised learning approaches, with its key components consisting of graph augmentation and contrastive learning. The former creates perturbation views of anchor graphs via various augmenting techniques, while the latter maximizes the similarity for two augmentations (positive pairs) of the same anchor and minimizes the similarity for those of two different anchors (negative pairs). The effectiveness of contrastive learning is dependent on the assumption that the augmented operations preserve the nature of data points and ensure that the augmented samples have consistent labels with the original ones.
However, graph data structures are discrete and their properties may vary significantly even with slight perturbations, which makes it much more challenging to design reasonable augmentations that guarantee the label-invariant assumption for graphs, in contrast to images or text. Additionally, the concept of "maximum similarity" in contrastive learning is difficult to measure since it is vague and lacks a clear indication of how much similarity should be maximized (or minimized) for a given pair of positive (or negative) views.
To address these challenges, an intuitive solution is to introduce the concept of distance from the input space, whereby the distance between the learned embeddings is forced to be equal to the distance between the corresponding input graphs. However, this requirement is challenging to achieve as the input objects (graphs) and output representations (vectors) are two distinct concepts, making it difficult to agree on their distance metrics. For instance, comparing the graph edit distance between graphs [15] and the Euclidean distance between vectors directly is inconceivable.
In response to the aforementioned challenges, in this paper, we propose a novel self-supervised learning method that seeks to align optimal transport plans [59] from graph space to node representation space, instead of transport distance [44]. This method is referred to as Graph Transport Learning (GTL), and it exploits the key concept in optimal transport (OT) theory, which aims to identify an optimal match (i.e., transport plan) between two data distributions (e.g., node sets of graphs) that minimizes the matching cost (i.e., transport distance). As shown in Figure 1, the transportation plan  $\pi$  explicitly determines how to match particles from the source distribution  $\mu$  to the target distribution  $\nu$ . Our approach involves several key steps. First, we obtain two graph views by augmenting a graph and generating node embeddings using a backbone model, such as Graph Neural Network (GNNs) [25]. Two optimal transport plans can be computed from the graph space and the representation space, respectively. To accurately capture the matching relationships in the original graph space, we enforce the backbone to learn representations that exhibit consistency with the optimal transport plans between the corresponding source graphs in the input space. This is
achieved by minimizing the discrepancy between the two plans. Unlike distance-based approaches, the transport plan alignment shows several advantages: (i) Direct comparability: Plans with the same dimension can be compared directly, regardless of differences between the graph and representation spaces. The value of the transport plan is dimensionless, representing the joint probability of two particles; (ii) Accurate match relationship: For discrete objects, the transport plan retains more accurate matching relation between data compared to distance; (iii) Label-variant: Importantly, our self-supervised model does not require differentiating between positive and negative samples after graph augmentation. This is due to the availability of corrective information from the input space, which eliminates the need for a label-invariant assumption. Notably, our experimental findings in Section 7 demonstrate that our model maintains robust performance even under high perturbation rates, such as when  $80\%$  of both edges and node attributes are destroyed during graph augmentation. In summary, we make the following contributions:
Figure 1: An illustration of discrete optimal transport plan.
- We propose a novel paradigm for self-supervised graph learning based on optimal plan alignment, GALOPA, which offers a distinct objective compared to contrastive learning methods. This approach eliminates the need for a label-invariant assumption.
- By constraining the discrepancy between the transport plans, we introduce a new loss to enable the sharing of the exact matching information from the graphs space to the representation space.
- Multiple comprehensive experiments, including distance v.s plan, node feature v.s edge, robustness test and comparison with state-of-the-art methods demonstrate remarkable performance of GALOPA.
Self-supervised Learning on Graphs. Graph self-supervised learning has been a promising paradigm for learning useful representations of unlabeled graph data. The node embedding methods aim to learn a low-dimensional vector embedding of vertex preserving different kinds of order proximity via factorization [3], random walk [17] or deep learning [8, 61]. Recently, graph contrastive learning has achieved unprecedented success [6, 28, 55, 64], where the goal is to ensure representations have high similarity between positive views of a graph and high dissimilarity between negative views. A common way of contrastive objective is contrasting views at the node level. For the
representations (positive views)  $\mathbf{h}_i^1$  and  $\mathbf{h}_i^2$  of same node  $i$  in two augmented graphs  $\mathcal{G}_1$  and  $\mathcal{G}_2$ , the pairwise contrastive loss can be defined as
$$
\mathcal {L} _ {\text {c o n t r a s t}} = \sum_ {i} \log \left(\frac {e ^ {\mathcal {D} \left(\mathbf {h} _ {i} ^ {1} , \mathbf {h} _ {i} ^ {2}\right) / \tau}}{e ^ {\mathcal {D} \left(\mathbf {h} _ {i} ^ {1} , \mathbf {h} _ {i} ^ {2}\right) / \tau} + \sum_ {k \neq i} e ^ {\mathcal {D} \left(\mathbf {h} _ {i} ^ {1} , \mathbf {h} _ {k} ^ {2}\right) / \tau}}\right) \tag {1}
$$
where  $\mathcal{D}$  is a discriminator that maps two views to an agreement (similarity) score, such as the inner product.  $\tau$  denotes the temperature. Obviously, this type of loss strongly relies on the label invariance assumption. In other words, it needs to know (or assume) beforehand that the two views are positive/negative samples, which is challenging for discrete graph structures. For example, there exist some graphs, such as molecular graphs, whose labels are very sensitive to perturbation/corruption. Some studies have explored the label-invariant augmentation [29, 31, 74], but such augmentations require very careful design and adjustment, and sometimes limit the power of graph augmentation. More recently, [22] utilizes the graph edit distance to train graph discriminator to predict whether a graph is an original graph or a perturbed one. However, this model still requires positive and negative samples and relies on the assumption of label invariance.
Optimal Transport. Optimal transport (OT) [59] is a mathematical tool for aligning probability distributions has received much attention in the machine learning community in many applications, e.g., computer vision [4, 51], generative adversarial network [2, 7, 30], domain adaptation [18]. OT aims to derive a transport plan between a source and a target distribution, such that the transport cost is minimized. [34] proposes the Gromov-Wasserstein (GW) distance between metrics defined within each space rather than between samples across different spaces, which has been used as a distance between graphs in several applications [56, 60, 67]. However, these studies can only perform graph classification on datasets with multiple graphs and cannot be applied to network analysis, such as node representation classification, where only one network is available. In this work, we combine optimal transport problem with graph neural network to form a novel self-supervised learning paradigm that allows both graph and network learning. See Appendix A for more details.
Notations. Let  $\mathcal{G} = (\mathcal{V},\mathbf{A},\mu)$  be a graph of  $n$  nodes where  $\mathcal{V}$  denotes the set of nodes,  $\mathbf{A} \in \{0,1\}^{n\times n}$  is the adjacency matrix.  $\mu \in \mathbb{R}^n$  is the empirical distribution of nodes in  $\mathcal{G}$ . Generally,  $\mu$  is a uniform discrete probability distribution. When the node label (or attribute) is available, we represent the node feature matrix as  $\mathbf{X} \in \mathbb{R}^{n\times d}$ , where  $d$  denotes the dimension of node attributes.  $\mathbf{X}^i$  represents the  $i$ -row of  $\mathbf{X}$ . Let  $[\boldsymbol{n}] = \{1,\dots ,n\}$ ,  $[\boldsymbol{n}]^2 = [\boldsymbol{n}] \times [\boldsymbol{n}]$  and  $\langle \cdot ,\cdot \rangle$  denotes the inner product for matrices. We denote  $\otimes$  the tensor-matrix multiplication.  $\mathbf{1}_n$  represents the vector with ones as all the  $n$  elements.  $|\cdot |$  denotes absolute value.
Plan and Optimal Transport. The optimal transport (OT) problem is pioneered by Monge [35] in order to seek the most cost-effective transport plan that transforms the mass of a pile of sand into another one. In particular, it studies how to find an optimal coupling or optimal plan  $\pi$  for transforming the distribution  $\mu$  to  $\nu$  with minimum total transport cost (i.e., optimal transport distance), where the element of  $\pi$  describes the probability of moving mass from one position to another. In this work, we mainly focus on the discrete case. Given two sets of features  $X_{1} = \{X_{1}^{i}\}_{i = 1}^{n}$  and  $X_{2} = \{X_{2}^{j}\}_{j = 1}^{m}$  where  $n$  and  $m$  are the number of features, respectively.  $\mu \in \mathbb{R}^n$  and  $\nu \in \mathbb{R}^m$  are the probability distributions of the entities in the two sets, respectively. The formulation of the OT distance is
$$
\mathcal {W} \left(\boldsymbol {X} _ {1}, \boldsymbol {X} _ {2}\right) = \min  _ {\pi \in \Pi (\mu , \nu)} \sum_ {i \in [ n ]} \sum_ {j \in [ m ]} c _ {\mathcal {X}} \left(\boldsymbol {X} _ {1} ^ {i}, \boldsymbol {X} _ {2} ^ {j}\right) \cdot \pi_ {i j} = \min  _ {\pi \in \Pi (\mu , \nu)} \left\langle \boldsymbol {K} \left(\boldsymbol {X} _ {1}, \boldsymbol {X} _ {2}\right), \pi \right\rangle \tag {2}
$$
$$
\Pi (\mu , \nu) = \left\{\pi \in \mathbb {R} ^ {n \times m} \mid \pi \mathbf {1} _ {m} = \mu , \mathbf {1} _ {n} \pi = \nu \right\} \tag {3}
$$
denotes all the joint distributions  $\pi$  with marginals  $\mu$  and  $\nu$ .  $K(X_{1},X_{2})_{ij} = c_{\mathcal{X}}(X_1^i,X_2^j)$  is the cost (work) of moving  $X_{1}^{i}$  to  $X_{2}^{j}$ , the cosine distance between  $X_{1}^{i}$  and  $X_{2}^{i}$  is a popular choice. The  $\pi \in \mathbb{R}^{n\times m}$  is called as transport plan. This distance is also known as the Wasserstein distance.
Optimal Transport for Graphs. The Wasserstein problem requires the two distributions of point sets to lie in the same space. But for graphs, it is difficult to measure the cost between two nodes on
different graphs without node label (attribute). Even if the cost between nodes could be calculated, the Wasserstein distance cannot take the edge information into account. To compare distributions that are not necessarily in the same space, [34] defines Gromov-Wasserstein distance between two graphs without node label  $\mathcal{G}_1 = (\pmb{A}_1,\mu)$  and  $\mathcal{G}_2 = (\pmb{A}_2,\nu)$  as follow
$$
\mathcal {W} _ {\mathrm {G W}} \left(\mathcal {G} _ {1}, \mathcal {G} _ {2}\right) = \min  _ {\pi_ {\mathcal {G}} \in \Pi (\mu , \nu)} \sum_ {i, k \in [ n ] ^ {2}} \sum_ {j, l \in [ m ] ^ {2}} c _ {\mathcal {A}} \left(\boldsymbol {A} _ {1} ^ {i k}, \boldsymbol {A} _ {2} ^ {j l}\right) \cdot \pi_ {i k} \pi_ {j l} = \min  _ {\pi \in \Pi (\mu , \nu)} \langle \boldsymbol {L} \left(\boldsymbol {A} _ {1}, \boldsymbol {A} _ {2}\right) \otimes \pi , \pi \rangle \tag {4}
$$
where  $\pmb{L}(\pmb{A}_1, \pmb{A}_2)$  is 4-D tensor and  $\pmb{L}(\pmb{A}_1, \pmb{A}_2)_{ijkl} = c_A (\pmb{A}_1^{ik}, \pmb{A}_2^{jl})$ . The cost function  $c_{\mathcal{A}}$  is commonly defined as  $c_{\mathcal{A}} (\pmb{A}_1^{ik}, \pmb{A}_2^{jl}) = |\pmb{A}_1^{ik} - \pmb{A}_2^{jl}|$ .
Consider two graphs with node attributes  $\mathcal{G}_1 = (A_1, X_1, \mu)$  and  $\mathcal{G}_2 = (A_2, X_2, \nu)$  where  $A_1 \in \mathbb{R}^{n \times n}$  and  $A_2 \in \mathbb{R}^{m \times m}$  denote their adjacency matrices,  $X_1 \in \mathbb{R}^{n \times d}$  and  $X_2 \in \mathbb{R}^{m \times d}$  are feature matrices. The fused Gromov-Wasserstein distance [56] between graphs  $\mathcal{G}_1$  and  $\mathcal{G}_2$  can be defined as
$$
\mathcal {W} _ {\mathrm {F G W}} \left(\mathcal {G} _ {1}, \mathcal {G} _ {2}\right) = \min  _ {\pi_ {\mathcal {G}} \in \Pi (\mu , \nu)} \sigma \sum_ {i j} c _ {\mathcal {X}} \left(\boldsymbol {X} _ {1} ^ {i}, \boldsymbol {X} _ {2} ^ {j}\right) \cdot \pi_ {i j} ^ {\mathcal {G}} + (1 - \sigma) \sum_ {i j k l} c _ {\mathcal {A}} \left(\boldsymbol {A} _ {1} ^ {i k}, \boldsymbol {A} _ {2} ^ {j l}\right) \cdot \pi_ {i j} ^ {\mathcal {G}} \pi_ {k l} ^ {\mathcal {G}} \tag {5}
$$
$$
\min  _ {\pi_ {\mathcal {G}} \in \Pi (\mu , \nu)} \left\langle \sigma \boldsymbol {K} \left(\boldsymbol {X} _ {1}, \boldsymbol {X} _ {2}\right) + (1 - \sigma) \boldsymbol {L} \left(\boldsymbol {A} _ {1}, \boldsymbol {A} _ {2}\right) \otimes \pi_ {\mathcal {G}}, \pi_ {\mathcal {G}} \right\rangle \tag {6}
$$
where  $\pmb{K}(\pmb{X}_1, \pmb{X}_2)_{ij} = c_{\mathcal{X}}(\pmb{X}_1^i, \pmb{X}_2^j)$ ,  $\pmb{L}(\pmb{A}_1, \pmb{A}_2)_{ijkl} = c_{\mathcal{A}}(A_1^{ik}, A_2^{jl})$  and  $\sigma \in [0; 1]$  denotes a trade-off parameter. Obviously, this definition is a fusion of Equations (2) and (4).
4 Graph Transport Alignment
As analyzed in the previous section, the recent graph contrastive techniques are deeply plagued by positive and negative sample generation, since graph properties could become completely different even with slight perturbations. To design a universal self-supervision scheme, we are motivated to calibrate the similarity between different representations in the output space using the matching signal between corresponding graphs from the input space. In particular, we first perturb the given graph to obtain two different graphs (e.g., a perturbed graph and the original one) and generate the node embeddings of the two graphs using the backbone model (e.g., GNNs). After computing the optimal transport plans for the graphs and the sets of node representations respectively (Sec
tion 4.1), we take the discrepancy between the two plans as the loss to calibrate the backbone for obtaining representation with rich geometry awareness and interpretable correspondences (Section 4.2). We compare the proposed graph self-supervised learning paradigm with graph contrastive learning in Section 4.3. The framework can be found in Figure 2.
Figure 2: A framework of graph transport self-supervised learning.
4.1 Optimal Transport Plan
In general, it is challenging to define two similarity metrics (e.g., distances), which can be directly compared, in two different spaces. This is especially the case for graphs and vectors, two distinct objects by nature. Fortunately, optimal transport theory offers a glimmer of hope, transportation plan, for such comparison. In this section, we present objectives that aim at finding the optimal plan for two graphs (or sets of vectors) to minimize the transport cost. For the graph, a natural idea is to jointly take into account both node attributes and explicit topology information (i.e., edges) in the transportation plan. In this work, we leverage [56] to present an objective function for calculating the optimal transport plan which integrates the edge structure and the feature information on nodes.
Specifically, the fused optimal transport plan  $\pi_{\mathcal{G}}^{*} \in \mathbb{R}^{n \times m}$  between two node-attribute graphs  $\mathcal{G}_1$  and  $\mathcal{G}_2$  can be defined as
$$
\pi_ {\mathcal {G}} ^ {*} = \underset {\pi_ {\mathcal {G}} \in \Pi (\mu , \nu)} {\operatorname {a r g m i n}} \left\langle \sigma \boldsymbol {K} \left(\boldsymbol {X} _ {1}, \boldsymbol {X} _ {2}\right) + (1 - \sigma) \boldsymbol {L} \left(\boldsymbol {A} _ {1}, \boldsymbol {A} _ {2}\right) \otimes \pi_ {\mathcal {G}}, \pi_ {\mathcal {G}} \right\rangle \tag {7}
$$
with the fused Gromov-Wasserstein distance (6). By tuning the parameter  $\sigma$  we can control the bias of the learned optimal plan between node attributes and edge structure. Intuitively we might think that combining more edge information could greatly benefit the expressiveness power of the model. However, we observe that the performance of the proposed method does not degrade significantly in the absence of edge information, and even increases rather than decreases on some datasets. It will be explained in detail later.
For the node representations of graph encoded by the backbone model (e.g., GNNs), we can either use Equation (2) directly to calculate the optimal plan  $\pi_{\mathcal{Z}}^{*} \in \mathbb{R}^{n \times m}$  or set  $\sigma = 1$  in Equation (7) as
$$
\pi_ {\mathcal {Z}} ^ {*} \left(\boldsymbol {Z} _ {1}, \boldsymbol {Z} _ {2}\right) = \underset {\pi_ {\mathcal {Z}} \in \Pi (\mu , \nu)} {\operatorname {a r g m i n}} \langle \boldsymbol {J} \left(\boldsymbol {Z} _ {1}, \boldsymbol {Z} _ {2}\right), \pi_ {\mathcal {Z}} \rangle \tag {8}
$$
where  $J(Z_{1},Z_{2})_{ij} = c_{\mathcal{Z}}(Z_{1}^{i},Z_{2}^{j}), Z_{1}$  and  $Z_{2}$  denote the node representations corresponding to  $\mathcal{G}_1$  and  $\mathcal{G}_2$ , respectively.
4.2 Optimal Transport Alignment
The previous section establishes the foundation that comparing the similarity metrics defined on graph space and vector space can be reduced into comparing the two optimal transport plans from these spaces, each of which can be solved using Equations (7) or (2). This is a valid comparison because the optimal transport plan  $\pi$  acts as a probabilistic matching of two distributions, while the two plan matrices have the same dimensions (i.e.,  $n \times m$ ).
Naturally, the encoder may succeed in obtaining a good representation for a graph if it learns the node embeddings that not only retain its structural information inside the graph, but also capture matching information with other graphs. This motivates us to force the encoder to preserve the matching relationship in the graph space by aligning the plan between the two graphs with the plan of their corresponding node representations. We define a match alignment loss by minimizing the discrepancy between the two transport plans as follows
$$
\mathcal {L} _ {\text {m a t c h}} = \Delta \left(\pi_ {\mathcal {G}} ^ {*}, \pi_ {\mathcal {Z}} ^ {*} \left(\boldsymbol {Z} _ {1}, \boldsymbol {Z} _ {2}\right)\right) \tag {9}
$$
where the discrepancy function  $\Delta (\cdot ,\cdot)$  can be any commonly used metric, e.g., the Frobenius-norm  $\| \cdot -\cdot \| _F$  or the divergence  $D(\cdot \| \cdot)$ .
In addition, to guide the encoder to learn a representation retaining structural information inside the graph, we also calibrate the cost matrix  $J(Z_{1},Z_{2})$ , which implies the implicit structure relationships between nodes, in the representation space as follow
$$
\mathcal {L} _ {\left(\mathrm {i m}\right) \mathrm {s t r c}} = \Delta \left(\sigma K \left(X _ {1}, X _ {2}\right) + (1 - \sigma) L \left(A _ {1}, A _ {2}\right) \otimes \pi_ {\mathcal {G}} ^ {*}, J \left(Z _ {1}, Z _ {2}\right)\right) \tag {10}
$$
To understand the concept of 'implicit structure' intuitively, here we consider the one-dimensional case. As shown in Figure 3, given a point set  $\mathcal{P} = \{Z_1^1,\dots ,Z_1^n\}$ , the location of each of its points is fixed. If the position of a point (yellow)  $Z_{2}^{j}$  is unknown but the transportation
Figure 3: An illustration of implicit structure in one-dimension.
cost from this point to all points (blue) in the point set  $\mathcal{P}$  is known, then the location of this point is determined with respect to the entire point set  $\mathcal{P}$ . The same holds true for another point  $Z_2^l$ . Thus the relative position relationship between points  $Z_2^j$  and  $Z_2^l$  can be captured implicitly by the transport costs matrix  $J(Z_1, Z_2)$ .
To this end, we define the overall graph transport alignment loss as
$$
\mathcal {L} _ {\mathrm {G A L O P A}} = \mathcal {L} _ {\text {m a t c h}} + \rho \mathcal {L} _ {(\mathrm {i m}) \mathrm {s t r c}} \tag {11}
$$
where  $\rho$  is the trade-off parameter.
4.3 Compare with Graph Contrastive Learning
Although our objective function is different from the contrastive loss as shown in Equation (1), we find that the algorithmic philosophy of both is very similar. Here we analyze at the node level. Given a graph  $\mathcal{G}$ , the perturbation graphs  $\mathcal{G}_1$  and  $\mathcal{G}_2$  are obtained by augmenting  $\mathcal{G}$ . If the attribute and context of node  $i$  in  $\mathcal{G}$  are corrupted in a similar way and obtained two node views in  $\mathcal{G}_1$  and  $\mathcal{G}_2$ . The cost between these two views is quite small (almost zero) and thus the optimal transport plan yields a high probability of matching between these two node views. With Equations (9) and (10), the model GALOPA calibrates the matching probability and cost of the corresponding node representations in the output space, which is actually making the representations of two similar nodes similar enough. And for two different nodes  $i$  and  $k$  of  $\mathcal{G}$ , the cost between their corresponding node views in graphs  $\mathcal{G}_1$  and  $\mathcal{G}_2$  is relatively larger. This leads to the opposite correction, i.e., making the representation of two dissimilar nodes sufficiently dissimilar.