Generalized Equivariance and Preferential Labeling
for GNN Node Classification

We now formulate the equivariance property of node classification tasks. It essentially asserts that for any permutation $\pi\in S_{n}$,

That is to say, if we permute the order of nodes, the solution should be changed correspondingly.

Suppose for every $X$ there exists a unique $Y$ satisfying $H(X,Y)$, the mapping from $X$ to $Y$ can be modeled by a function $h$ and the equivariance property becomes the form that we commonly see

for every permutation $\pi\in S_{n}$.

In the above formulation, we define the equivariance property of a node classification task. In fact, equivariance can also be said in terms of GNN output $f(X)$, given by

We first consider distinguishing different nodes in a graph. The state-of-the-art approaches (Li, Chen, and Koltun 2018; Selsam et al. 2019; Zhang et al. 2020) assign all nodes with the same embedding, and thus, the model cannot distinguish different nodes effectively. Consider a common graph convolutional network (GCN), which learns the hidden representation for a node by encoding the node vector with the neighbors via a set of fully-connected layers. In the example given by Figure 2, all four nodes will have the same hidden representation, because every node is represented by the same embedding and all nodes have the same neighboring information.

We now consider the equivariance property of node classification for unattributed graphs, which is believed to be important for various GNN applications (Chen, Li, and Bruna 2018; Azizian et al. 2020).

In node classification for unattributed graphs, if the node index changes, the output would change accordingly, shown in Eqn (1). Thus, it is tempting to design an equivariant GNN satisfying Eqn (3) for node classification tasks, as suggested by Wu et al. 2021 (Wu et al. 2021). Otherwise, the GNN would be sensitive to labeling artefacts (Allamanis, Brockschmidt, and Khademi 2018; Wei et al. 2020), if it does not satisfy some form of equivariance. Previous work achieves the equivariance property (3) by using the same embeddings for all nodes (Li, Chen, and Koltun 2018; Selsam et al. 2019; Zhang et al. 2020).

However, we hereby show that an equivariant GNN satisfying (3) will fail on node classification problems, where multiple outputs are appropriate. In other words, the mapping from $X$ to $Y$ is not a function, and given an input $X$, there exists multiple $Y$ such that $H(X,Y)$ holds. Usually, GNN predicts one appropriate $Y$ by a function $Y=f(X)$.

We show the drawback of equivariant GNNs with an example of a non-trivial auto-isomorphic graph, i.e., there exists a non-identity permutation $\pi$ such that $\pi(X)=X$. If (3) holds, then $\pi(X)=X$ implies $f(X)=\pi(f(X))$. This means that GNN output must be the same for all corresponding nodes shuffled by $\pi$.

This, unfortunately, may be a bad solution for various tasks. Consider the maximum independent set (MIS) problem that selects the largest number of vertices that are not directly connected. In Figure 2, for example, $\{1,3\}$ is an MIS and $\{2,4\}$ is also an MIS. However, an equivariant GNN cannot predict any MIS in this example, because there exists a permutation $\pi$ (e.g., $\pi:1\mapsto 2,2\mapsto 3,3\mapsto 4,4\mapsto 1$) essentially tying the output of all nodes.

It should be mentioned that we show the limitation by an example of non-trivial auto-isomorphism. Our analysis is suggestive to real applications, where graphs often have similar local structures.

To address the limitation of the node distinction, a naïve idea is still to assign node embeddings by randomly labeling the nodes, but to use an ensemble of different node labelings to eliminate artefacts (Murphy et al. 2019; Sato, Yamada, and Kashima 2021). For training, node labels are assigned randomly; this serves as a way of data augmentation, and can be thought of as training an ensemble over epochs. During inference, it assigns multiple random labels and uses an average ensemble for prediction.

However, it is not appropriate if we directly apply such a naïve idea to equivariant node classification. The standard cross-entropy training is essentially to

