MuseGNN: Interpretable and Convergent Graph Neural Network Layers at
  Scale
Among the many variants of graph neural network (GNN) architectures capable
of modeling data with cross-instance relations, an important subclass involves
layers designed such that the forward pass iteratively reduces a
graph-regularized energy function of interest. In this way, node embeddings
produced at the output layer dually serve as both predictive features for
solving downstream tasks (e.g., node classification) and energy function
minimizers that inherit desirable inductive biases and interpretability.
However, scaling GNN architectures constructed in this way remains challenging,
in part because the convergence of the forward pass may involve models with
considerable depth. To tackle this limitation, we propose a sampling-based
energy function and scalable GNN layers that iteratively reduce it, guided by
convergence guarantees in certain settings. We also instantiate a full GNN
architecture based on these designs, and the model achieves competitive
accuracy and scalability when applied to the largest publicly-available node
classification benchmark exceeding 1TB in size.
Haitian Jiang, Renjie Liu, Xiao Yan, Zhenkun Cai, Minjie Wang, David Wipf
# MuseGNN: Interpretable and Convergent Graph Neural Network Layers at Scale

###### Abstract

Among the many variants of graph neural network (GNN) architectures capable of modeling data with cross-instance relations, an important subclass involves layers designed such that the forward pass iteratively reduces a graph-regularized energy function of interest. In this way, node embeddings produced at the output layer dually serve as both predictive features for solving downstream tasks (e.g., node classification) and energy function minimizers that inherit desirable inductive biases and interpretability. However, scaling GNN architectures constructed in this way remains challenging, in part because the convergence of the forward pass may involve models with considerable depth. To tackle this limitation, we propose a sampling-based energy function and scalable GNN layers that iteratively reduce it, guided by convergence guarantees in certain settings. We also instantiate a full GNN architecture based on these designs, and the model achieves competitive accuracy and scalability when applied to the largest publicly-available node classification benchmark exceeding 1TB in size.

## 1 Introduction

Graph neural networks (GNNs) are a powerful class of deep learning models designed specifically for graph-structured data. Unlike conventional neural networks that primarily operate on independent samples, GNNs excel in capturing the complex cross-instance relations modeled by graphs (Hamilton et al., 2017; Kearnes et al., 2016; Kipf and Welling, 2017; Velickovic et al., 2018). Foundational to GNNs is the notion of message passing (Gilmer et al., 2017), whereby nodes iteratively gather information from neighbors to update their representations. In doing so, information can propagate across the graph in the form of node embeddings, reflecting both local patterns and global network effects, which are required by downstream tasks such as node classification.

Among many GNN architectures, one promising subclass is based on graph propagation layers derived to be in a one-to-one correspondence with the descent iterations of a graph-regularized energy function (Chen et al., 2022; Ahn et al., 2022; Yang et al., 2021; Ma et al., 2021; Gasteiger et al., 2019; Pan et al., 2020; Zhang et al., 2020; Zhu et al., 2021; Xue et al., 2023). For these models, the layers of the GNN forward pass computes increasingly-refined approximations to a minimizer of the aforementioned energy. Importantly, if such energy minimizers possess interpretable/explainable properties or inductive biases, the corresponding GNN architecture naturally inherits them, unlike traditional GNN constructions that may be less transparent. More broadly, this association between the GNN forward pass and optimization can be exploited to introduce targeted architectural enhancements (e.g., forming explainable predictions, robustly handling graphs with spurious edges, etc.). Borrowing from (Chen and Eldar, 2021; Yang et al., 2022), we
will henceforth refer to models of this genre as _unfolded_ GNNs, given that the layers are derived from a so-called unfolded (in time/iteration) energy descent process.

Despite their merits w.r.t. interpretability, unfolded GNNs face non-trivial scalability challenges. This is in part because they can be constructed with arbitrary depth (i.e., arbitrary descent iterations) while still avoiding undesirable oversmoothing effects (Oono and Suzuki, 2020; Li et al., 2018), and the computational cost and memory requirements of this flexibility are often prohibitively high. Even so, there exists limited prior work explicitly tackling the scalability of unfolded GNNs, or providing any complementary guarantees regarding convergence on large graphs. Hence most unfolded GNN models are presently evaluated on relatively small benchmarks.

To address these shortcomings, we propose a scalable unfolded GNN model that incorporates efficient subgraph sampling within the fabric of the requisite energy function. Our design of this model is guided by the following three desiderata: (i) maintain the characteristic interpretable properties and extensible structure of full-graph unfolded GNN models, (ii) using a consistent core architecture, preserve competitive accuracy across datasets of varying size, including the very largest publicly-available benchmarks, and (iii) do not introduce undue computational or memory overheads beyond what is required to train the most common GNN alternatives. Our proposed model, which we will later demonstrate satisfies each of the above, is termed _MuseGNN_ in reference to a GNN architecture produced by the _minimization_ of an _unfolded sampling_-based _energy_. Building on background motivations presented in Section 2, our MuseGNN-related contributions herein are three-fold:

* In Sections 3 and 4, we expand a widely-used unfolded GNN framework to incorporate offline sampling into the architecture-inducing energy function design itself, as opposed to a post hoc application of sampling to existing GNN methods. The resulting model, which we term MuseGNN, allows us to retain the attractive interpretability and extensibility properties of unfolded GNNs, with the scalability of canonical GNN architectures.
* We analytically demonstrate in Section 5 that MuseGNN possesses desirable convergence properties regarding both upper-level (traditional model training) and lower-level (interpretable energy function descent) optimization processes. In doing so we increase our confidence in reliable performance when moving to new problem domains that may deviate from known benchmarks.
* Finally, in Section 6 we provide complementary empirical support that MuseGNN performance is stable in practice, preserving competitive accuracy and scalability across task size. En route, we achieve SOTA performance w.r.t. homogeneous graph models applied to the largest, publicly-available node classification datasets from OGB and IGB exceeding 1TB.

## 2 Background and Motivation

This section provides background details regarding existing unfolded GNN architectures. Later, we discuss why GNNs formed in this way are relevant, followed by their scalability challenges.

### Interpretable GNN Architectures from Unfolded Optimization

Notation.Let \(\mathcal{G}=\{\mathcal{V},\mathcal{E}\}\) denote a graph with \(n=|\mathcal{V}|\) nodes and edge set \(\mathcal{E}\). We define \(D\) and \(A\) as the degree and adjacency matrices of \(\mathcal{G}\) such that the corresponding graph Laplacian is \(L=D-A\). Furthermore, associated with each node is both a \(d\)-dimensional feature vector, and a \(d^{\prime}\)-dimensional label vector, the respective collections of which are given by \(X\in\mathbb{R}^{n\times d}\) and \(T\in\mathbb{R}^{n\times d^{\prime}}\).

GNN Basics.Given a graph defined as above, canonical GNN architectures are designed to produce a sequence of node-wise embeddings \(\{Y^{(k)}\}_{k=1}^{K}\) that are increasingly refined across \(K\) model layers according to the rule \(Y^{(k)}=h(Y^{(k-1)};W,A,X)\in\mathbb{R}^{n\times d}\). Here \(h\) represents a function that updates \(Y^{(k-1)}\) based on trainable model weights \(W\) as well as \(A\) (graph structure) and optionally \(X\) (input node features). To facilitate downstream tasks such as node classification, \(W\) may be trained, along with additional parameters \(\theta\) of any application-specific output layer \(g:\mathbb{R}^{d}\rightarrow\mathbb{R}^{d^{\prime}}\), to minimize a loss of the form
\[\mathcal{L}(W,\theta)=\sum_{i=1}^{n^{\prime}}\mathcal{D}\bigg{(}g\left[Y^{(K)}\left( W\right)_{i};\theta\right],T_{i}\bigg{)}. \tag{1}\]

In this expression, \(Y_{i}^{(K)}\equiv Y^{(K)}\left(W\right)_{i}\) reflects the explicit dependency of GNN embeddings on \(W\) and the subscript \(i\) denotes the \(i\)-th row of a matrix. Additionally, \(n^{\prime}\) refers to the number of nodes in \(\mathcal{G}\) available for training, while \(\mathcal{D}\) is a discriminator function such as cross-entropy. We will sometimes refer to the training loss (1) as an _upper-level_ energy function to differentiate its role from the lower-level energy defined next.

Moving to Unfolded GNNs.An _unfolded_ GNN architecture ensues when the functional form of \(h\) is explicitly chosen to align with the update rules that minimize a second, _lower-level_ energy function denoted as \(\ell(Y)\). In this way, it follows that \(\ell\left(Y^{(k)}\right)=\ell\left(h\left[Y^{(k-1)};W,A,X\right]\right)\leq \ell\left(Y^{(k-1)}\right)\), This restriction on \(h\) is imposed to introduce desirable inductive biases on the resulting GNN layers that stem from properties of the energy \(\ell\) and its corresponding minimizers (see Section 2.2 ).

While there are many possible domain-specific choices for \(\ell(Y)\) in the growing literature on unfolded GNNs (Chen et al., 2022; Ahn et al., 2022; Yang et al., 2021; Ma et al., 2021; Gasteiger et al., 2019; Pan et al., 2020; Zhang et al., 2020; Zhu et al., 2021; Xue et al., 2023)., we will focus our attention on a particular form that encompasses many existing works as special cases, and can be easily generalized to cover many others. Originally inspired by Zhou et al. (2003) and generalized by Yang et al. (2021) to account for node-wise nonlinearities, we consider

\[\ell(Y):=\|Y-f(X;W)\|_{F}^{2}+\lambda\operatorname{tr}(Y^{\top}LY)+\sum_{i=1}^ {n}\zeta(Y_{i}), \tag{2}\]

where \(\lambda>0\) is a trade-off parameter, \(f\) represents a trainable base model parameterized by \(W\) (e.g., a linear layer or MLP), and the function \(\zeta\) denotes a (possibly non-smooth) penalty on individual node embeddings. The first term in (2) favors embeddings that resemble the input features as processed by the base model, while the second encourages smoothness across graph edges. And finally, the last term is included to enforce additional node-wise constraints (e.g., non-negativity).

We now examine the form of \(h\) that can be induced when we optimize (2). Although \(\zeta\) may be non-smooth and incompatible with vanilla gradient descent, we can nonetheless apply proximal gradient descent in such circumstances (Combettes and Pesquet, 2011), leading to the descent update

