Geometric instability of graph neural networks on large graphs
We analyse the geometric instability of embeddings produced by graph neural
networks (GNNs). Existing methods are only applicable for small graphs and lack
context in the graph domain. We propose a simple, efficient and graph-native
Graph Gram Index (GGI) to measure such instability which is invariant to
permutation, orthogonal transformation, translation and order of evaluation.
This allows us to study the varying instability behaviour of GNN embeddings on
large graphs for both node classification and link prediction.
Emily Morris, Haotian Shen, Weiling Du, Muhammad Hamza Sajjad, Borun Shi
# Geometric instability of graph neural networks on large graphs

###### Abstract

We analyse the geometric instability of embeddings produced by graph neural networks (GNNs). Existing methods are only applicable for small graphs and lack context in the graph domain. We propose a simple, efficient and graph-native Graph Gram Index (GGI) to measure such instability which is invariant to permutation, orthogonal transformation, translation and order of evaluation. This allows us to study the varying instability behaviour of GNN embeddings on large graphs for both node classification and link prediction.

## 1 Introduction

Graph representation learning [1] has seem many recent success in solving problems over relational data. The importance of stochastic effects is evident in graph learning. Earlier methods such as DeepWalk [2] and node2vec [3], commonly known as shallow embedding, leverage random walk. Many recent Graph Neural Network (GNN)[4] models contain stochastic components such as sampling and batching. From an implementation point of view, this allows training models on large real-world graphs. From a theoretical point of view, adding randomness is proven to improve expressiveness and performance [5][6]. The impact of randomness common in standard machine learning frameworks also carries over to the graph domain, such as in weight initialization, gradient descent, dataset splits[7] etc. It is therefore important to understand how all the stochastic components together affect graph learning models at a granular level.

Stability of embeddings and models have been extensively studied outside of the domain of graphs[8][9][10][11]. Yet the practical impact of randomness on GNNs has been very rarely studied. [12] recently empirically showed that final predictions given by GCN [13] and GAT[14] vary significantly on individual nodes despite relatively stable overall accuracies. At a finer level, [15] showed the instability in embeddings produced by several shallow methods and GraphSAGE [16] and [17] specifically for node2vec. To the best of our knowledge, all previous studies are carried out on small graphs for node classification and stability indices from other domains are directly borrowed.

Additional motivations of understanding embedding geometric stability include improving reproducibility and model reliability [11], reducing retraining effort when embeddings are used by multiple downstream systems by measuring drift [18], and understanding how a model works by studying its embedding space.

The main contributions of this paper are:

* We formalise geometric stability index for embeddings and examine existing methods.
* We propose a simple, time and space efficient and graph-native stability index, the Graph Gram index (GGI). We motivate and show that GGI is invariant to permutation, orthogonal transformation, translation and the order of evaluation.
Section 4: We show the geometric instability of embeddings of several popular baseline GNNs on large graphs for both node classification and link prediction.

## 2 Geometric stability

**Notation.** Let \(G=(V,E)\) denote a graph, \(|V|\) the number of nodes, \(|E|\) the number of relationships. Each node has a input feature in \(\mathbb{R}^{n}\), and a final embedding in \(\mathbb{R}^{d}\). For a given node at position \(i\) in some ordered list, we refer to that node as \(v_{i}\). We refer to the final embedding of that node as \(z_{i}\in\mathbb{R}^{1\times d}\). Whenever we need to refer to the embedding of some node \(v_{i}\) by name, we will abuse the notation slightly as \(z_{v_{i}}\). Let \(Z\in\mathbb{R}^{|V|\times d}\) denote the embeddings of all nodes in the graph. Since we are concerned with multiple embeddings produced from different configurations, let superscript \(Z^{k}\) denotes the embeddings produced from configuration \(k\), unless otherwise stated. Let \(N\) denote the overall number of configurations.

Stability indices that have been used to evaluate the geometric stability of embeddings of graph models are either borrowed from the NLP domain [15], or from analytic topology [17]. We give the definitions of one similarity index from each area and refer to the remaining ones in Appendix A.

