Improvements on Uncertainty Quantification for Node Classification via Distance-Based Regularization

We define a graph with attributed node-level features $\mathcal{G}=(\mathcal{V},\mathcal{E},X,Y_{\mathbb{L}})$, where $\mathcal{V}$ is a set of nodes on the graph with cardinality $N$ and $\mathcal{E}\subset\mathcal{V}\times\mathcal{V}$ denotes a set of graph edges that can be represented by an adjacency matrix $W$. A feature matrix is denoted by $X=[\mathbf{x}_{1},\dots,\mathbf{x}_{N}]^{T}\in\mathbb{R}^{N\times d}$, in which each row $\mathbf{x}_{i}\in\mathbb{R}^{d}$ is a feature vector of node $i$ with dimension $d$. Under the semi-supervised learning setting, a set of labels is available, denoted by $Y_{\mathbb{L}}=\{y_{i}\mid i\in\mathbb{L}\}$, where $\mathbb{L}\subset\mathcal{V}$ and $y_{i}\in\{1,\dots,K\}$ for $K$ classes.

Our goal is to design and learn a deterministic GNN based on $\mathcal{G}$ that takes the feature matrix ${X}$ and the adjacency matrix $W$ as INPUT and predicts the parameters of a Dirichlet distribution for each node $i\in\mathcal{V}$ as OUTPUT, denoted as ${\bm{\alpha}}_{i}$, which is often referred to as the concentration parameters. Therefore, the network function $F_{\Theta}$ can be expressed as: $\mathcal{A}=F_{\Theta}({X},W)$, where $\mathcal{A}:=[{\bm{\alpha}}_{i}]_{i\in\mathcal{V}}$ is a matrix and $\Theta$ refers to network parameters. The statistical relations between the class label ${\bf y}_{i}$, the vector of class probabilities ${\bf p}_{i}$, and the Dirichlet parameters ${\bm{\alpha}}_{i}$ can be represented as:

Based on the predictions of $\mathcal{A}$, the expected vector of class probabilities $\bar{\bf p}_{i}:=\mathbb{E}[{\bf p}_{i}|{\bm{\alpha}}_{i}]=[\alpha_{i1}/\alpha_{i0},\cdots,\alpha_{iK}/\alpha_{i0}]^{T}$, where $\alpha_{i0}=\sum_{k=1}^{K}\alpha_{ik}$ is called the Dirichlet strength. The aleatoric and epistemic uncertainties about the classification of each node $i$ can be calculated as:

respectively. The aleatoric uncertainty is measured by the negative of the largest class probability in $\bar{\bf p}_{i}$. This uncertainty is higher when the largest class probability in $\bar{\bf p}_{i}$ is lower, which implies that the model is less confident and the probabilities are more evenly spread across classes. On the other hand, the epistemic uncertainty is measured by the negative of the Dirichlet strength $\alpha_{i0}$, whose value is higher when the Dirichlet strength is lower, meaning that the model is unfamiliar with the feature vector of node $i$ and the predicted Dirichlet distribution is less concentrated around a specific point or set of points on the probability simplex [22]. We note that a high aleatoric uncertainty may not indicate a high epistemic uncertainty and vice versa. For example, two evidence parameters ${\bm{\alpha}}_{i}=[1,\cdots,1]$ and ${\bm{\alpha}}_{j}=[1000,\cdots,1000]$ have the same aleatoric uncertainty: $-1/K$, since they have the same projected class probabilities; but their epistemic uncertainties differ drastically: $K$ versus $1000K$. Please refer to [32, 33] for rationales of aleatoric and epistemic uncertainties in (2).

Our framework is based on graph posterior network (GPN) [32], which extends posterior network (PN) [7] to semi-supervised node classification. GPN consists of three main steps. First, a feature encoder maps the original features onto a low-dimensional latent space with a simple two-layer multi-layer perception (MLP) encoder. Second, a radial normalizing flow [27] estimates the density of the latent space per class. Lastly, a personalized page rank message passing scheme [13] diffuses the pseudo counts (density multiplied by the number of training nodes) by taking the graph structure into account. We summarize the three steps with notations as follows,

1. 1.

   Multi-layer perceptron for representation learning: $\mathbf{z}_{i}=f(\mathbf{x}_{i};{\bm{\theta}})$ or $f_{\bm{\theta}}$ in short.

