SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Vignesh Kothapalli * New York University
Tom Tirer Bar-Ilan University
Joan Bruna  
New York University
Graph neural networks (GNNs) have become increasingly popular for classification tasks on graph-structured data. Yet, the interplay between graph topology and feature evolution in GNNs is not well understood. In this paper, we focus on node-wise classification, illustrated with community detection on stochastic block model graphs, and explore the feature evolution through the lens of the "Neural Collapse" (NC) phenomenon. When training instance-wise deep classifiers (e.g. for image classification) beyond the zero training error point, NC demonstrates a reduction in the deepest features' within-class variability and an increased alignment of their class means to certain symmetric structures. We start with an empirical study that shows that a decrease in within-class variability is also prevalent in the node-wise classification setting, however, not to the extent observed in the instance-wise case. Then, we theoretically study this distinction. Specifically, we show that even an "optimistic" mathematical model requires that the graphs obey a strict structural condition in order to possess a minimizer with exact collapse. Interestingly, this condition is viable also for heterophilic graphs and relates to recent empirical studies on settings with improved GNNs' generalization. Furthermore, by studying the gradient dynamics of the theoretical model, we provide reasoning for the partial collapse observed empirically. Finally, we present a study on the evolution of within- and between-class feature variability across layers of a well-trained GNN and contrast the behavior with spectral methods.
Graph neural networks [52] employ message-passing mechanisms to capture intricate topological relationships in data and have become de-facto standard architectures to handle data with non-Euclidean geometric structure [11, 12, 22, 24, 33, 60, 63, 66, 70]. However, the influence of topological information on feature learning in GNNs is yet to be fully understood [40, 64, 69, 72].
In this paper, we study the feature evolution in GNNs in a node-wise supervised classification setting. In order to gain insights into the role of topology, we focus on the controlled environment of the prominent stochastic block model (SBM) [1-3, 29, 45]. The SBM provides an effective framework to control the level of sparsity, homophily, and heterophily in the random graphs and facilitates analysis of GNN which relies solely on structural information [8, 16, 32, 40, 42, 49]. While inductive supervised learning on graphs is a relatively more difficult problem than transductive learning, it aligns with practical scenarios where nodes need to be classified in unseen graphs [24], and is also amenable to training GNNs that are deeper than conventional shallow Graph Convolution Network (GCN) models [14, 16, 32, 35, 47, 62, 62, 70].
The empirical and theoretical study of GNNs' feature evolution in this paper employs a "Neural Collapse" perspective [48]. When training Deep Neural Networks (DNNs) for classification, it
*Correspondence to: Vignesh Kothapalli (vk2115@nyu.edu)
is common to continue optimizing the networks' parameters beyond the zero training error point [9, 28, 39], a stage that was referred to in [48] as the "terminal phase of training" (TPT). Papyan, Han, and Donoho [25, 48] have empirically shown that a phenomenon, dubbed Neural Collapse (NC), occurs during the TPT of plain DNNs $^2$  on standard instance-wise classification datasets. NC encompasses several simultaneous properties: (NC1) The within-class variability of the deepest features decreases (i.e., outputs of the penultimate layer for training samples from the same class tend to their mean); (NC2) After subtracting their global mean, the mean features of different classes become closer to a geometrical structure known as a simplex equiangular tight frame; (NC3) The last layer's weights exhibit alignment with the classes' mean features. A consequence of NC1-3 is that the classifier's decision rule becomes similar to the nearest class center in the feature space. We refer to [34] for a review on this topic.
The common approach to theoretically study the NC phenomenon is the "Unconstrained Features Model" (UFM) [31, 43]. The core idea behind this "optimistic" mathematical model is that the deepest features are considered to be freely estimable. This idea has facilitated a recent surge of theoretical works in an effort to understand the global optimality conditions and gradient dynamics of these features and the last layer's weights in DNNs [18, 25, 37, 43, 54-56, 61, 65, 71, 74]. In our work, we extend NC analysis to settings where relational information in data is paramount, and creates a tension with the 'freeness' associated with the UFM model. In essence, we highlight the key differences when analyzing NC in GNNs by identifying structural conditions on the graphs, under which the global minimizers of the training objective exhibit full NC1. Interestingly, the structural conditions that we rigorously establish in this paper are aligned with the neighborhood conditions on heterophilic graphs that have been empirically hypothesized to facilitate learning by Ma et al. [40].
Our main contributions can be summarized as follows:
- We conduct an extensive empirical study that shows that a decrease in within-class variability is prevalent also in the deepest features of GNNs trained for node classification on SBMs. However, not to the extent observed in the instance-wise setting.
- We propose and analyze a graph-based UFM to understand the role of node neighborhood patterns and their community labels on NC dynamics. We prove that even this optimistic model requires a strict structural condition on the graphs in order to possess a minimizer with exact variability collapse. Then, we show that satisfying this condition is a rare event, which theoretically justifies the distinction between observations for GNNs and plain DNNs.
- Nevertheless, by studying the gradient dynamics of the graph-based UFM, we provide theoretical reasoning for the partial collapse during GNNs training.
- Finally, we study the evolution of features across the layers of well-trained GNNs and contrast the decrease in NC1 metrics along depth with a NC1 decrease along power iterations in spectral clustering methods.
2 Preliminaries and Problem Setup
We focus on supervised learning on graphs for inductive community detection. Formally, we consider a collection of  $K$  undirected graphs  $\{\mathcal{G}_k = (\mathcal{V}_k,\mathcal{E}_k)\}_{k = 1}^K$ , each with  $N$  nodes,  $C$  non-overlapping balanced communities and a node labelling ground truth function  $y_{k}:\mathcal{V}_{k}\to \{\mathbf{e}_{1},\ldots ,\mathbf{e}_{C}\}$ . Here,  $\forall c\in [C],\mathbf{e}_c\in \mathbb{R}^C$  indicates the standard basis vector, where we use the notation  $[C] = \{1,\dots ,C\}$ . The goal is to learn a parameterized GNN model  $\psi_{\Theta}(.)$  which minimizes the empirical risk given by:
$$
\min  _ {\Theta} \frac {1}{K} \sum_ {k = 1} ^ {K} \mathcal {L} \left(\psi_ {\Theta} \left(\mathcal {G} _ {k}\right), y _ {k} \left(\mathcal {V} _ {k}\right)\right) + \frac {\lambda}{2} \| \Theta \| _ {F} ^ {2}, \tag {1}
$$
where  $\| \cdot \|_F$  represents the Frobenius norm,  $\mathcal{L}$  is the loss function that is invariant to label permutations [16], and  $\lambda > 0$  is the penalty parameter. We choose  $\mathcal{L}$  based on the mean squared error (MSE) as:
$$
\mathcal {L} \left(\psi_ {\Theta} \left(\mathcal {G} _ {k}\right), y _ {k}\right) = \min  _ {\pi \in S _ {C}} \frac {1}{2 N} \| \psi_ {\Theta} \left(\mathcal {G} _ {k}\right) - \pi \left(y _ {k} \left(\mathcal {V} _ {k}\right)\right) \| _ {2} ^ {2}, \tag {2}
$$
where  $\pi$  belongs to the permutation group over  $C$  elements. Using the MSE loss for training DNN classifiers has become increasingly popular recently. For example, Hui and Belkin [30] have
2Throughout the paper, by (plain) DNNs we mean networks that output an instance-wise prediction (e.g., image class rather than pixel class), while by GNNs we mean networks that output node-wise predictions.
performed an extensive empirical study that shows that training with MSE loss yields performance that is similar to (and sometimes even better than) training with CE loss. This choice also facilitates theoretical analyses [25, 55, 71].
2.1 Data model
We employ the Symmetric Stochastic Block Model (SSBM) to generate graphs  $\{\mathcal{G}_k = (\mathcal{V}_k,\mathcal{E}_k)\}_{k = 1}^K$ . Stochastic block models (originated in [29]) are classical random graph models that have been extensively studied in statistics, physics, and computer science. In the SSBM model that is considered in this paper, each graph  $\mathcal{G}_k$  is associated with an adjacency matrix  $\mathbf{A}_k\in \mathbb{R}^{N\times N}$ , degree matrix  $\mathbf{D}_k = \mathrm{diag}(\mathbf{A}_k\mathbf{1})\in \mathbb{R}^{N\times N}$ , and a random node features matrix  $\mathbf{X}_k\in \mathbb{R}^{d\times N}$ , with entries sampled from a normal distribution. Formally, if  $\mathbf{P}\in \mathbb{R}^{C\times C}$  represents a symmetric matrix with diagonal entries  $p$  and off-diagonal entries  $q$ , a random graph  $\mathcal{G}_k$  is considered to be drawn from the distribution  $\mathrm{SSBM}(N,C,p,q)$  if an edge between vertices  $v_{i},v_{j}$  is formed with probability  $(\mathbf{P})_{y_k(v_i),y_k(v_j)}^3$ . We choose the regime of exact recovery [1-3, 45] in sparse graphs where  $p = \frac{a\ln(N)}{N}$ ,  $q = \frac{b\ln(N)}{N}$  for parameters  $a,b\geq 0$  such that  $|\sqrt{a} -\sqrt{b}| > \sqrt{C}$ . The need for exact recovery (information-theoretically) stems from the requirement that  $\psi_{\Theta}$  should be able to reach TPT (Appendix B).
2.2 Graph neural networks
Inspired by the widely studied model of higher-order GNNs by Morris et al. [44], we design  $\psi_{\Theta}$  based on a family of graph operators  $\mathcal{F} = \{\mathbf{I}, \widehat{\mathbf{A}}_k\}$ ,  $\forall k \in [K]$ , and denote it as  $\psi_{\Theta}^{\mathcal{F}}$ . Formally, for a GNN  $\psi_{\Theta}^{\mathcal{F}}$  with  $L$  layers, the node features  $\mathbf{H}_k^{(l)} \in \mathbb{R}^{d_l \times N}$  at layer  $l \in [L]$  is given by:
$$
\mathbf {X} _ {k} ^ {(l)} = \mathbf {W} _ {1} ^ {(l)} \mathbf {H} _ {k} ^ {(l - 1)} + \mathbf {W} _ {2} ^ {(l)} \mathbf {H} _ {k} ^ {(l - 1)} \widehat {\mathbf {A}} _ {k}, \tag {3}
$$
$$
\mathbf {H} _ {k} ^ {(l)} = \sigma (\mathbf {X} _ {k} ^ {(l)}),
$$
where  $\mathbf{H}_k^{(0)} = \mathbf{X}_k$ , and  $\sigma(\cdot)$  represents a point-wise activation function such as ReLU.  $\mathbf{W}_1^{(l)}, \mathbf{W}_2^{(l)} \in \mathbb{R}^{d_l \times d_{l-1}}$  are the weight matrices and  $\widehat{\mathbf{A}}_k = \mathbf{A}_k \mathbf{D}_k^{-1}$  is the normalized adjacency matrix, also known as the random-walk matrix. We also consider a simpler family without the identity operator  $\mathcal{F}' = \{\widehat{\mathbf{A}}_k\}, \forall k \in [K]$  and analyze the GNN  $\psi_\Theta^{\mathcal{F}'}$  with only graph convolution functionality. Formally, the node features  $\mathbf{H}_k^{(l)} \in \mathbb{R}^{d_l \times N}$  for  $\psi_\Theta^{\mathcal{F}'}$  is given by:
$$
\mathbf {X} _ {k} ^ {(l)} = \mathbf {W} _ {2} ^ {(l)} \mathbf {H} _ {k} ^ {(l - 1)} \widehat {\mathbf {A}} _ {k}, \tag {4}
$$
$$
\mathbf {H} _ {k} ^ {(l)} = \sigma (\mathbf {X} _ {k} ^ {(l)}).
$$
Here, the subscript for the weight matrix  $\mathbf{W}_2^{(l)}$  is retained to highlight that it acts on  $\mathbf{H}_k^{(l-1)}\widehat{\mathbf{A}}_k$ . Finally, we employ the training strategy of Chen et al. [16] and apply instance-normalization [58] on  $\sigma(\mathbf{X}_k^{(l)}), \forall l \in \{1, \dots, L-1\}$  to prevent training instability.
2.3 Tracking neural collapse in GNNs
In our setup, reaching zero training error (TPT) implies that the network perfectly classifies all the nodes (up to label permutations) in all the training graphs. To this end, we leverage the NC metrics introduced in [48, 55, 56, 74] and extend them to GNNs in an inductive setting. To begin with, let us consider a single graph  $\mathcal{G}_k = (\mathcal{V}_k,\mathcal{E}_k), k\in [K]$  with a normalized adjacency matrix  $\hat{\mathbf{A}}_k$ . Additionally, we denote  $\mathbf{H}_k^{(l)}\in \mathbb{R}^{d_l\times N}$  as the output of layer  $l\in [L - 1]$ , irrespective of the GNN design. Now, by dropping the subscript and superscript for notational convenience, we define the class means and the global mean of  $\mathbf{H}$  as follows:
$$
\overline {{\mathbf {h}}} _ {c} := \frac {1}{n} \sum_ {i = 1} ^ {n} \mathbf {h} _ {c, i}, \forall c \in [ C ], \quad \overline {{\mathbf {h}}} _ {G} := \frac {1}{C n} \sum_ {c = 1} ^ {C} \sum_ {i = 1} ^ {n} \mathbf {h} _ {c, i}, \tag {5}
$$
where  $n = N / C$  represents the number of nodes in each of the  $C$  balanced communities, and  $\mathbf{h}_{c,i}$  is the feature vector (a column in  $\mathbf{H}$ ) associated with  $v_{c,i} \in \mathcal{V}$ , i.e., the  $i^{th}$  node belonging to class
In our setup, the nodes of sampled SSBM graphs are allowed to have self-edges.
$c \in [C]$ . Next, let  $\mathcal{N}(v_{c,i})$  denote all the neighbors of  $v_{c,i}$  and let  $\mathcal{N}_{c'}(v_{c,i})$  denote only the neighbors of  $v_{c,i}$  that belong to class  $c' \in [C]$ . We define the class means and global mean of  $\mathbf{H}\widehat{\mathbf{A}}$ , which is unique to the GNN setting as follows:
$$
\overline {{\mathbf {h}}} _ {c} ^ {\mathcal {N}} := \frac {1}{n} \sum_ {i = 1} ^ {n} \mathbf {h} _ {c, i} ^ {\mathcal {N}}, \forall c \in [ C ], \quad \overline {{\mathbf {h}}} _ {G} ^ {\mathcal {N}} := \frac {1}{C n} \sum_ {c = 1} ^ {C} \sum_ {i = 1} ^ {n} \mathbf {h} _ {c, i} ^ {\mathcal {N}}, \tag {6}
$$
where  $\mathbf{h}_{c,i}^{\mathcal{N}} = \left(\sum_{v_{c,j}\in \mathcal{N}_c(v_{c,i})}\mathbf{h}_{c,j} + \sum_{v_{c',j}\in \mathcal{N}_{c'\neq c}(v_{c,i})}\mathbf{h}_{c',j}\right) / |\mathcal{N}(v_{c,i})|$ .
- Variability collapse in features  $\mathbf{H}$ : For a given features matrix  $\mathbf{H}$ , let us define the within- and between-class covariance matrices,  $\Sigma_W(\mathbf{H})$  and  $\Sigma_B(\mathbf{H})$ , as:
$$
\boldsymbol {\Sigma} _ {W} (\mathbf {H}) := \frac {1}{C n} \sum_ {c = 1} ^ {C} \sum_ {i = 1} ^ {n} \left(\mathbf {h} _ {c, i} - \overline {{\mathbf {h}}} _ {c}\right) \left(\mathbf {h} _ {c, i} - \overline {{\mathbf {h}}} _ {c}\right) ^ {\top}, \tag {7}
$$
$$
\boldsymbol {\Sigma} _ {B} (\mathbf {H}) := \frac {1}{C} \sum_ {c = 1} ^ {C} \left(\overline {{\mathbf {h}}} _ {c} - \overline {{\mathbf {h}}} _ {G}\right) \left(\overline {{\mathbf {h}}} _ {c} - \overline {{\mathbf {h}}} _ {G}\right) ^ {\top}. \tag {8}
$$
To empirically track the within-class variability collapse with respect to the between-class variability, we define two NC1 metrics:
$$
\mathcal {N C} _ {1} (\mathbf {H}) = \frac {1}{C} \operatorname {T r} \left(\boldsymbol {\Sigma} _ {W} (\mathbf {H}) \boldsymbol {\Sigma} _ {B} ^ {\dagger} (\mathbf {H})\right), \quad \widetilde {\mathcal {N C}} _ {1} (\mathbf {H}) = \frac {\operatorname {T r} \left(\boldsymbol {\Sigma} _ {W} (\mathbf {H})\right)}{\operatorname {T r} \left(\boldsymbol {\Sigma} _ {B} (\mathbf {H})\right)}, \tag {9}
$$
where  $\dagger$  denotes the Moore-Penrose pseudo-inverse and  $\mathrm{Tr}(\cdot)$  denotes the trace of a matrix. Although  $\mathcal{NC}_1$  is the original NC1 metric used by Papyan et al. [48], we consider also  $\widetilde{\mathcal{NC}}_1$ , which has been proposed by Tirer et al. [56] as an alternative metric that is more amenable to theoretical analysis.
- Variability collapse in neighborhood-aggregated features  $\mathbf{H}\widehat{\mathbf{A}}$ : Similarly to the above, we track the within- and between-class variability of the "neighborhood-aggregated" features matrix  $\mathbf{H}\widehat{\mathbf{A}}$  by  $\boldsymbol{\Sigma}_{W}(\mathbf{H}\widehat{\mathbf{A}})$  and  $\boldsymbol{\Sigma}_{B}(\mathbf{H}\widehat{\mathbf{A}})$  (computed using  $\overline{\mathbf{h}}_c^{\mathcal{N}}$  and  $\overline{\mathbf{h}}_G^{\mathcal{N}}$ ), as well as  $\mathcal{NC}_1(\mathbf{H}\widehat{\mathbf{A}})$  and  $\widetilde{\mathcal{NC}}_1(\mathbf{H}\widehat{\mathbf{A}})$ . (See Appendix C for formal definitions.) Finally, we follow a simple approach and track the mean and variance of  $\mathcal{NC}_1(\mathbf{H}), \widetilde{\mathcal{NC}}_1(\mathbf{H}), \mathcal{NC}_1(\mathbf{H}\widehat{\mathbf{A}}), \widetilde{\mathcal{NC}}_1(\mathbf{H}\widehat{\mathbf{A}})$  across all  $K$  graphs in our experiments.
As the primary focus of our paper is the analysis of feature variability during training and inference, we defer the definition and examination of metrics based on NC2 and NC3 to Appendix C, H.
3 Evolution of penultimate layer features during training
In this section, we explore the evolution of the deepest features of GNNs during training. In Section 3.1, we present empirical results of GNNs in the setup that is detailed in Section 2, showing that a decrease in within-class feature variability is present in GNNs that reach zero training error, but not to the extent observed with plain DNNs. Then, in Section 3.2, we theoretically study a mathematical model that provides reasoning for the empirical observations.
3.1 Experiments
Setup. We focus on the training performance of GNNs  $\psi_{\Theta}^{\mathcal{F}}, \psi_{\Theta}^{\mathcal{F}'}$  on sparse graphs and generate a dataset of  $K = 1000$  random SSBM graphs with  $C = 2, N = 1000, p = 0.025, q = 0.0017$ . The networks  $\psi_{\Theta}^{\mathcal{F}}, \psi_{\Theta}^{\mathcal{F}'}$  are composed of  $L = 32$  layers with graph operator, ReLU activation, and instance-normalization functionality. The hidden feature dimension is set to 8 across layers. They are trained for 8 epochs using stochastic gradient descent (SGD) with a learning rate 0.004, momentum 0.9, and a weight decay of  $5 \times 10^{-4}$ . During training, we track the NC1 metrics for the penultimate layer features  $\mathbf{H}_k^{(L-1)}$ , by computing their mean and standard deviation across  $k \in [K]$  graphs after every epoch. To measure the performance of the GNN, we compute the 'overlap' [16] between predicted communities and ground truth communities (up to permutations):
$$
\operatorname {o v e r l a p} (\hat {y}, y) := \max  _ {\pi \in S _ {C}} \left(\frac {1}{N} \sum_ {i = 1} ^ {N} \delta_ {\hat {y} (v _ {i}), \pi (y (v _ {i}))} - \frac {1}{C}\right) / \left(1 - \frac {1}{C}\right) \tag {10}
$$
(b) overlap
(c) NC1: H
(d) NC1:  $\hat{\mathbf{H}}\hat{\mathbf{A}}$
Figure 1: GNN  $\psi_{\Theta}^{\mathcal{F}}$ : Illustration of loss, overlap, and  $\mathcal{NC}_1$  plots for  $\mathbf{H}$ ,  $\mathbf{H}\hat{\mathbf{A}}$  during training.
(a) loss
(b) overlap
Figure 2: GNN  $\psi_{\Theta}^{\mathcal{F}'}$ : Illustration of loss, overlap, and  $\mathcal{NC}_1$  plots for  $\mathbf{H}$ ,  $\mathbf{H}\hat{\mathbf{A}}$  during training.
where  $\hat{y}$  is the node labelling function based on GNN design and  $\frac{1}{N}\sum_{i=1}^{N}\delta_{\hat{y}(v_i),\pi(y(v_i))}$  is the training accuracy ( $\delta$  denotes the Kronecker delta). The overlap allows us to measure the improvements in performance over random guessing while retaining the indication that the GNN has reached TPT. Formally, when  $\frac{1}{N}\sum_{i=1}^{N}\delta_{\hat{y}(v_i),\pi(y(v_i))} = 1$  (zero training error), then  $\mathrm{overlap}(\hat{y},y) = 1$ . We illustrate the empirical results in Figures 1 and 2, and present extensive experiments (showing similar behavior) along with infrastructure details in Appendix  $\mathbf{H}^4$ .
Observation: The key takeaway is that  $\mathcal{NC}_1(\mathbf{H}_k^{(L - 1)})$ ,  $\widetilde{\mathcal{NC}_1} (\mathbf{H}_k^{(L - 1)})$  tend to reduce and plateau during TPT in  $\psi_{\Theta}^{\mathcal{F}}$  and  $\psi_{\Theta}^{\mathcal{F}'}$ . Notice that even though we consider a controlled SSBM-based setting, the  $\mathcal{NC}_1$  values observed here are higher than the values observed in the case of plain DNNs on real-world instance-wise datasets [48, 74]. Additionally, we can observe that trends for  $\mathcal{NC}_1(\mathbf{H}_k^{(L - 1)}\widehat{\mathbf{A}}_k)$ ,  $\widetilde{\mathcal{NC}_1} (\mathbf{H}_k^{(L - 1)}\widehat{\mathbf{A}}_k)$  are similar to those of  $\mathcal{NC}_1(\mathbf{H}_k^{(L - 1)})$ ,  $\widetilde{\mathcal{NC}_1} (\mathbf{H}_k^{(L - 1)})$ .
3.2 Theoretical analysis
In this section, we provide a theory for this empirical behavior. Most, if not all, of the theoretical papers on NC, adopt the UFM approach, which treats the features as free optimization variables – disconnected from data [18, 25, 43, 55, 56, 74]. Here, we consider a graph-based adaptation of this approach, that we dubbed as gUFM. We consider GNNs of the form of  $\psi_{\Theta}^{\mathcal{F}'}$ , which is more tractable for mathematical analysis. Formally, by considering  $\mathcal{L}$  to be the MSE loss, treating  $\{\mathbf{H}_k^{(L-1)}\}_{k=1}^K$  as freely estimizable variables, and representing  $\mathbf{W}_2^{(L)} \in \mathbb{R}^{C \times d_{L-1}}$ ,  $\mathbf{H}_k^{(L-1)} \in \mathbb{R}^{d_{L-1} \times N}$  as  $\mathbf{W}_2$ ,  $\mathbf{H}_k$  (for notational convenience), the empirical risk based on the gUFM can be formulated as follows:
$$
\widehat {\mathcal {R}} ^ {\mathcal {F} ^ {\prime}} \left(\mathbf {W} _ {2}, \left\{\mathbf {H} _ {k} \right\} _ {k = 1} ^ {K}\right) := \frac {1}{K} \sum_ {k = 1} ^ {K} \left(\frac {1}{2 N} \left\| \mathbf {W} _ {2} \mathbf {H} _ {k} \widehat {\mathbf {A}} _ {k} - \mathbf {Y} \right\| _ {F} ^ {2} + \frac {\lambda_ {H _ {k}}}{2} \| \mathbf {H} _ {k} \| _ {F} ^ {2}\right) + \frac {\lambda_ {W _ {2}}}{2} \| \mathbf {W} _ {2} \| _ {F} ^ {2} \tag {11}
$$
where  $\mathbf{Y} \in \mathbb{R}^{C \times N}$  is the target matrix, which is composed of one-hot vectors associated with the different classes, and  $\lambda_{W_2}, \lambda_{H_k} > 0$  are regularization hyperparameters. To simplify the analysis, let us assume that  $\mathbf{Y} = \mathbf{I}_C \otimes \mathbf{1}_n^\top$ , where  $\otimes$  denotes the Kronecker product. Namely, the training data is balanced (a common assumption in UFM-based analyses in literature) with  $n = N / C$  nodes per
$^{4}$ Code is available at: https://github.com/kvignesh1420/gnn_collapse
class in each graph and (without loss of generality) organized class-by-class. Note that for  $K = 1$  (which allows omitting the graph index  $k$ ) and no graphical structure, i.e.,  $\widehat{\mathbf{A}} = \mathbf{I}$  (since  $\mathbf{A} = \mathbf{I}$ ), (11) reduces to the plain UFM that has been studied in [25, 55, 71]. In this case, it has been shown that any minimizer  $(\mathbf{W}_2^*, \mathbf{H}^*)$  is collapsed, i.e., its features have exactly zero within-class variability:
$$
\mathbf {h} _ {c, 1} ^ {*} = \dots = \mathbf {h} _ {c, n} ^ {*} = \overline {{\mathbf {h}}} _ {c} ^ {*}, \quad \forall c \in [ C ], \tag {12}
$$
which implies  $\Sigma_W(\mathbf{H}^*) = \mathbf{0}$ . We will show now that the situation in gUFM is significantly different.
Considering the  $K = 1$  case, we start by showing that, to have minimizers of (11) that possess the property in (12), the graph must obey a strict structural condition. For  $K > 1$ , having a minimizer  $(\mathbf{W}_2^*, \{\mathbf{H}_k^*\})$  where, for some  $j \in [K]$ ,  $\mathbf{H}_j^*$  is collapsed directly follows from having the structural condition satisfied by the  $j$ -th graph (as shown in our proof, the sufficiency of the condition does not depend on the shared weights  $\mathbf{W}_2$ ). On the other hand, generalizing the necessity of the structural condition to the case of  $K > 1$  is technically challenging (see the appendix for details). For that reason, we state the condition in the following theorem only for  $K = 1$ . Note also that, showing that the condition is unlikely to be satisfied per graph is enough for explaining the plateaus above zero of NC metrics (computed over multiple graphs), which are demonstrated in Section 3.1.