where $\omega$ denotes the trainable parameters in a GNN, $\mathcal{D}$ is the training set, and $D_{\mathrm{KL}}$ is the Kullback–Leibler divergence between the predictions and the ground truth. $(\pi(Y))_{i}$ and $f_{i}(\pi(X))$ are the $i$th row in a matrix, representing the target and predicted distributions of a node. However, the training objective will enforce the GNN to learn the same prediction of nodes under auto-isomorphism, because for every $\pi\in S_{n}$, the KL objective requires $\pi(Y)=f(\pi(X))$. For some auto-isomorphic permutation $\tau$, i.e., $\tau(X)=X$, this implies $f(X)=f(\tau(X))=\tau(Y)$. Since $Y=f(X)$, we will have $f(X)=\tau(f(X))$. This is precisely the limitation that we have analyzed in Section 2.2, namely, auto-isomorphism tying the prediction of a GNN.

To address this issue, we propose a desired generalized equivariance property.

An equivariant GNN satisfying (3) fails for equivariant node classification, because it unreasonably assumes the output is a function of input, i.e., approximating (1) by (2).

We relax this constraint for multi-output node classification and analyze the desired form of equivariance in this setting. We denote $\mathcal{H}(X)=\{Y:H(X,Y)\}$ be the set of all correct outputs given a graph $X$. The training set typically provides one groundtruth solution $Y_{*}\in\mathcal{H}(X_{*})$ for a specific graph $X_{*}$, as usually one solution suffices in real applications (and this coincides with GNNs whose output is a function of input).

We would define $\mathcal{H}_{*}(\cdot)|_{\mathcal{D}}$ as a minimal equivariant subset of $\mathcal{H}(\cdot)$ such that $Y_{*}\in\mathcal{H}_{*}(X_{*})$, with the domain restricted to $\mathcal{D}=\{X:X=\pi(X_{*})\text{ for some }\pi\in S_{n}\}$.

This starts from defining

which essentially endorses multiple correct outputs other than the given $Y_{*}$ due to self-isomorphism $\gamma$.

Then, the equivariance property suggests $\mathcal{H}_{*}(X)=\pi\left(\mathcal{H}_{*}(X_{*})\right)$, if $X=\pi(X_{*})$ for some $\pi$. Here, we abuse the notation as $\pi(\mathcal{H}_{*}(X_{*}))\overset{\Delta}{=}\{\pi(Y):Y\in\mathcal{H}_{*}(X_{*})\}$.

We would like to design a neural network predicting a correct solution, i.e.,

Due to the equivariance of $\mathcal{H}_{*}$, we have

By the definition of $\mathcal{H}_{*}$ in (4), Eqn (5) implies that there exists $\gamma_{1}\in S_{n}$ such that $\gamma_{1}(X_{*})=X_{*}$ and $f(X_{*})=\gamma_{1}(Y_{*})$. Likewise, Eqn (6) implies that there exists $\gamma_{2}\in S_{n}$ such that $\gamma_{2}(X_{*})=X_{*}$ and $f(\pi(X_{*}))=\pi\gamma_{2}(Y_{*})$. Combining these and denoting $\gamma_{2}\gamma_{1}^{-1}$ by $\gamma$, we see that there exists $\gamma\in S_{n}$ such that

We call (7) the generalized equivariance property. In fact, (3) is a special case of (7), where $\gamma$ is an identity permutation. However, we relax (3) by allowing an additional auto-isomorphic permutation $\gamma$ in the solution space, and thus, it does not suffer from the limitation in Section 2.2.

The analysis shows that an ideal GNN for unattributed node classification should satisfy (7) rather than (2).

Inspired by the above analysis, we propose a simple yet effective approach, Preferential Labeling, which asymptotically satisfies (7). The overview of our approach is shown in Figure 3.

For training, Preferential Labeling assigns nodes with a random permutation of labels. A node is represented by a table-lookup embedding based on the assigned label. To satisfy (7), we dynamically sample multiple label assignments for an input graph in each epoch, but only train the GNN with the best labeling (i.e., having the lowest loss). When processing the graph in the next epoch, we re-assign node labels and lookup for a (possibly) different preferred embedding by other random permutations.

Formally, we allocate $\bm{e}_{1},\cdots,\bm{e}_{N}$ as embedding parameters in our GNN model, where $N$ is the total number of embeddings; these embeddings have not been associated with any graph or node. In each epoch when processing a graph that has nodes $V=\{v_{1},\cdots,v_{n}\}$ with $n\leq N$, we randomly sample a permutation $\pi\in S_{n}$. By the convention of our paper, $\pi$ operates on adjacency matrices; consequently, a node $v_{i}$ is now represented by $\bm{e}_{\pi^{-1}(i)}$ for GNN processing.