\[Y^{(k)}=h\left(Y^{(k-1)};W,A,X\right)=\operatorname{prox}_{\zeta}\left(Y^{(k- 1)}-\alpha\left[(I+\lambda L)Y^{(k-1)}-f(X;W)\right]\right), \tag{3}\]

where \(\alpha\) controls the learning rate and \(\operatorname{prox}_{\zeta}(u):=\operatorname{arg\,min}_{y}\frac{1}{2}\|u-y \|_{2}^{2}+\zeta(y)\) denotes the proximal operator of \(\zeta\). Analogous to more traditional GNN architectures, (3) contains a \(\zeta\)-dependent nonlinear activation applied to the output of an affine, graph-dependent filter. With respect to the former, if \(\zeta\) is chosen as an indicator function that assigns an infinite penalty to any embedding less than zero, then \(\operatorname{prox}_{\zeta}\) reduces to standard ReLU activations; we will adopt this choice for MuseGNN.

### Why Unfolded GNNs?

There are a variety of reasons why it can be advantageous to construct \(h\) using unfolded optimization steps as in (3) or related. Of particular note here, unfolded GNN node embeddings inherit interpretable characteristics of the lower-level energy, especially if \(K\) is sufficiently large such that \(Y^{(K)}\) approximates a minimum of \(\ell(Y)\). For example, it is well-known that an oversmoothing effect can sometimes cause GNN layers to produce node embeddings that converge towards a non-informative constant (Oono and Suzuki, 2020; Li et al., 2018). However, it is transparent how to design \(\ell(Y)\) such that minimizers do not degenerate in this way, and hence the induced \(h\) as in (3) can be iterated indefinitely without risk of oversmoothing, even while spreading information across the graph.

As another representative scenario, in many customer-facing applications, it is desirable to provide explanations for why a particular prediction was made by a GNN model (Ying et al., 2019). In these situations, the embeddings of unfolded GNNs contain additional information when contextualized w.r.t. their role in
the lower-level energy. For example, if for some node \(Y_{i}^{(K)}\) is very far from \(f(X_{i};W)\) relative to other nodes, it increases our confidence that subsequent predictions may be based on network effects rather than the local feature \(X_{i}\). More broadly, the interpretability of unfolded GNNs facilitates bespoke modifications of (3) for addressing heterophily and/or adversarial attacks (Yang et al., 2021), handling long-range dependencies (Xue et al., 2023), forming connections with the gradient dynamics of physical systems (Di Giovanni et al., 2023) and deep equilibrium models (Gu et al., 2020; Yang et al., 2022), and exploiting the robustness of boosting algorithms (Sun et al., 2019) among other things.

### Scalability Challenges and Candidate Solutions

As benchmarks continue to expand in size (see Table 1 below), it is no longer feasible to conduct full-graph GNN training using only a single GPU or even single machine, especially so for relatively deep unfolded GNNs. To address such GNN scalability challenges, there are presently two dominant lines of algorithmic workarounds. The first adopts various sampling techniques to extract much smaller computational subgraphs upon which GNN models can be trained in mini-batches. Examples include neighbor sampling (Hamilton et al., 2017; Ying et al., 2018), layer-wise sampling (Chen et al., 2018; Zou et al., 2019), and graph-wise sampling (Chiang et al., 2019; Zeng et al., 2021). For each of these, there exist both online and offline versions, where the former involves randomly sampling new subgraphs during each training epoch, while the latter (Zeng et al., 2021; Gasteiger et al., 2022) is predicated on a fixed set of subgraphs for all epochs; MuseGNN will ultimately be based on a novel integration and analysis of this approach as will be introduced in Section 3.

The second line of work exploits the reuse of historical embeddings, meaning the embeddings of nodes computed and saved during the previous training epoch. In doing so, much of the recursive forward and backward computations required for GNN training, as well as expensive memory access to high-dimensional node features, can be reduced (Chen et al., 2017; Fey et al., 2021; Huang et al., 2023). This technique has recently been applied to training unfolded GNN models (Xue et al., 2023), although available performance results do not cover large-scale graphs, there are no convergence guarantees, and no code is available (at least as of the time of this submission) for enabling further comparisons. Beyond this, we are unaware of prior work devoted to the scaling and coincident analysis of unfolded GNNs.

## 3 Graph-Regularized Energy Functions Infused with Sampling

Our goal of this section is to introduce a convenient family of energy functions formed by applying graph-based regularization to a set of subgraphs that have been sampled from a given graph of interest. For this purpose, we first present the offline sampling strategy which will undergird our approach, followed by details of energy functional form we construct on top of it for use with MuseGNN. Later, we conclude by analyzing special cases of these sampling-based energies, further elucidating relevant properties and connections with full-graph training.

### Offline Sampling Foundation

