SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Gleb Bazhenov*
HSE University, Yandex Research
Yandex Research, Skoltech
Andrey Malinin†
Isomorphic Labs
Yandex Research, HSE University
Liudmila Prokhorenkova
In reliable decision-making systems based on machine learning, models have to be robust to distributional shifts or provide the uncertainty of their predictions. In node-level problems of graph learning, distributional shifts can be especially complex since the samples are interdependent. To evaluate the performance of graph models, it is important to test them on diverse and meaningful distributional shifts. However, most graph benchmarks considering distributional shifts for node-level problems focus mainly on node features, while structural properties are also essential for graph problems. In this work, we propose a general approach for inducing diverse distributional shifts based on graph structure. We use this approach to create data splits according to several structural node properties: popularity, locality, and density. In our experiments, we thoroughly evaluate the proposed distributional shifts and show that they can be quite challenging for existing graph models. We also reveal that simple models often outperform more sophisticated methods on the considered structural shifts. Finally, our experiments provide evidence that there is a trade-off between the quality of learned representations for the base classification task under structural distributional shift and the ability to separate the nodes from different distributions using these representations.
Recently, much effort has been put into creating decision-making systems based on machine learning for various high-risk applications, such as financial operations, medical diagnostics, autonomous driving, etc. These systems should comprise several important properties that allow users to rely on their predictions. One such property is robustness, the ability of an underlying model to cope with distributional shifts when the features of test inputs become different from those encountered in the training phase. At the same time, if a model is unable to maintain high performance on a shifted input, it should signal the potential problems by providing some measure of uncertainty. These properties are especially difficult to satisfy in the node-level prediction tasks on graph data since the elements are interdependent, and thus the distributional shifts may be even more complex than for the classic setup with i.i.d. samples.
To evaluate graph models in the node-level problems, it is important to test them under diverse, complex, and meaningful distributional shifts. Unfortunately, most existing graph datasets split the nodes into train and test parts uniformly at random. Rare exceptions include the Open Graph Benchmark (OGB) [12] that creates more challenging non-random data splits by dividing the nodes
*Correspondence to gy-bazhenov@yandex-team.ru
†Work done while at Yandex Research
according to some domain-specific property. For instance, in the OGB-Arxiv dataset, the papers are divided according to their publication date, which is a realistic setup. Unfortunately, not many datasets contain such meta-information that can be used to split the data. To overcome this issue, Gui et al. [10] propose the Graph OOD Benchmark (GOOD) designed specifically for evaluating graph models under distributional shifts. The authors distinguish between two types of shifts: concept and covariate. However, creating such shifts is non-trivial, while both types of shifts are typically present simultaneously in practical applications. Also, the splitting strategies in GOOD are mainly based on the node features and do not take into account the graph structure. $^3$
In our work, we fill this gap and propose a universal method for inducing structural distributional shifts in graph data. Our approach allows for creating diverse, complex, and meaningful node-level shifts that can be applied to any graph dataset. In particular, we introduce the split strategies that focus on such node properties as popularity, locality, and density. Our framework is flexible and allows one to easily extend it with other structural shifts or vary the fraction of nodes available for training and testing. We empirically show that the proposed distributional shifts are quite challenging for existing graph methods. In particular, the locality-based shift appears to be the most difficult in terms of the predictive performance for most considered OOD robustness methods, while the density-based shift is extremely hard for OOD detection by uncertainty estimation methods. Our experiments also reveal that simple models often outperform more sophisticated approaches on structural distributional shifts. In addition, we investigate some modifications of graph model architectures that may improve their OOD robustness or help in OOD detection on the proposed structural shifts. Our experiments provide evidence that there is a trade-off between the quality of learned representations for base classification task under structural distributional shift and the ability to separate the nodes from different distributions using these representations.
2.1 Graph problems with distributional shifts
Several research areas in graph machine learning investigate methods for solving node-level prediction tasks under distributional shifts, and they are primarily different in what problem they seek to overcome. One such area is adversarial robustness, which requires one to construct a method that can handle artificial distributional shifts that are induced as adversarial attacks for graph models in the form of perturbed and contaminated inputs. The related approaches often focus on designing various data augmentations via introducing random or learnable noise [30, 42, 40, 35, 22].
Another research area is out-of-distribution generalization. The main task is to design a method that can handle real-world distributional shifts on graphs and maintain high predictive performance across different OOD environments. Several invariant learning and risk minimization techniques have been proposed to improve the robustness of graph models to such real-world shifts [1, 19, 20, 38].
There is also an area of uncertainty estimation, which covers various problems. In error detection, a model needs to provide the uncertainty estimates that are consistent with prediction errors, i.e., assign higher values to potential misclassifications. In the presence of distributional shifts, the uncertainty estimates can also be used for out-of-distribution detection, where a model is required to distinguish the shifted OOD data from the ID data [24, 36, 2, 27, 18].
The structural distributional shifts proposed in our paper can be used for evaluating OOD generalization, OOD detection, and error detection since they are designed to replicate the properties of realistic graph data.
2.2 Uncertainty estimation methods
Depending on the source of uncertainty, it is usually divided into data uncertainty, which describes the inherent noise in data due to the labeling mistakes or class overlap, and knowledge uncertainty, which accounts for the insufficient amount of information for accurate predictions when the distribution of the test data is different from the training one [5, 24, 23].
<sup>3</sup>In Appendix C, we provide a detailed discussion of the GOOD benchmark.
General-purpose methods The most simple approaches are standard classification models that predict the parameters of softmax distribution. For these methods, we can define the measure of uncertainty as the entropy of the predictive categorical distribution. This approach, however, does not allow us to distinguish between data and knowledge uncertainties, so it is usually inferior to other methods discussed below.
Ensembling techniques are powerful but expensive approaches that providing decent predictive performance and uncertainty estimation. The most common example is Deep Ensemble [21], which can be formulated as an empirical distribution of model parameters obtained after training several instances of the model with different random seeds for initialization. Another way to construct an ensemble is Monte Carlo Dropout [6], which is usually considered as the baseline in uncertainty estimation literature. However, it has been shown that this technique is commonly inferior to deep ensembles, as the obtained predictions are dependent and thus lack diversity. Importantly, ensembles allow for a natural decomposition of total uncertainty in data and knowledge uncertainty [23].
There is also a family of Dirichlet-based methods. Their core idea is to model the point-wise Dirichlet distribution by predicting its parameters for each input individually using some model. To get the parameters of the categorical distribution, one can normalize the parameters of the Dirichlet distribution by their sum. This normalization constant is called evidence and can be used to express the general confidence of a Dirichlet-based model. Similarly to ensembles, these methods are able to distinguish between data and knowledge uncertainty. There are numerous examples of Dirichlet-based methods, one of the first being Prior Network [24] that induces the behavior of Dirichlet distribution by contrastive learning against the OOD samples. Although this method is theoretically sound, it requires knowing the OOD samples in the training stage, which is a significant limitation. Another approach is Posterior Network [2], where the behavior of the Dirichlet distribution is controlled by Normalizing Flows [15, 13], which estimate the density in latent space and reduce the Dirichlet evidence in the regions of low density without using any OOD samples.
Graph-specific methods Recently, several uncertainty estimation methods have been designed specifically for node-level problems. For example, Graph Posterior Network [33] is an extension of the Posterior Network framework discussed above. To model the Dirichlet distribution, it first encodes the node features into latent representations with a graph-agnostic model and then uses one flow per class for density estimation. Then, the Personalized Propagation scheme [17] is applied to the Dirichlet parameters to incorporate the network effects. Another Dirichlet-based method is Graph-Kernel Dirichlet Estimation [41]. In contrast to Posterior Networks, its main property is a compound training objective, which is focused on optimizing the node classification performance and inducing the behavior of the Dirichlet distribution. The former is achieved via knowledge distillation from a teacher GNN, while the latter is performed as a regularisation against the prior Dirichlet distribution computed via graph kernel estimation. However, as shown by Stadler et al. [33], this approach is inferior to Graph Posterior Network while having a more complex training procedure and larger computational complexity.
2.3 Methods for improving robustness
Improving OOD robustness can be approached from different perspectives, which, however, share the same idea of learning the representations that are invariant to undesirable changes of the input distribution. For instance, the main claim in domain adaptation is that to achieve an effective domain transfer and improve the robustness to distributional shift, the predictions should depend on the features that can not discriminate between the source and target domains. Regarding the branch of invariant learning, these methods decompose the input space into different environments and focus on constructing robust representations that are insensitive to their change.
General-purpose methods A classic approach for domain adaptation is Domain-Adversarial Neural Network [7] that promotes the emergence of features that are discriminative for the main learning task on the source domain but do not allow to detect the distributional shift on other domains. This is achieved by jointly optimizing the underlying representations as well as two predictors operating on them: the main label predictor that solves the base classification task and the domain classifier that discriminates between the source and the target domains during training. Another simple technique in the class of unsupervised domain adaptation methods is Deep Correlation Alignment
[34], which trains to align the second-order statistics of the source and target distributions that are produced by the activation layers of a deep neural model that solves the base classification task.
A representative approach in invariant learning is Invariant Risk Minimization [1], which searches for data representations providing decent performance across all environments, while the optimal classifier on top of these representations matches for all environments. Another method is Risk Extrapolation [20], which targets both robustness to covariate shifts and invariant predictions. In particular, it targets the forms of distributional shifts having the largest impact on performance in training domains.
Graph-specific methods More recently, several graph-specific methods for improving OOD robustness have been proposed. One of them is an invariant learning technique Explore-to-Extrapolate Risk Minimization [38], which leverages multiple context explorers that are specified by graph structure editors and adversarially trained to maximize the variance of risks across different created environments. There is also a simple data augmentation technique Mixup [39] that trains neural models on convex combinations of pairs of samples and their corresponding labels. However, devising such a method for solving the node-level problems is not straightforward, as the inputs are connected to each other. Based on this technique, Wang et al. [37] have designed its adaptation for graph data: instead of training on combinations of the initial node features, this method exploits the intermediate representations of nodes and their neighbors that are produced by graph convolutions.
3 Structural distributional shifts
3.1 General approach
As discussed above, existing datasets for evaluating the robustness and uncertainty of node-level problems mainly focus on feature-based distributional shifts. Here, we propose a universal approach that produces non-trivial yet reasonable structural distributional shifts. For this purpose, we introduce a node-level graph characteristic  $\sigma_{i}$  and compute it for every node  $i\in \mathcal{V}$ . We sort all nodes in ascending order of  $\sigma_{i}$  — those with the smallest values of  $\sigma_{i}$  are considered to be ID, while the remaining ones are OOD. As a result, we obtain a graph-based distributional shift where ID and OOD nodes have different structural properties. The type of shift depends on the choice of  $\sigma_{i}$ , and several possible options are described in Section 3.2 below.
We further split the ID nodes uniformly at random into the following parts:
- Train contains nodes  $\nu_{\mathrm{train}}$  that are used for regular training of models and represent the only observations that take part in gradient computation.
- Valid-In enables us to monitor the best model during the training stage by computing the validation loss for nodes  $\mathcal{V}_{\mathrm{valid - in}}$  and choose the best checkpoint.
- Test-In is used for testing on the remaining ID nodes  $\nu_{\mathrm{test-in}}$  and represents the simplest setup that requires a model to reproduce in-distribution dependencies.
The remaining OOD nodes are split into Valid-Out and Test-Out subsets based on their  $\sigma_{i}$ :
- Test-Out is used to evaluate the robustness of models to distributional shifts. It consists of nodes with the largest values of  $\sigma_{i}$  and thus represents the most shifted part  $\nu_{\mathrm{test - out}}$
- Valid-Out contains OOD nodes with smaller values of  $\sigma_{i}$  and thus is less shifted than Test-Out. This subset  $\mathcal{V}_{\mathrm{valid - out}}$  can be used for monitoring the model performance on a shifted distribution. Our experiments assume a more challenging setup when such shifted data is unavailable during training. However, the presence of this subset allows us to further separate  $\mathcal{V}_{\mathrm{test - out}}$  from  $\mathcal{V}_{\mathrm{train}}$  — the larger  $\mathcal{V}_{\mathrm{valid - out}}$  we consider, the more significant distributional shift is created between the train and OOD test nodes.
Our general framework is quite flexible and allows one to easily vary the size of the training part and the type of distributional shift. Let us now discuss some particular shifts that we propose in this paper.
(a) Popularity
(b) Locality
(c) Density
Figure 1: Visualization of structural shifts for AmazonPhoto dataset: ID is blue, OOD is red.
3.2 Proposed distributional shifts
To define our data splits, we need to choose a node property  $\sigma_{i}$  as a splitting factor. We consider diverse graph characteristics covering various distributional shifts that may occur in practice. In a standard non-graph ML setup, shifts typically happen only in the feature space (or, more generally, the joint distribution of features and targets may shift). However, in graph learning tasks, there can be shifts specifically related to the graph structure. We discuss some representative examples below.
Popularity-based The first strategy represents a possible bias towards popularity. In some applications, it is natural to expect the training set to consist of more popular items. For instance, in the web search, the importance of pages in the internet graph can be measured via PageRank [29]. For this application, the labeling of pages should start with important ones since they are visited more often. Similar situations may happen for social networks, where it is natural to start labeling with the most influential users, or citation networks, where the most cited papers should be labeled first. However, when applying a graph model, it is essential to make accurate predictions on less popular items. Motivated by that, we introduce a popularity-based split based on PageRank. The vector of PageRank values  $\pi_{i}$  describes the stationary distribution of a random walk with restarts, for which the transition matrix is defined as the normalized adjacency matrix  $\mathbf{AD}^{-1}$  and the probability of restart is  $\alpha$ :
$$
\boldsymbol {\pi} = (1 - \alpha) \mathbf {A} \mathbf {D} ^ {- 1} \boldsymbol {\pi} + \alpha \boldsymbol {p}. \tag {1}
$$
The vector of restart probabilities  $\pmb{p}$  is called a personalization vector, and  $p_i = 1 / n$  by default, which describes the uniform distribution over nodes.
To construct a popularity-based split, we use PageRank (PR) as a measure of node importance. Thus, we compute PR for every node  $i$  and set  $\sigma_{i} = -\pi_{i}$ , which means that the nodes with smaller PR values (i.e., less important ones) belong to the OOD subsets. Note that instead of PR, one can potentially use any other measure of node importance, e.g., node degree, or betweenness centrality.
Figure 1a illustrates that the proposed shift separates the most important nodes that belong to the cores of large clusters and the structural periphery, which consists of less important nodes in terms of their PR values. We also observe that such popularity-based split agrees well with some natural distributional shifts. For instance, Figure 11a in Appendix shows the distribution of PageRank in train and test parts of the OGB-Arxiv dataset [12]. Here, the papers are split according to their publication date. Thus, older papers naturally have more citations and, therefore, larger PageRank. Our proposed split allows one to mimic this behavior for datasets without timestamps available. We refer to Appendix A for more details and additional analysis.
Locality-based Our next strategy is focused on a potential bias towards locality, which may happen when labeling is performed by exploring the graph starting from some node. For instance, in web search applications, a crawler has to explore the web graph following the links. Similarly, the information about the users of a social network can usually be obtained via an API, and new users are discovered following the friends of known users. To model such a situation, one could divide nodes based on the shortest path distance to a given node. However, graph distances are discrete, and the number of nodes at a certain distance may grow exponentially with distance. Thus, such an approach
does not provide us with the desired flexibility in varying the size of the train part. Instead, we use the concept of Personalized PageRank (PPR) [29] to define a local neighborhood of a node. PPR is the stationary distribution of a random walk that always restarts from some fixed node  $j$ . Thus, the personalization vector  $p$  in (1) is set to the one-hot-encoding of  $j$ . The associated distributional shift naturally captures locality since a random walk always restarts from the same node.
For our experiments, we select the node  $j$  with the highest PR score as a restarting one. Then, we compute the PPR values  $\pi_{i}$  for every node  $i$  and define the measure  $\sigma_{i} = -\pi_{i}$ . The nodes with high PPR, which belong to the ID part, are expected to be close to the restarting node, while far away nodes go to the OOD subset. Figure 1b confirms that the locality is indeed preserved, as the ID part consists of one compact region around the restarting node. Thus, the OOD subset includes periphery nodes as well as some nodes that were previously marked as important in the PR-based split but are far away from the restarting node. Our analysis in Appendix A also provides evidence for this behavior: the PPR-based split strongly affects the distribution of pairwise distances within the ID/OOD parts as the locality bias of the ID part makes the OOD nodes more distant from each other.
While locality-based distributional shifts are natural, we are unaware of publicly available benchmarks focusing on such shifts. We believe that our approach will be helpful for evaluating the robustness of GNNs under such shifts. Our empirical results in Section 4 demonstrate that locality-based splits are the most challenging for graph models and thus may require special attention.
Density-based The next distributional shift we propose is based on density. One of the most simple node characteristics that describe the local density in a graph is the local clustering coefficient. Considering some node  $i$ , let  $d_{i}$  be its degree and  $\gamma_{i}$  be the number of edges connecting the neighbors of  $i$ . Then, the local clustering coefficient is defined as the edge density within the one-hop neighborhood:
$$
c _ {i} = \frac {2 \gamma_ {i}}{d _ {i} \left(d _ {i} - 1\right)} \tag {2}
$$
For our experiments, we consider nodes with the highest clustering coefficient to be ID, which implies that  $\sigma_{i} = -c_{i}$ .
This structural property might be particularly interesting for inducing distributional shifts since it is defined through the number of triangles, the substructures that most standard graph neural networks are unable to distinguish and count [4]. Thus, it is interesting to know how changes in the clustering coefficient affect the predictive performance and uncertainty estimation.
Figure 1c visualizes the density-based split. We see that the OOD part includes both the high-degree central nodes and the periphery nodes of degree one. Indeed, the nodes of degree one naturally have zero clustering coefficient. On the other hand, for the high-degree nodes, the number of edges between their neighbors usually grows slower than quadratically. Thus, such nodes tend to have a vanishing clustering coefficient.
Finally, we note that the existing datasets with realistic splits may often have the local clustering coefficient shifted between the train and test parts. Figures 12c and 13c in Appendix show this for two OGB datasets. Depending on the dataset, the train subset may be biased towards the nodes with either larger (Figure 13c) or smaller (Figure 12c) clustering. In our experiments, we focus on the former scenario.