**Second-order cosine.** Given two embedding matrices \(Z^{l}\), \(Z^{m}\). For each node \(v_{i}\), defined an ordered list \(\{u_{1},\ldots,u_{K}\}=\mathcal{N}^{l}_{k}(v_{i})\cup\mathcal{N}^{m}_{k}(v_{i})\) where \(\mathcal{N}^{l}_{k}(v_{i})\) is the k nearest neighbour of node \(i\) under configuration \(l\). Let \(s^{l}(v_{i})\) denote a vector about \(v_{i}\) where the \(j^{th}\) entry is \(s^{l}_{j}(v_{i})=cossim(z^{l}_{i},z^{l}_{u_{j}})\). Similarly we can define \(s^{m}(v_{i})\). Let \(s^{cos[l,m]}_{k}(z^{l}_{i},z^{m}_{i})=cossim(s^{l}_{v_{i}},s^{m}_{v_{i}})\). Cosine distance are sometimes equivalently used. The second order cosine similarity of \(N\) configurations is then averaging over all nodes and all configuration pairs, i.e \(\frac{1}{N^{2}\times|\mathcal{V}|}\sum_{l,m\in N\times N}\sum_{i\in V}s^{cos[l, m]}_{k}(z^{l}_{i},z^{m}_{i})\). This stability index was originally used to detect words semantic shifts over time [19].

**Wasserstein distance.** Given two embedding matrices \(Z^{l}\), \(Z^{m}\), the Wasserstein distance between them is \(W(Z^{l},Z^{m})=(\inf_{\eta:V\sim V}\sum_{j\in V}\lVert z_{i}-z_{\eta(i)} \rVert_{2}^{2})^{1/2}\). Where \(\inf_{\eta:V\to V}\) is over all bijections between nodes. The average across all possible configuration pairs are taken as the stability index.

While these indices have shown preliminary instability results, there are several improvements to be made.

First all of, most of the existing similarity indices are inefficient to compute, evident with the common _averaging over all nodes and all configuration pairs_. Finding k nearest neighbours (kNN) itself is a costly operation. While there are performant kNN libraries[20], employing them adds additional engineering cost. Opting for approximate kNN speed up the process[21], at the cost of introducing confounding randomness which makes it difficult to draw concrete conclusions on the embeddings. Several of the neighbour-based indices contain hyperparameters of their own (\(k\)), which for the same reason should ideally be avoided. We give the full time and space complexity in Appendix B.

For any index (such as aligned cosine) that directly compares embeddings across configurations (i.e to compare \(z^{l}_{i}\), \(z^{m}_{i}\)), an alignment step such as the solving the orthogonal Procrustes problem[22] (definition in Appendix C) needs to be done first. While embedded node neighbours can still carry important information, the interpretation of semantically similar words no longer applies. Indices from analytic topology (such as Wasserstein distance) are originally used to measure the overlap of probability distributions over metric space, which again does not map to any natural interpretation for graph embedding. Same drawback applies for Hausdorff distance (we elaborate more in Appendix D ).

We propose that a stability index for embeddings produced by graph machine learning models should ideally be time and space efficient to compute, free of hyperparameters of its own, and have an intuitive and simple interpretation in the context of graphs. Having such an index enables easier analysis of embeddings produced by graph models on large real-world graphs.

## 3 Graph Gram Index

We propose a simple index called the Graph Gram Index (GGI).

\(Z^{l}Z^{l}{}^{T}\) is indeed the Gram matrix, which contains all covariance information between \(z^{l}_{i}\), \(z^{l}_{j}\). Using the Gram matrix avoids solving the Procrustes alignment step which would be required if we calculate
\(z_{i}^{l}z_{j}^{m}\). Intuitively, this is achieved because we _compare within one set of embeddings to generate a structure summary, and compare the summaries across configurations._ The commutative diagram gives a pictorial illustration. All previous methods follow the top-right path while we follows the bottom-left.

An analogous approach underpins Centered Kernel Alignment, a _similarity_ measure to compare internal representations across general neural networks (of any domain). We refer to Appendix E for a definition of CKA and how it inspires our method.

Applying hadamard product of adjacency matrix over the Gram matrix means that we only capture the node pairs that are actual edges in the graph. Intuitively, a stable model produces embeddings whose node pairs that correspond to an edge should remain similar, in their respective embedding spaces, _up to any equivariant transformation satisfied by the specific model_. Note that a node pair corresponding to an edge could potentially lie far apart in the embedding space, which previous stability indices would not consider.