We repeat this sampling process $K$ times, and compute the loss of these permutations when fed to GNN. Finally, we select the permutation that has the lowest loss as the final permutation for training. In different training epochs, these permutations are re-sampled even for the same data sample.

The training process can be described as

where we train the model with the best labeling $\pi$, which is a permutation of both $X$ and $Y$ matrices (Section 2.1).

Consider the inference of a $k$-way classification problem. The GNN with permutation $\pi_{m}$ outputs a probability as $\big(p_{1}^{(m)},\cdots,p_{k}^{(m)}\big)$. We pick the prediction that has the highest joint predicted probability (product of node probabilities). Eventually, our Preferential Labeling approach predicts a label by $c=\operatorname{argmax}_{i}\{p_{i}\}$.

Our Preferential Labeling does not suffer from the limitation of equivariance, because our network is not an equivariant function as in (2). However, we will achieve the generalized equivariance property (7) asymptotically, because the labeling of $X$ is optimized out (except for auto-isomorphism) during training and inference by taking the preferred $\pi$, detailed below.

We show that the inference of our Preferential Labeling asymptotically satisfies the generalized equivariance property for node classification, if we have enough sampled permutations. We also draw a connection with Expectation–Maximization (EM) algorithms.

During our inference, we assign multiple labelings to a graph, and pick the prediction that has the highest predicted probability as our output. Formally, we consider the joint predicted probability of a graph $X$ with labeling $\tau$ as

where $f$ denotes a GNN function, outputing an $n\times k$ matrix ($n$: the number of nodes, $k$: the number of category). The $i$th row is the predicted distribution of the $i$th node.

During inference, we have multiple labelings $\tau$. We pick the best one that maximizes $s(X,\tau)$, given by

The prediction of our Preferential Labeling is

The formula follows our convention of adjacency matrix representations. When we perform node labeling, we permute the indexes of both $X$ and $Y$ by $\tau_{*}^{(X)}$. After GNN processing, we need an inverse permutation $\big(\tau_{*}^{(X)}\big)^{-1}$ to obtain the predictions for $X$, because our prediction should be corresponding to the original graph $X$, rather than $\tau_{*}^{(X)}(X)$.

Our Preferential Labeling is also related to EM algorithms.

This easy theorem provides further intuition on our Preferential Labeling approach. EM algorithms are known for handling multi-modal mixtures of distributions, similar to multi-output node classification. The training of our Preferential Labeling is also analogous to the E-step, which determines the fitness of the sample to a mixture component. Our approach adopts a hard EM variant that selects a single best permutation, because full marginalization of $S_{n}$ is intractable. Also, we assume a uniform prior for $S_{n}$, which is particularly suitable for eliminating labeling artefacts.

The static labeling assigns an embedding based on the identity of a node (e.g., $x_{1}$ and $c_{1}$ in Figure 1), although such identity does not represent meaningful semantics in different samples. Static learning is widely applied in previous work (Boldi et al. 2012; Allamanis, Brockschmidt, and Khademi 2018).

This baseline assigns all nodes in the unattributed graph with the same embedding. This is adopted in previous work (Li, Chen, and Koltun 2018; Zhang et al. 2020) to eliminate the artefacts of node labeling.

The random labeling assigns an embedding randomly during training and inference. This is a special case of our approach with $K=1$, and no actual preferential training is performed.

An intuitive way to encode a node without labeling artefact is by its degree information, which captures some local information of the node. In this baseline, we use $1/(d_{v}+1)$ as a one-dimensional, non-learnable embedding feature, where $d_{v}$ is the degree of node $v$. We use the above formula so that the feature is in $(0,1]$.

The drawbacks of using degree information as a single numerical feature are its low-dimensionality and non-learnability. In this baseline, we extend the idea by embedding the ranking of node degrees. Specifically, we sort all nodes by the degrees in descending order, and a node having $i$th largest degree is encoded by $i$th embedding vector $\bm{e}_{i}$.