In the present context, offline sampling refers to the case where we sample a fixed set of subgraphs from \(\mathcal{G}\) once and store them, which can ultimately be viewed as a form of preprocessing step. More formally, we assume an operator \(\Omega:\mathcal{G}\rightarrow\{\mathcal{G}_{s}\}_{s=1}^{m}\), where \(\mathcal{G}_{s}=\{\mathcal{V}_{s},\mathcal{E}_{s}\}\) is a subgraph of \(\mathcal{G}\) containing \(n_{s}=|\mathcal{V}_{s}|\) nodes, of which we assume \(n_{s}^{\prime}\) are target training nodes (indexed from \(1\) to \(n_{s}^{\prime}\)), and \(\mathcal{E}_{s}\) represents the edge set. The corresponding feature and label sets associated with the \(s\)-th subgraph are denoted \(X_{s}\in\mathbb{R}^{n_{s}\times d}\) and \(T_{s}\in\mathbb{R}^{n_{s}^{\prime}\times d^{\prime}}\), respectively.

There are several key reasons we employ offline sampling as the foundation for our scalable unfolded GNN architecture development. Firstly, conditioned on the availability of such pre-sampled/fixed subgraphs, there is no additional randomness when we use them to replace energy functions dependent on \(\mathcal{G}\) and \(X\) (e.g., as in (2)) with surrogates dependent on \(\{\mathcal{G}_{s},X_{s}\}_{s=1}^{m}\). Hence we retain a deterministic energy substructure contributing to a more transparent bilevel (upper- and lower-level) optimization process, as well as more interpretable node-wise embeddings that serve as attendant minimizers. Secondly, offline sampling allows us to conduct formal convergence analysis that is agnostic to the particular sampling operator \(\Omega\). In this
way, we need not compromise flexibility in choosing a practically-relevant \(\Omega\) in order to maintain desirable convergence guarantees. For example, as will be detailed later, ShadowKHop (Zeng et al., 2021) sampling pairs well with our approach and directly adheres to our convergence analysis in Section 5. And lastly, offline sampling facilitates an attractive balance between model accuracy and efficiency within the confines of unfolded GNN architectures. As will be shown in Section 6, we can match the accuracy of full-graph training with an epoch time similar to online sampling methods, e.g., neighbor sampling.

### Energy Function Formulation

To integrate offline sampling into a suitable graph-regularized energy function, we first introduce two sets of auxiliary variables that will serve as more flexible input arguments. Firstly, to accommodate multiple different embeddings for the same node appearing in multiple subgraphs, we define \(Y_{s}\in\mathbb{R}^{n_{s}\times d}\) for each subgraph index \(s=1,\ldots,m\), as well as \(\mathbb{Y}=\{Y_{s}\}_{s=1}^{m}\in\mathbb{R}^{(\sum n_{s})\times d}\) to describe the concatenated set. Secondly, we require additional latent variables that, as we will later see, facilitate a form of controllable linkage between the multiple embeddings that may exist for a given node (i.e., when a given node appears in multiple subgraphs). For this purpose, we define the latent variables as \(M\in\mathbb{R}^{n\times d}\), where each row can be viewed as a shared summary embedding associated with each node in the original/full graph.

We then define our sampling-based extension of (2) for MuseGNN as

\[\ell_{\text{muse}}\left(\mathbb{Y},M\right):=\sum_{s=1}^{m}\left[\|Y_{s}-f(X _{s};W)\|_{F}^{2}+\lambda\operatorname{tr}(Y_{s}^{\top}L_{s}Y_{s})+\gamma\|Y _{s}-\mu_{s}\|_{F}^{2}+\sum_{i=1}^{n_{s}}\zeta(Y_{s,i})\right], \tag{4}\]

where \(L_{s}\) is the graph Laplacian associated with \(\mathcal{G}_{s}\), \(\gamma\geq 0\) controls the weight of the additional penalty factor, and each \(\mu_{s}\in\mathbb{R}^{n_{s}\times d}\) is derived from \(M\) as follows. Let \(I(s,i)\) denote a function that maps the index of the \(i\)-th node in subgraph \(s\) to the corresponding node index in the full graph. For each subgraph \(s\), we then define \(\mu_{s}\) such that its \(i\)-th row satisfies \(\mu_{s,i}=M_{I(s,i)}\); per this construction, \(\mu_{s,i}=\mu_{s^{\prime},j}\) if \(I(s,i)=I(s^{\prime},j)\). Consequently, \(\{\mu_{s}\}_{s=1}^{m}\) and \(M\) represent the same overall set of latent embeddings, whereby the former is composed of repeated samples from the latter aligned with each subgraph.

Overall, there are three prominent factors which differentiate (4) from (2):

1. \(\ell_{\text{muse}}\left(\mathbb{Y},M\right)\) involves a deterministic summation over a fixed set of sampled subgraphs, where each \(Y_{s}\) is unique while \(W\) (and elements of \(M\)) are shared across subgraphs.
2. Unlike (2), the revised energy involves both an expanded set of node-wise embeddings \(\mathbb{Y}\) as well as auxiliary summary embeddings \(M\). When later forming GNN layers designed to minimize \(\ell_{\text{muse}}\left(\mathbb{Y},M\right)\), we must efficiently optimize over _all_ of these quantities, which alters the form of the final architecture.
3. The additional \(\gamma\|Y_{s}-\mu_{s}\|_{F}^{2}\) penalty factor acts to enforce dependencies between the embeddings of a given node spread across different subgraphs.