The framework of GGI is flexible enough to allow various extension, for example by using non-Gram matrix based summary structure, using other the \(S^{l}\to s^{l}\) pooling step such as the Frobenius norm (notably, this step is the same as many graph pooling steps), and replacing std with other summary statistics. We leave them as future work.

### Invariance properties

To the best of our knowledge, there is no prior study about the geometric invariances a stability index should satisfy. It is important to define such invariances, because the meaningful geometric instability any index captures should be the ones that consumers of the embeddings (final GNN layer or downstream model) are not equivariant to. At the same time, it should ideally not be only invariant to an overly specific set of transformations that few models satisfy.

Hence, we propose that it is desirable to be invariant to node permutation, orthogonal transformation and translation. In addition it should be invariant to the order of evaluating a given set of configurations. kNN-Jaccard and second-order cossim for example are not invariant to permutation due to it's own nondeterminism from sampling.

**Lemma**.: _GGI is invariant to node permutation, orthogonal transformation and translation of embeddings, and order of evaluation._

Proof follows from simple applications of properties of matrix multiplication and isometries which we include in Appendix F.
## 4 Experiments

Code for the experiments and our implementation of GGI is available on github.2

Footnote 2: [https://github.com/brs96/geometric-instability-gnn-large-graphs](https://github.com/brs96/geometric-instability-gnn-large-graphs)

**Node classification.** We use a machine with 60GB of vRAM and 8 vCPU, with.one NVIDIA T4 GPU. We evaluate four popular baseline GNN models, GCN, GraphSAGE, GAT, GIN[23]. We use Cora as a small baseline comparison with existing methods, and additional calculate Graph Gram Index on ogbn-arxiv[24]. For each model, we train for 30 different random seeds with embedding dimension 64 on Cora and 128 on ogbn-arxiv. Further training details in Appendix H.

Computing second order cossim on Cora takes around 20 minutes, with 20 nearest neighbours for all nodes precomputed already. On the other hand, computing GGI across 30 sets of embeddings takes only **10 seconds** using the single GPU.

The GGI indices for the four models are _1.7%, 3%, 4% and 6.7%_ respectively on Cora, and _0.7%, 1.2%, 1.6% and 3.1%_ on ogbn-arxiv. We take advantage the flexibility of the final step for GGI and replace taking standard deviation by taking the box plot. GraphSAGE and GCN produces embeddings whose edge source and target nodes are assigned more orthogonal vectors, as illustrated in 1(a), 1(b). Box plots for 20NN-Jaccard and second order cossim give contradictory results between themselves which we include in Appendix G, which further shows the difficulty of drawing consistent conclusions based on these methods.

It takes close to one hour to simply compute 20NN-Jaccard for one configuration, and prohibitively long for second order cossim on ogbn-arxiv. It however only take **less than 20 minutes** to calculate GGI with peak GPU vRAM less than **10 GB**, doable on any commodity machine.

**Link prediction** To the best of our knowledge there's no previous literature studying the geometric stability of embeddings produced by link prediction models. We evaluate batched versions of GraphSAGE and GCN (with GraphSAINT sampling[25]) on ogbl-citation2[24] with embedding dimension 128. Standard training setup is used as detailed in Appendix H.

Footnote 2: [https://github.com/brs96/geometric-instability-gnn-large-graphs](https://github.com/brs96/geometric-instability-gnn-large-graphs)

We observe GGI for batched GraphSAGE and GCN to be _2.5%, 2.6%_ respectively. Similar box plot is shown in Appendix E. Interestingly, both GGI index and box-plot statistics align with the observations we see during full-batch training of respective models, on different datasets and even tasks. it hence indicates the embedding geometric stability as well as the shape of the learnt point clouds seem to be inherent to the model architecture.

## 5 Conclusion and future work

Our results indicate there seems to be a true level of embedding geometric stability produced by one model architecture. GGI and the framework behind it offers a convenient tool for researchers and practitioners to evaluate embedding geometric stability on large graphs. One could leverage the flexibility to adopt suitable indices that capture desirable geometric properties in the embedding space. Implementation of GGI can be easily further improved by leveraging sparse-dense multiplication[26] (for \(A\circ Gram\)) and split multiplications into chunks to fit arbitrarily large graphs into parallelizable GPU workload.

Figure 1: Box plots of stability measures for Cora.