2. 2.

   Normalizing flow for density estimation: $g_{\bm{\phi}}$ for short, and more specifically

   $$g_{\bm{\phi}}(\mathbf{z}_{i})_{k}=N_{k}\cdot\mathbb{P}(\mathbf{z}_{i}|k;{\bm{\phi}}),$$

   (3)

   where $N_{k}$ is the number of training nodes belonging to the class $k$, $\mathbf{z}_{i}$ is the embedding vector of node $i$ obtained via the first step, $\mathbb{P}(\mathbf{z}_{i}|k;{\bm{\phi}})$ is the conditional density per class $k$ estimated by a normalizing flow module, and ${\bm{\phi}}$ denotes the parameters of this module. GPN also includes the evidence computed prior to the graph aggregation, defined by

   $$\bm{\alpha}_{i}^{\text{feat}}=g_{\bm{\phi}}(\mathbf{z}_{i})+\mathbf{1}.$$

   (4)

3. 3.

   Personalized page rank (PPR) for evidence diffusion: $\bm{\beta}^{\text{aggr}}_{i,k}=\sum_{j\in\mathcal{V}}{\bm{\Pi}}_{i,j}^{\text{ppr}}g_{\bm{\phi}}(\mathbf{z}_{j})_{k},$ where ${\bm{\Pi}}_{i,j}^{\text{ppr}}$ refer to the dense PPR scores implicitly reflecting the importance of node $j$ from the perspective of node $i$. Then we can get the predicted concentrate parameters $\bm{\alpha}$ with a uniform prior $\mathbf{1}$ for a non-degenerated Dirichlet distribution, i.e.,

   $$\bm{\alpha}_{i}=\bm{\beta}_{i}^{\text{aggr}}+\mathbf{1}.$$

   (5)

As opposed to GPN, PN is designed for uncertainty estimation for i.i.d. inputs, which only considers the first two steps to predict the Dirichlet distribution $\text{Dir}(\bm{\alpha}_{i}^{\text{feat}})$.

Given the labels of training nodes: $Y_{\mathbb{L}}=\{y_{i}\mid i\in\mathbb{L}\}$, GPN is trained by minimizing the following Bayesian loss:

The first term in (6), called uncertainty cross entropy (UCE) [5], is defined by

where $Y=[{\bf y}_{i}]_{i\in[N]}$, ${\bf y}_{i}\in\{0,1\}^{K}$ is the one-hot encoded ground-truth class of the node $i$, and $\Psi$ is the digamma function, in terms of the Gamma function by: $\Psi(x)=\frac{\Gamma^{\prime}(x)}{\Gamma(x)}$. Minimizing UCE is known to increase confidence in classifying observed data (training nodes in this context). The second term in (6) is based on the entropy of each node-level Dirichlet distribution $\text{Dir}(\bm{\alpha}_{i})$ that favors smooth distributions. For more details on GPN, please refer to Appendix C.

The loss function plays a pivotal role in learning effective representation functions and density estimations. In this context, we establish several theorems (Theorems 1 and 4) to describe some demanding assumptions on $f_{\theta}$ and $g_{\phi}$ that achieve the minimum UCE loss. We then describe a limitation of UCE in separating ID and OOD nodes in Theorem 6 and Corollary 7 for PN, which means that we only consider the first two steps in GPN. The main conclusion of our analysis is that the UCE loss function alone is insufficient to learn a representation space that separates OOD from ID nodes. We take graph connectivity into account in Theorem 8 to study some scenarios where GPN is ineffective for OOD detection. Although our theorems do not completely characterize graph learning, they provide some insights into the behavior of the network parameters in PN/GPN when minimizing the UCE loss.

For the proofs, please refer to Appendix B. Here, we elaborate on the ideal configuration that satisfies the conclusion of Theorem 1. We assume the MLP function $f_{\theta}$ is arbitrarily complex such that it maps $\mathbf{x}_{i}\in\mathcal{X}_{k}$ into a bounded ball in the representational space, i.e.,

where each ball $B(\mathbf{z}_{k},r_{k})$ is centered at a point $\mathbf{z}_{k}$ and bounded in size with $r_{k}<r$ for a positive value $r$. Furthermore, we choose the normalizing flow $g_{\phi}$ to be