With respect to the latter, it is elucidating to consider minimization of \(\ell_{\text{muse}}\left(\mathbb{Y},M\right)\) over \(\mathbb{Y}\) with \(M\) set to some fixed \(M^{\prime}\). In this case, up to an irrelevant global scaling factor and additive constant, the energy can be equivalently reexpressed as

\[\ell_{\text{muse}}\left(\mathbb{Y},M=M^{\prime}\right)\equiv\sum_{s=1}^{m} \left[\|Y_{s}-\left[f^{\prime}(X_{s};W)+\gamma^{\prime}\mu_{s}\right]\|_{F}^{ 2}+\lambda^{\prime}\operatorname{tr}(Y_{s}^{\top}L_{s}Y_{s})+\sum_{i=1}^{n_{s} }\zeta^{\prime}(Y_{s,i})\right], \tag{5}\]

where \(f^{\prime}(X_{s};W):=\frac{1}{1+\gamma}f(X_{s};W)\), \(\gamma^{\prime}:=\frac{\gamma}{1+\gamma}\), \(\lambda^{\prime}:=\frac{\lambda}{1+\gamma}\), and \(\zeta^{\prime}(Y_{s,i}):=\frac{1}{1+\gamma}\zeta(Y_{s,i})\). From this expression, we observe that, beyond inconsequential rescalings (which can be trivially neutralized by simply rescaling the original choices for \(\{f,\gamma,\lambda,\zeta\}\)), the role of \(\mu_{s}\) is to refine the initial base predictor \(f(X_{s};W)\) with a corrective factor reflecting embedding summaries from other subgraphs sharing the same nodes. Conversely, when we minimize \(\ell_{\text{muse}}\left(\mathbb{Y},M\right)\) over \(M\) with \(\mathbb{Y}\) fixed, we find that the optimal \(\mu_{s}\) for every \(s\) is equal to the _mean_ embedding for each constituent node across all subgraphs. Hence \(M^{\prime}\) chosen in this way via alternating minimization has a natural interpretation as grounding each \(Y_{s}\) to a shared average representation reflecting the full graph structure.
This reformulation provides additional entry points for interpretability as well. For example, if optimal solutions are such that \(Y_{s}\) and \(\mu_{s}\) are similar while \(f(X_{s};W)\) is significantly different, we have a likely indication that, at least within this subgraph, graph structure and network effects are more dominant than node features. In contrast, if \(Y_{s}\) and \(f(X_{s};W)\) are similar with \(\mu_{s}\) much less so, feature utility may be selectively strong for a given subgraph; however, if all three quantities are similar the situation is more indicative of strong overall node features. Beyond these considerations, we now consider two limiting special cases of \(\ell_{\text{muse}}\left(\mathbb{Y},M\right)\) that provide complementary contextualization.

### Notable Limiting Cases

We first consider setting \(\gamma=0\), in which case the resulting energy completely decouples across each subgraph such that we can optimize each

\[\ell_{\text{muse}}^{s}(Y_{s}):=\|Y_{s}-f(X_{s};W)\|_{F}^{2}+\lambda\operatorname {tr}(Y_{s}^{\top}L_{s}Y_{s})+\sum_{i=1}^{n_{s}}\zeta(Y_{s,i})\ \ \forall s \tag{6}\]

independently. Under such conditions, the only cross-subgraph dependency stems from the shared base model weights \(W\) which are jointly trained. Hence \(\ell_{\text{muse}}^{s}(Y_{s})\) is analogous to a full graph energy as in (2) with \(\mathcal{G}\) replaced by \(\mathcal{G}_{s}\). In Section 5 we will examine convergence conditions for the full bilevel optimization process over all \(Y_{s}\) and \(W\) that follows from the \(\gamma=0\) assumption. We have also found that this simplified setting performs well in practice; see Appendix B.1 for ablations.

At the opposite extreme when \(\gamma=\infty\), we are effectively enforcing the constraint \(Y_{s}=\mu_{s}\) for all \(s\). As such, we can directly optimize all \(Y_{s}\) out of the model leading to the reduced \(M\)-dependent energy

\[\ell_{\text{muse}}(M):=\sum_{s=1}^{m}\left[\|\mu_{s}-f(X_{s};W)\|_{F}^{2}+ \lambda\operatorname{tr}(\mu_{s}^{\top}L_{s}\mu_{s})+\sum_{i=1}^{n_{s}}\zeta( \mu_{s,i})\right]. \tag{7}\]

Per the definition of \(\{\mu_{s}\}_{s=1}^{m}\) and the correspondence with unique node-wise elements of \(M\), this scenario has a much closer resemblance to full graph training with the original \(\mathcal{G}\). In this regard, as long as a node from the original graph appears in at least one subgraph, then it will have a single, unique embedding in (7) aligned with a row of \(M\).

Moreover, we can strengthen the correspondence with full-graph training via the suitable selection of the offline sampling operator \(\Omega\). In fact, there exists a simple uniform sampling procedure such that, at least in expectation, the energy from (7) is equivalent to the original full-graph version from (2), with the role of \(M\) equating to \(Y\). More concretely, we present the following (all proofs are deferred to Appendix D):

**Proposition 3.1**.: _Suppose we have \(m\) subgraphs \((\mathcal{V}_{1},\mathcal{E}_{1}),\ldots,(\mathcal{V}_{m},\mathcal{E}_{m})\) constructed independently such that \(\forall s=1,\ldots,m,\forall u,v\in\mathcal{V},\operatorname{Pr}[v\in \mathcal{V}_{s}]=\operatorname{Pr}[v\in\mathcal{V}_{s}\ |\ u\in\mathcal{V}_{s}]=p;(i,j)\in\mathcal{E}_{s}\iff i\in \mathcal{V}_{s},j\in\mathcal{V}_{s},(i,j)\in\mathcal{E}\). Then when \(\gamma=\infty\), we have \(\mathbb{E}[\ell_{\text{muse}}(M)]=mp\ \ell(M)\) with the \(\lambda\) in \(\ell(M)\) rescaled to \(p\lambda\)._

Strengthened by this result, we can more directly see that (7) provides an intuitive bridge between full-graph models based on (2) and subsequent subgraph models we intend to build via (4). Even so, we have found that relatively small \(\gamma\) values nonetheless work well in practice (although in principle it is actually possible to handle \(\gamma=\infty\) using techniques like ADMM (Boyd et al., 2011)).

## 4 From Sampling-based Energies to the MuseGNN Framework

Having defined and motivated a family of sampling-based energy functions vis-a-vis (4), we now proceed to derive minimization steps that will serve as GNN model layers that define MuseGNN forward and backward passes as summarized in Algorithm 1. Given that there are two input arguments, namely \(\mathbb{Y}\) and \(M\), it is natural to adopt an alternating minimization strategy whereby we fix one and optimize over the other, and vice versa.

With this in mind, we first consider optimization over \(\mathbb{Y}\) with \(M\) fixed. Per the discussion in Section 3.2, when conditioned on a fixed \(M\), \(\ell_{\text{muse}}\left(\mathbb{Y},M\right)\) decouples over subgraphs. Consequently, optimization can proceed using subgraph-independent proximal gradient descent, leading to the update rule

\[Y_{s}^{(k)}=\operatorname{prox}_{\zeta}\left[Y_{s}^{(k-1)}-\alpha\left([(1+ \gamma)I+\lambda L_{s}]Y_{s}^{(k-1)}-[f(X_{s};W)+\gamma\mu_{s}]\right)\right], \ \ \forall s. \tag{8}\]
Here the input argument to the proximal operator is given by a gradient step along (4) w.r.t. \(Y_{s}\). We also remark that execution of (8) over \(K\) iterations represents the primary component of a single forward training pass of our proposed MuseGNN framework as depicted on lines 6-8 of Algorithm 1.

We next turn to optimization over \(M\) which, as mentioned previously, can be minimized by the mean of the subgraph-specific embeddings for each node. However, directly computing these means is problematic for computational reasons, as for a given node \(v\) this would require the infeasible collection of embeddings from all subgraphs containing \(v\). Instead, we adopt an online mean estimator with forgetting factor \(\rho\). For each node \(v\) in the full graph, we maintain a mean embedding \(M_{v}\) and a counter \(c_{v}\). When this node appears in the \(s\)-th subgraph as node \(i\) (where \(i\) is the index within the subgraph), we update the mean embedding and counter via

\[M_{I(s,i)}\leftarrow\frac{\rho c_{I(s,i)}}{c_{I(s,i)}+1}M_{I(s,i)}+\frac{(1- \rho)c_{I(s,i)}+1}{c_{I(s,i)}+1}Y_{s,i}^{(K)},\qquad c_{I(s,i)}\gets c_{I(s,i)}+1. \tag{9}\]

Also shown on line 9 of Algorithm 1, we update \(M\) and \(c\) once per forward pass, which serves to refine the effective energy function observed by the core node embedding updates from (8).

For the backward pass, we compute gradients of

\[\mathcal{L}_{\text{muse}}(W,\theta):=\sum_{s=1}^{m}\sum_{i=1}^{n_{s}^{\prime}} \mathcal{D}\left(g\left[Y_{s}^{(K)}\left(W\right)_{i};\theta\right],T_{s,i}\right) \tag{10}\]

w.r.t. \(W\) and \(\theta\) as listed on line 10 of Algorithm 1, where \(\mathcal{L}_{\text{muse}}(W,\theta)\) is a sampling-based modification of (1). For \(W\) though, we only pass gradients through the calculation of \(Y_{s}^{(K)}\), not the full online \(M\) update; however, provided \(\rho\) is chosen to be sufficiently large, \(M\) will change slowly relative to \(Y_{s}^{(K)}\) such that this added complexity is not necessary for obtaining reasonable convergence.