where Vol$(\cdot)$ refers to the volume of the ball. The conclusion in Theorem 1 states that for every $\epsilon,$ there exists a suitable upper bound $r$ of all the balls such that $\text{UCE}(\mathcal{A},Y)<\epsilon$.

An implication of Theorem 1 is that UCE is not sufficient to separate OOD from ID nodes. Example 2 illustrates a scenario that OOD nodes can be close to ID nodes in the learned representation space even though they can be separated in the feature space based on OOD-specific features, which are unfortunately discarded. In other words, the learned representation space by GPN based on the UCE loss is not guaranteed to preserve the distance between OOD and ID nodes in its representation learning step.

Limitations and discussions $-$ The first assumption in Theorem 1 regarding the distinct class-specific distributions of feature vectors might not be realistic in practice since certain ID classes may not be clearly distinguishable due to noise in features or class labels. Nevertheless, our intuition suggests that if the UCE loss is inadequate for separating OOD from ID nodes in situations where they are separable, it is even more likely to falter in the more complex, non-separable cases. As for the second assumption in Theorem 1, it is true that arbitrary complexity of $f_{\theta}$ and $g_{\phi}$ does not fully respect the inductive bias [24] of the network design, such as the MLP layers with ReLU for the feature encoder $f({\bf x};{\bm{\theta}})$, our analysis remains insightful and informative about the structures these networks are likely to exhibit. For example, we may expect from Theorem 1 that the representation of each class may favor an embedded space that compresses OOD-specific features, while density estimation $g_{\phi}$ tends to have higher peaks over smaller volumes as the model consolidates the representation space.

We note that in [7, Theorem 1] and [32, Theorem], the authors demonstrated that PN/GPN is able to achieve reasonable uncertainty estimation when the feature encoder is a ReLU network and PPR diffusion is removed to disregard network effects. Unfortunately, the analysis is based on extreme node features, specifically as $\delta\rightarrow\infty,$

for any node $i$ with a high probability. Moreover, Theorem 1 in the GPN paper [32] holds even when the supports of disjoint ID and OOD classes in the latent space overlap. In other words, this result does not prevent distant nodes from having similar representations. In summary, the analysis based on extreme node features may provide limited insights about the issues of GPN studied in this section. See Appendix D for detailed discussions.

By taking $\epsilon\rightarrow 0$, Theorem 1 reduces to Corollary 3 where UCE is equal to 0.

Next, we aim for the construction of a specific case to make UCE equal to zero. As it is challenging to analyze the joint minimization on $\theta$ and $\phi$, we assume that the normalizing flow can be chosen optimally. For this purpose, we consider a simplified problem where the true density function is assumed to be known for a given $\theta$ and hence it can be used to replace the learning of the normalizing flow. As a result, the problem reduces to the learning of the representation network $\theta$.

Notice that the true analytical solution $\hat{\phi}$ in Theorem 4 is a function of $\theta$, i.e., $\hat{\phi}=\phi(\theta)$. It is possible to know an analytical form of $\phi(\theta)$. For example, if the data points belonging to each class are sampled from a known Gaussian distribution in the original feature space and the representation network is a linear projection function, then the true density of the projected data points belonging to each class can be derived based on any configuration of the known Gaussian distribution in the original feature space.

Before discussing OOD detection, we start with the definition of OOD nodes.

In order to detect the OOD nodes, we need a good representation, e.g., $f_{\theta}(\mathcal{X}_{[K]})\bigcap f_{\theta}(\mathcal{X}-\mathcal{X}_{[K]})=\varnothing$. However, it is possible that $\cup_{k=1}^{K}f_{\theta}^{-1}(\mathcal{Z}_{k})=\mathcal{X}$, which implies that any OOD node in the feature space is mapped to the same distribution of an ID class in the representation space. To this end, we require the preimage, $f_{\theta}^{-1}$, to be well-behaved in the sense that $f_{\theta}^{-1}(\mathcal{Z}_{k})$ must be contained within a bounded region near the support of the true distribution $\mathcal{X}_{k}$. Explicitly, we require there exists a constant $\delta>0$ such that $d(\mathbf{x},\hat{\mathbf{x}})<\delta,\forall\hat{\mathbf{x}}\in f_{\theta}^{-1}(\mathcal{Z}_{k})$ and $\forall\mathbf{x}\in\mathcal{X}_{k}$ for each $k$. The choice of $\delta$ depends on additional information about the dataset. An excessively large $\delta$ causes overlap between classes, thus increasing the likelihood of improperly identifying the OOD nodes as ID. On the other hand, a relatively small $\delta$ forces the model to overfit, which labels the non-training nodes as OOD.