```
1:Input: \(\{\mathcal{G}_{s}\}_{s=1}^{m}\): subgraphs, \(\{X_{s}\}_{s=1}^{m}\): features, \(K\): # unfolded layers, \(E\): # epochs
2:Randomly initialize \(W\) and \(\theta\); initialize \(c\in\mathbb{R}^{n}\) and \(M\in\mathbb{R}^{n\times d}\) to be zero
3:for\(e=1,2,\ldots,E\)do
4:for\(s=1,2,\ldots,m\)do
5:\(\mu_{s,i}\gets M_{I(s,i)},i=1,2,\ldots,n_{s}\)
6:\(Y_{s}^{(0)}\gets f(X_{s};W)\)
7:for\(k=1,2,\ldots,K\)do
8: Update \(Y_{s}^{(k)}\) from \(Y_{s}^{(k-1)}\) by (8) using \(\mu_{s}\)
9:endfor
10: Update \(M\) and \(c\) using (9) with \(Y_{s,i}^{(K)},i=1,2,\ldots,n_{s}\)
11: Update \(W\) and \(\theta\) via SGD over \(\mathcal{L}_{\text{muse}}\) from (10)
12:endfor
13:endfor
```

**Algorithm 1** Training procedure for MuseGNN

## 5 Convergence Analysis of MuseGNN

Global Convergence with \(\gamma=0\).In the more restrictive setting where \(\gamma=0\), we derive conditions whereby the entire bilevel optimization pipeline converges to a solution that jointly minimizes both lower- and upper-level energy functions in a precise sense to be described shortly. We remark here that establishing convergence is generally harder for bilevel optimization problems relative to more mainstream, single-level alternatives (Colson et al., 2005). To describe our main result, we first require the following definition:

**Definition 5.1**.: Assume \(f(X;W)=XW\), \(\gamma=0\), \(\zeta(y)\)=0, \(g(y;\theta)=y\), and that \(\mathcal{D}\) is a Lipschitz continuous convex function. Given the above, we then define \(\mathcal{L}_{\text{muse}}^{(k)}(W)\) as (10) with \(K\) set to \(k\). Analogously, we also define \(\mathcal{L}_{\text{muse}}^{*}(W)\) as (10) with \(Y_{s}^{(K)}\) replaced by \(Y_{s}^{*}:=\operatorname*{arg\,min}_{Y_{s}}\ell_{\text{muse}}^{s}(Y_{s})\) for all \(s\).
**Theorem 5.2**.: _Let \(W^{*}\) be the optimal value of the loss \(\mathcal{L}^{*}_{\text{muse}}(W)\) per Definition 5.1, while \(W^{(t)}\) denotes the value of \(W\) after \(t\) steps of stochastic gradient descent over \(\mathcal{L}^{(k)}_{\text{muse}}(W)\) with diminishing step sizes \(\eta_{t}=O(\frac{1}{\sqrt{t}})\). Then provided we choose \(\alpha\in\left(0,\min_{s}\left\|I+\lambda L_{s}\right\|_{2}^{-1}\right]\) and \(Y^{(0)}_{s}=f(X_{s};W)\), for some constant \(C\) we have that_

\[\mathbb{E}\left[\mathcal{L}^{(k)}_{\text{muse}}(W^{(t)})\right]-\mathcal{L}^{ *}_{\text{muse}}(W^{*})\leq O\left(\frac{1}{\sqrt{t}}+e^{-Ck}\right).\]

Given that \(\mathcal{L}^{*}_{\text{muse}}(W^{*})\) is the global minimum of the combined bilevel system, this result guarantees that we can converge arbitrarily close to it with adequate upper- and lower-level iterations, i.e., \(t\) and \(k\) respectively. Note also that we have made no assumption on the sampling method. In fact, as long as offline sampling is used, convergence is guaranteed, although the particular offline sampling approach can potentially impact the convergence rate; see Appendix C.2 for further details.

Lower-Level Convergence with arbitrary \(\boldsymbol{\gamma}\).We now address the more general scenario where \(\gamma\geq 0\) is arbitrary. However, because of the added challenge involved in accounting for alternating minimization over both \(\mathbb{Y}\) and \(M\), it is only possible to establish conditions whereby the lower-level energy (4) is guaranteed to converge as follows.

**Theorem 5.3**.: _Assume \(\zeta(y)=0\). Suppose that we have a series of \(\mathbb{Y}^{(k)}\) and \(M^{(k)}\), \(k=0,1,2,\ldots\) constructed following the updating rules \(\mathbb{Y}^{(k)}:=\arg\min_{Y}\ell_{\text{muse}}(\mathbb{Y},M^{(k-1)})\) and \(M^{(k)}:=\arg\min_{M}\ell_{\text{muse}}(\mathbb{Y}^{(k)},M)\), with \(\mathbb{Y}^{(0)}\) and \(M^{(0)}\) initialized arbitrarily. Then_

\[\lim_{k\to\infty}\ell_{\text{muse}}(\mathbb{Y}^{(k)},M^{(k)})=\inf_{\mathbb{Y },M}\ell_{\text{muse}}(\mathbb{Y},M). \tag{11}\]

While technically this result does not account for minimization over the upper-level MuseGNN loss, we still empirically find that the entire process outlined by Algorithm 1, including the online mean update, is nonetheless able to converge in practice. See Appendix C.1 for an illustration.

## 6 Experiments

We now seek to empirically show that MuseGNN serves as a reliable unfolded GNN model that:

1. Preserves competitive accuracy across datasets of widely varying size, including the very largest publicly-available graph benchmarks, with a single fixed architecture, and
2. Operates with comparable computational complexity relative to common alternatives that are also capable of scaling to the largest graphs.

For context though, we note that graph benchmark leaderboards often include top-performing entries based on complex compositions of existing models and training tricks, at times dependent on additional features or external data not included in the original designated dataset. Although these approaches have merit, our goal herein is not to compete with them, as they typically vary from dataset to dataset. Additionally, they are more frequently applied to smaller graphs using architectures that have yet to be consistently validated across multiple, truly large-scale benchmarks.

Datasets.We evaluate the performance of MuseGNN on node classification tasks from the Open Graph Benchmark (OGB) [19, 20] and the Illinois Graph Benchmark (IGB) [17]. Table 1 presents the details of these datasets, which are based on homogeneous graphs spanning a wide range of sizes. Of particular note is IGB-full, currently the largest publicly-available graph benchmark, along with a much smaller version called IGB-tiny that we include for comparison purposes. We also point out that while MAG240M is originally a heterogenous graph, we follow the common practice of homogenizing it [19]; similarly for IGB datasets, we adopt the provided homogeneous versions.
MuseGNN Design.For _all_ experiments, we choose the following fixed MuseGNN settings: Both \(f(X;W)\) and \(g(Y;\theta)\) are 3-layer MLPs, the number of unfolded layers \(K\) is 8 (rationaleic discussed later below), and the embedding dimension \(d\) is 512, and the forgetting factor \(\rho\) for the online mean estimation is 0.9. For offline subgraph sampling, we choose a variation on neighbor sampling called ShadowKHop (Zeng et al., 2021), which loosely approximates the conditions of Proposition 3.1. See Appendix A for further details regarding the MuseGNN implementation and hyperparameters.

Baseline Models.With the stated objectives of this section in mind, we compare MuseGNN with GCN (Kipf and Welling, 2017), GAT (Velickovic et al., 2018), and GraphSAGE (Hamilton et al., 2017), in each case testing with both neighbor sampling (NS) (Hamilton et al., 2017) and GNNAutoScale (GAS) (Fey et al., 2021) for scaling to large graphs. As a point of reference, we note that GAT with neighbor sampling is currently the top-performing homogeneous graph model on both the MAG240M and IGB-full leaderboards. We also compare with an analogous full-graph (FG) unfolded GNN (UGNN) with the same architecture as MuseGNN but without scalable sampling.

Accuracy Comparisons.As shown in Table 2, MuseGNN achieves similar accuracy to a comparable full-graph unfolded GNN model on the small datasets (satisfying our first objective from above), while the latter is incapable of scaling to even the mid-sized ogbn-products benchmark. Meanwhile, MuseGNN is generally better than the other GNN baselines, particularly so on the largest dataset, IGB-full. Notably, MuseGNN exceeds the performance of GAT with neighbor sampling, the current SOTA for homogeneous graph models on both MAG240M and IGB-full. Please also refer to Appendix B.1 for MuseGNN accuracy ablations involving \(\gamma\), where we show that while even \(\gamma=0\) achieves strong results, increasing \(\gamma>0\) can lead to further improvement. Likewise, Appendix B.2 contains experiments showing that the improvement of MuseGNN is not merely a product of using ShadowKHop sampling instead of more traditional neighbor sampling with baseline models; there we show that when switched to offline ShadowKHop sampling, the baseline GNNs degrade further, possibly because they are not anchored to an integrated energy.

Timing Comparisons.Turning to model complexity, Table 3 displays the training speed of MuseGNN relative to two common baselines. From these results, we observe that MuseGNN executes with a similar epoch time (satisfying our second objective from above).

Convergence Illustration.In Figure 1 we show the empirical convergence of (4) w.r.t. \(\mathbb{Y}\) during the forward pass of MuseGNN models on ogbn-papers100M for differing values of \(\gamma\). In all cases the energy converges within 8 iterations, supporting our choice of \(K=8\) for experiments.

## 7 Conclusion

In this work, we have proposed MuseGNN, an unfolded GNN model that scales to large datasets by incorporating offline graph sampling into the design of its lower-level, architecture-inducing energy function. In so doing, MuseGNN readily handles graphs with \(\sim 10^{8}\) or more nodes and high-dimensional node features, exceeding 1TB in total size, all while maintaining interpretable layers with concomitant convergence guarantees.

\begin{table}
\begin{tabular}{l r r r r r} \hline \hline
**Dataset** & \(|\mathcal{V}|\) & \(|\mathcal{E}|\) & **Dim.** & **\# Class** & **Dataset Size** \\ \hline ogbn-arxiv (Hu et al., 2020) & 0.17M & 1.2M & 128 & 40 & 182MB \\ IGB-tiny (Khatua et al., 2023) & 0.1M & 0.5M & 1024 & 19 & 400MB \\ ogbn-products (Hu et al., 2020) & 2.4M & 123M & 100 & 47 & 1.4GB \\ ogbn-papers100M (Hu et al., 2020) & 111M & 1.6B & 128 & 172 & 57GB \\ MAG240M (Hu et al., 2021) & 244.2M & 1.7B & 768 & 153 & 377GB\({}^{1}\) \\ IGB-full (Khatua et al., 2023) & 269.3M & 4.0B & 1024 & 19 & 1.15TB \\ \hline \hline \end{tabular}
\end{table}
Table 1: Dataset details, including node feature dimension (Dim.) and number of classes (# Class).