We characterize in Theorem 6 that under certain conditions, OOD can be detected correctly if they are far away from the ID distribution. One such condition is that the function $f_{\theta}$ is well-fit, meaning that it maps the support of $\mathcal{X}_{k}$ inside $B(\mathbf{z}_{k},r_{k})$ for every $k$ in $[K]$. We denote a set of all well-fit functions by $\bm{\Gamma}$. Please refer to Figure 1 for a geometric illustration of these relevant quantities.

The main purpose of Theorem 6 is to show a desired behavior for OOD detection is induced by point-wise effects. In practice, we suggest the careful construction of $f_{\theta}$’s topology coupled with proper regularization terms to achieve the point-wise effect.

Lastly, we consider the setting of GPN to use a graph layer after the feature-wise evidence predictions.

Limitations and Discussions. In Theorem 8, we provide a special situation where the GPN architecture may fail to detect OOD nodes by predicting large evidence values for belonging to class 0. In addition, we show in Appendix B that if GPN has bad initial feature predictions, even ideal graph construction coupled with an ideal graph neural network (with a homophily degree 1.0) fails to prevent the OOD nodes from being misclassified as ID nodes. For homophily graphs with degrees less than 1.0, the majority of nodes for each class have similar evidence values after graph diffusion layers, except that the nodes between some pairs of classes may have different evidence values.

As discussed in Section 4.1, the UCE loss function alone is insufficient to learn a representation space that separates OOD from ID nodes using the GPN model. We propose a heuristic remedy that enforces distance minimization on the graph. Ideally, we should design a distance formula that can preserve the distance relationship among all the feature vectors. However, distance preservation likely increases variation in the latent space as we cannot compress the classes’ support in the representation space to be arbitrarily small, while simultaneously preserving distances.

Instead, we consider distance minimization, as it helps prevent the model from discarding relevant features while decreasing variation between nodes in the representation space. Here we give two formulations directly. The simpler, theorem-motivated term, is the distance regularization on the latent space,

The regularization encourages nearby points in the graph representation to remain nearby in the latent space. In other words, this regularization (9) discourages the overlap $f_{\theta}(\mathcal{X}_{[K]})\bigcap f_{\theta}(\mathcal{X}-\mathcal{X}_{[K]})$.

We also minimize the "distance" on the produced evidence through a divergence-based regularization,

where $\text{div}_{\text{kl}}$ refers to the Kullback–Leibler divergence and two symmetric terms are considered. This divergence-based formulation likely decreases the variation in evidence between neighboring nodes, because high variation in the latent space between neighboring nodes need not be mapped to similar evidence.

In summary, we augment the GPN mode with either one of the proposed regularization terms in (9) and (10), thus leading to the objective function as follows,

with two positive parameters $\lambda_{1},\lambda_{2}.$ The first term is the standard UCE loss function. The second term is regarded by GPN as an entropy regularizer. The last term, R, is chosen to be either $\text{R}_{D}$ or $\text{R}_{\alpha}$, a decision implemented through hyperparameter tuning. We have included a theoretical result in Appendix B that provides a rationale for the proposed distance regularizations.

Our proposed model differs from the GPN model in three main aspects. First, we use validation cross entropy (CE) instead of hold-out datasets to select hyperparameters. Second, we consider the activation function for the MLP layers as one of the hyperparameters for selection. We expect feature value rescaling through non-linear activation of feature values to affect the density predictions. Third, we incorporate one of the proposed distance-based regularization terms to the loss function used in GPN.We conduct an ablation study to demonstrate the contribution of these three components. The results for CiteSeer and PubMed are shown in Table 3 and the remaining datasets are included in Appendix E. Hyperparameter tuning using validation cross-entropy improves GPN’s performance, especially in cases where the choice of activation function has a significant impact on specific datasets. Additionally, we consistently observe performance enhancements from the distance-based regularization in both datasets, demonstrating the effectiveness of the proposed distance awareness regularization term.