MuseGNN: Forming Scalable, Convergent GNN Layers that Minimize a Sampling-Based Energy
Authors: Haitian Jiang
Renjie Liu
Zengfeng Huang
Yichuan Wang
Xiao Yan
Zhenkun Cai
Minjie Wang
David Wipf
Venue: ICLR 2025 Poster (2025)
Keywords: Graph Neural Networks
Energy-based Models
Scalable Training
Bi-level Optimization
Primary area: learning on graphs and other geometries & topologies

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 transparent, exploitable 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. Our source code is available at https://github.com/haitian-jiang/MuseGNN.

Published as a conference paper at ICLR 2025
MUSEGNN: FORMING SCALABLE, CONVERGENT GNN LAYERS THAT MINIMIZE A BASSED ENERGY
Haitian Jiang"
Renjie Liu?
Zengfeng Hunag 3
Yichuan Wang*
Xiao Yan
Zhenkun Cai
Minjie Wang
David Wipf
NNww York University 'Sutherr University of Science and Technology
Fudan University
*UC Berkeley "Centre for Perceptual and Interactive Intelligence AAmzon Web Services
"haitian. hhattian.jianggyyuu jiangenyu.edu
Tdavi.dwipf@emmal.com
ABSTRACT
Among the many variants of graph neural network (GNN) architectures capable of modeling data with cross-instance relations, 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 transparent, exploitable inductive biases and interpretability. However, scaling GNN architectures con- structed in this way remains challenging, in part because the convergence of the forward pass may involve models with considerable depth. To tackle this limita- tion, 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-avan availallee node classificattin benchmark exceeding ITB size, Our source code available at hhttps://githhhi com wwwaiiiiiiei. jiang/ jiang/MuseGng MuseGNN
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 indepen- dent samples, GNNs excel in capturing the complex cross-instance relations modeled by graphs (Hamilton et al. [0177; Kearnes et al.] 2016)] Kipf & Welling [017] Velickovic et al.) 2018). Foun- dational to GNNs is the notion message passing (Gilmer et al.] 2017). whereby nodes iteratively gather information from neighbors to update their representations. In doing information can propagate across the graph in the form of node embeddings, reflecting both local patterns and global network effects, which required by downstream tasks such node classification.
Among many GNN architectures, one promising subclass is based on graph propagation layers de- rived to be in one-to-one correspondence with the descent iterations of graph-regularized energy function (Ahn et al.
2022] Chen et al. 2022a 2021 Yang et al.) 2021 Ma et al. 2021 Gasteiger al. 2019 Pan al.
2020 Zhang et al.) 2020 Zhu al. 2021 Xue al 2023). For these models, approximations to minimizer
the layers of the GNN forward pass computes increasingly-ren
of the aforementioned energy function. Importantly, if such energy minimizers possess interpretable properties or inductive biases, the corresponding GNN architecture naturally inherits them (Zheng al. 2024). 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., robustly handling graphs with spurious edges and/or heterophily
Fu et al. 2023: Yang et al., 2021). Borrowing from Yang et al. (2022), we will refer models of
"Work partially completed during an internship at Amazon Web Services.
1 'Coresponding author.

---

Published as a conference paper at ICLR 2025
this genre as unfolded GNNs, UGNNs for short, given that the layers are derived from so-called unfolded (in time/iteration) energy descent process.
Despite their merits w.L.t interpretability, UGNNs face non-trivial scalability challenges. This is in part because they can be constructed with arbitrary depth (i.c., arbitrary descent iterations) while still avoiding undesirable oversmoothing effects (Oono & 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 UGNNs, or providing any complementary guarantees regarding convergence on large graphs. Hence most UGNN models are presently evaluated on relatively small benchmarks.
To address these shortcomings, we propose scalable UGNN model that incorporates efficient sub- graph 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 properties, interpretability, 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 overhead 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 architec- ture produced by the mininization of an unfolded sampling -basdd energy. Building on background motivations presented in Section 2
our MuseGNN-rellted contributions are three-fold:
1. In Sections B and 4 we expand widely-used UGNN framework to incorporate offline sampling into the centtetuur-inndueinn energy function design itself, as opposed to a post hoc application of sampling existing GNN methods. The resulting model, which we term MuseGNN, allows us to combine the attractive properties of UGNNs with the scalability of canonical GNNs. No- tably, by design MuseGNN can exploit an unbiased estimator of the original full-graph energy if desired, or else sampling-basee
alternatives that are provably more expressive.
2. We prove in Section that MuseGNN possesses desirable convergence properties regarding both upper-level (traditional model training) and lower-level (interpretable energy function descent) optimization processes. A supporting empirical example of the latter is illustrated in Figure above. Critically, our convergence gaurantees are agnostic to the particular offline sampling operator that is adopted, and any future options can be seamlessly integrated. These attributes increase our confidence in reliable performance when moving to new problem domains that may deviate from known benchmarks.
3. 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 ITB in size.
Table D contextualizes MuseGNN w.r.t. prior work, with details to follow in subsequent sections. Specifically, compare with those commonly-used GNN architectures that have been success- fully scaled and validated on the largest public graphs using some form of neighbor sampling (NS) (Hamilton et al., 2017; Huang et al. 2024, Waleffe et al. 2023), or alternatively, GNNAu- toScale (GAS) (Fey et al. [202]] Yu et al. 2022). Note that most GNNs populating competitive leaderboards such as OGB do not presently scale to the largest graphs, nor do they possess prop- erties of UGNNs, and hence are not our focus. And finally, we compare with existing full-graph (FG) UGNN approaches, as well as USP (Li et al. (2022) and LazyGNN (Xue et al. (023), two recent frameworks explicitly designed for scaling UGNNs. Appendix B. also considers additional scalable GNN models.
BACKGROUND AND MOTIVATION
2.1 GNN ARCUITECTRRE FROM UNFOLDED OPTIMIZATION
Notation.
Let = (V.E) denote a graph with n = | nodes and edge set E. We define D and A as the degree and adjacency matrices of 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' dimeensiona label vector, the respective collections of which are given by X E R and T E R=add"
2

---

Published as a conference paper at ICLR 2025
Table 1: MuseGNN vs. existing methods on largest graphs (LG), where "top acc.' refers to top LG accuracy. Note that the convergence guarantee and greater energy expressivity are specifically defined w.r.t. UGNN models, hence 'N/A' for non-UGNNs without lower-level energy.
y-
y=
Log-Energy
Y=2.0
==3.0
energy converge handle top 1 energy LGs acc. expressivity
107
descent guarantec
GNN+NS
X
NA
/
N/A
GNN+GAS
X
N/A
X
N/A
0
2
Layer
UGNN(FG)
V
V
X
X
USP
/
partial
X
X
Figure 1: MuseGNN forward pass convergence, energy descent on pabn-pppers 100M dataset with varying (to be introduced in Section 3.2)
LazyGNN
X
X
MuseGNN
V
GNN Basics. Given graph defined as above, canonical GNN architectures designed to pro- duce a sequence of node- wise embeddings (Y k== that are increasingly refined across K model layers according to the rule y(k) h(Y W, A, X) E Raxd, Here h represents function that updates (k-1) based trainable model weights W well as A (graph structure) and option- ally X (input node features). To facilitate downstream tasks such as node classification, W may be trained, along with additional parameters of any application-specffi output layer Rd + Rd minimize loss of the form
n2' C(W,0)- = Σ ED ED(o ( .
(1)
1=1
In this expression, y,()) (()) (W); reflects the explicit dependency of GNN embeddings on W and the subscript i denotes the i-th row of a matrix. Additionally, n' refers to the number of nodes in G available for training, while D discriminator function such as cross-entropy. We will sometimes refer to the training loss (T) 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 explicitly chosen to align with the update rules that minimize a second, lower-level energy function denoted as ((Y). In this way, follows that E ()()) l (h (k W, A, x]) S & (k-))))) This restriction on h is imposed to introduce desirable inductive biases on the resulting GNN layers that stem from properties of the energy l and its corresponding minimizers (see Section 2.2 ).
While there are many possible domain-specine folded GNNs (Ahn et al., 2022] Chen et al. 2022a Ma al. [022 Gasteiger et al.) 2019 Pan al.
choices for {(Y) in the growing literature on un-
2020 Yang al. 2021 Zhang et al.) 2020 Zhu et al. 202 1 Xue et al. (023).
will focus
attention on a particular lar form that encompasses many existing works as special cases, and can be easily generalized to cover many others. Originally inspired by Zhou al. (2003) and generalized
by] Yang et al. (222)) to account for node-wise nonlinearities, we consider
n
{Y))= f(X; W) X Ar LLY) Σ 5(0).
where À > 0 is a trade-off parameter, f represents a trainable base model parameterized by W (e.g., linear layer or MLP), and the function C denotes (possibly non-smooth) penalty individual node embeddings. The first term (2) favors embeddings that resemble the input features as pro- cessed 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 nog-ttgaaii
We now examine the form of h that can be induced when optimize (2). Although may be non- smooth and incompatible with vanilla gradient descent, we can nonetheless apply proximal gradient descent in such circumstances (Combettes & Pesquet 2011), leading to the descent step
(+) h (y (y(I-1), W, A, x) prox
(x-1-1)
[ (x - s(x;w)).
where a controls the learning rate and prox; (u) := arg min,, 111u- 1uu-9 -112 C(y) denotes the proximal c
operator of Ç. Analogous to more traditional GNN architectures, 3 contains (-dependent nonlin-
3

---

Published as a conference paper at ICLR 2025
activation applied to the output of affine, graph-dependent filter. With respect to the former, if Ç is chosen as an indicator function that assigns an infinite penalty to any embedding less than zero, then prox; reduces to standard ReLU activations; we will adopt this choice for MuseGNN.
2.2 WHY UNFOLDED GNNS?
There are a variety of reasons why it can be advantageous to construct h using unfolded opti- mization steps as in (3) or related. Of particular note here, UGNN node embeddings inherit ex- ploitable/interpretbble characteristics of the lower-level energy (Chen et al., 2021 Zheng et al.] 2024), especially if K sufficiently large such that Y(K) approximates a minimum of {(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 & Suzuki 2020 Li ct al.
2018). This can be understood through the lens of minimizing the the second term of
isolation Cai & Wang 2020). The latter can be driven to zero whenever Y((K) y,(K) for i.j E V.
since (r(()(()) [TY(K) = En.jEE (| - 1 However, is clear how to design {(Y)
such that minimizers do not degenerate in this way, e.g., by adding the first term in (2) or related generalizations, it has been previously established that oversmoothing is effectively miligated, even
while spreading information across the graph (Fu et al. 2023 Ma et al. 2021 Pan et al. 2020 Yang
et al. 2021 Zhang et al. 2020 Zhu et al. 2021).
UGNNs provide other transparent entry points for customization as well. For example, if an energy function is insensitive to spurious graph edges, then a corresponding GNN architecture constructed via energy function descent is likely be robust against corrupted graphs from adversarial attacks or
heterophily (Fu al.] [0233] Yang et al.] (02I]. More broadly, the flexibility of UGNNs facilitates
bespoke modifications of (3) for graph signal denoising (Chen et al. 2021). handling long-range
dependencies (Xue et al. 2023), forming connections with the gradient dynamics of physical sys-
tems (Di Giovanni et al. 2023 and deep equilibrium models (Gu et al. 20200 Yang et al.) (022),
exploiting the robustness of boosting algorithms (Sun al., 2019). or differentiating the relative
importance of features versus network effects in making predictions (Yoo et al. (023).
2.3 SCALABILITY CHALLENGES AND CANDIDATE SOLUTIONS
As benchmarks continue to expand in size (see Table 4p. it is no longer feasible to conduct full- graph GNN training using only single GPU or even single machine, especially so for relatively deep UGNNs. 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. Relevant exam- ples include neighbor sampling (Hamilton et al. 2017] Ying et al.) 2018) layer- sampling Chen cct al. 2018 Zou et al., 2019). and graph-wise sampling (Chiang et al. [019 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 fixed set of subgraphs for all epochs. Recent work specific to scaling
et al. (2022) is predicated
UGNNs (Li et al. (022) has incorporated subgraphs based on graph partitioning, i.c., splitting the
original full graph into disjoint subgraphs, and using these to approximate gradients involving the
full-graph energy. While convergence is established for the forward pass, referred to as an unbiased stochastic proximal-solve
(USP), there are no guarantees for the joint forward- forward-backward backward passes
during training, nor empirical comparisons on large graphs; see Appendix B.5 for further discussion of USP and how differs from MuseGNN.
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. 202T Huang et al.) 2024 This technique has recently been applied to training UGNN models via the LazyGNN frame- work (ue et al. (023). as well as somewhat related implicit GNN models (Chen et al. 2022b), although available performance results do not cover large-scale graphs and there are no convergence guarantees. Beyond this (and USP from above), we are unaware of prior work specifically devoted to the scaling and coincident analysis of unfolded GNNs.
4

---

Published as a conference paper at ICLR 2025
3
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-basee energies, further elucidating relevant properties and connections with full-graph training.
3.1 OFFLINE SAMPLING FOUNDATION
In the present context, offline sampling refers to the case where we sample fixed set of subgraphs from G once and store them, which can ultimately be viewed as a form of preprocessing step. More formally, we assume an operator f2 9 { where G. (V., E subgraph of g containing = | nodes, of which we assume n' are target training nodes (indexed from 1 to n²), and & represents the edge set. The corresponding feature and label sets associated with the s-th subgraph are denoted X, E I xd and T, E R xd' respectively.
There are several key reasons employ offline sampling as the foundation for scalable un- folded GNN architecture development. Firstly, conditioned on the availability of such pre-sampled subgraphs, there is no additional randomness when we them to replace energy functions depen- dent on G and X (e.g., as in ( with surrogates dependent on {g., X. ).)=- Hence we retain a deterministic energy substructure contributing more transparent bilevel (upper- and lower-level) optimization process and attendant node-wise embeddings that serve as minimizers. Secondly, of- fline sampling allows us conduct formal convergence analysis that agnostic to the particular sampling operator Ω. In this way. we need not compromise flexibility in choosing a practically- relevant in order to maintain desirable convergence guarantees. And lastly, offline sampling facil- itates an attractive balance between model accuracy and efficiency within the confines of unfolded GNN architectures. As will be shown Section 6 we can match the accuracy of full full-graph training with epoch time similar to online sampling methods, e.g., neighbor sampling.
3.2 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, E RT, xd for each subgraph index s = 1....,,. m, as well as Y == (Y,), E R )x)xd to describe the concatenated set. Secondly, we require additional latent variables that, we will later see, facilitate a form of controllable linkage between the multiple embeddings that may exist for given node (i.e., when a given node appears multiple subgraphs). For this purpose, we define the latent variables as M E xxd where each row can be viewed as shared summary embedding associated with each node in the original/full graph.
We then define our sampling-based extension of 2 for MuseGNN
Eusoe M) = Σ |r= f(X:: W))l + Atr(Y, L.Y.) + + - Polla + Σ ((Y..) (4)
1=1
where L. is the graph Laplacian associated with G.. 7 2 controls the weight of the additional penalty factor, and each μa xd is derived from M as follows. Let [(s, i) denote a function that maps the index of the i-th node in subgraph to the corresponding node index in the full graph. For each subgraph s, we then define μl, such that its i-th row satisfies M4,i construction, Me,i Ma'.j if is) I(s', j). Consequently, { 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.
Mr( per this
)m
Overall, there are three prominent factors which differentiate from
1. mmsse (Y,M) involves a deterministic summation over a fixed set of sampled subgraphs, where each Y, is unique while W (and elements of M) are shared across subgraphs. As will be dis- cussed further below, this energy actually has the potential to be more expressive relative to the full graph form in 2
5

---

Published as a conference paper at ICLR 2025
2. Unlike (2) the revised energy involves both an expanded set of node- embeddings Y as well as auxiliary summary embeddings M. When later forming GNN layers designed to minimize Lmsse (", Mf), we must efficiently optimize over all of these quantities, which alters the form of the final architecture.
3. The additional - 2 penalty factor acts to enforce dependencies between the embeddings of given node spread across different subgraphs.
With respect to the latter, it is elucidating to consider minimization of mmsse (Y, M) over Y with M set to some fixed M'. In this case, up to irrelevant global scaling factor and additive constant, the energy can be equivalently reexpressed as
y/
Cuse (1, M = M1) = Σ
|r, - +) + tTT( L.Y.) + Σ s'(Ya))
(5) 14. and ('YY,,) := + ((Y,,). From this expression, we observe that, beyond inconsequential rescalings (which can be trivially neutralized by simply rescaling the original choices for {f. 7. ()). the role of μa to refine the initial base predictor s(X,; W) with corrective factor reflecting embedding summaries from other subgraphs sharing the same nodes. Conversely, when we minimize (use (Y, .)) over M with Y fixed, we find that the optimal H4, for every s is equal to the mean embedding for each constituent node across all subgraphs. Hence M' chosen this way via alternating minimization has natural interpretation as grounding each Y, to a shared average representation reflecting the full graph structure. We now consider two limiting special cases muse (Y, M) that provide complementary contextualization.
where /
3.3 NOTABLE LIMITING CASES
We first consider setting 7 = 0, in which case the resulting energy completely decouples across each subgraph such that we can optimize each
The
e muxe ) = YY f(X,; W) + Atr(Y, L.Y.) + Σ ((Y,.), Vs
(6)
i=1
independently, 8 1,.... m. Under such conditions, the only cross-subgraph dependency stems from the shared base model weights W which are jointly trained. Hence ( is analogous to a full graph energy as in (2) with G replaced by G.. We remark that there exists non-isomorphic graphs such that the induced full-graph energies from (2) have equivalent minima, and yet the cor- responding ( (Y,) )== from (6) do not; see Appendix Corr additional analysis and examples. Hence in this sense the sampling-he based energy has the potential to be more expressive than the orig- inal, and need not be viewed as merely an approximation thereof. In Section (5) we will examine convergence conditions for the full bilevel optimization process over all Y, and W that follows the = O assumption. We have also found that this simplified setting performs well in practice; see Appendix B.2 for ablations.
At the opposite extreme when 7 = co, we are effectively enforcing the constraint Y, μa, Vs. As such, we can directly optimize all Y. out of the model leading to the reduced M M-dependent energy
Cmuse (M)
Il,, S(X:; W) Atr(u, Latts) +
((ua,))
Per the definition of { and the correspondence with unique node-wise elements of M, this scenario has a much closer resemblance to full graph training with the original G. In this regard, as long as node from the original graph appears in at least one subgraph, then it will have single, unique embedding in aligned with row of M.
Moreover, we can strengthen the correspondence with full-graph training via the suitable selection of the offline sampling operator Ω. fact, there exists 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
Proposition 3.1. Suppose we have m subgraphs Ei)....... (Yn E, constructed independently such that Vs 1,..., m. Vu, V, Pr(u V.) Pr(u E V, u E V.J pr (i.j) E. > i

---

Published as a conference paper at ICLR 2025
Va. E Va, (i.j) E. Then when y oc, have B[lmuse (M)) mp ((M) with the X in E(M) rescaled to pd.
Strengthened by this result, we full-graph models based on (2) and subsequent subgraph models we intend to build via 4 with the later serving unbiased estimator of the former. Even so, have found that relatively small values nonetheless work well in practice, possibly (at least in some situations) because the sampling- energy itself may be more expressive as mentioned earlier.
directly that 7 provides intuitive bridge between
FROM SAMPLING-BASED ENERGIES TO THE MUSEGNN FRAMEWORK
Having defined and motivated family of sampling-based energy functions vis-a-vis 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 [ Given that there two input argu- ments, namely 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 Y with M fixed. Per the discussion in Section 3.2 when conditioned on fixed M, muse (Y, M) decouples over subgraphs. Consequently, opti- mization can proceed using subgraph-inddeeeden proximal gradient descent, leading to the update rule
(( & (y(k-1) -a a -a(1))aaaa111 +7)1 ALg) Y [f(x: [f(X,:
+ +wml vs.
(8)
Here the input argument to the proximal operator is given by a gradient step along w.r.t. Y.. We also remark that execution of (8) over K iterations represents the primary component of single forward training pass of our proposed MuseGNN framework as depicted on lines 6-8 of Algorithm
We next turn optimization over M which, mentioned previously, can be minimized by the mean of the subgraph-speciff embeddings for each node. However, directly computing these means is problematic for computational reasons, for given node this would require the infeasible col- lection of embeddings from all subgraphs containing c. Instead, we adopt an online mean estimator with forgetting factor p. For each node in the full graph, maintain mean embedding M, and a counter Cy. When this node appears in the s-th subgraph as node i (where is the index within the subgraph), we update the mean embedding and counter via
pCr(s.)
(1 (1- ()C)))((.))) + (K)
Mr(s.))
MI(s.))
CI(s.)) + Ci(a.)) +
(9)
Ci(s,1) +1
CI(a.)) +1
we update M and e once per forward pass, which serves to refine the effective energy function observed by the core node embedding updates from (8)
Also shown on line of Algorithm [
For the backward pass, we compute gradients of
Liusse (W,0) = Σ Σ D g (()) (W),; .T..))
(10)
w.r.t. W and 0 as listed on line 10 of Algorithm where Luuse (W. 0) is sampling modifica- tion of (). For W though, we only pass gradients through the calculation of (K) not the full online M update; however, provided p chosen be sufficiently large, M will change slowly relative to y-(k) such that this added complexity is not necessary for obtaining reasonable convergence.
5 CONVERGENCE ANALYSIS OF MUSEGNN
Global Convergence with 7 In the more restrictive setting where ς 0, we derive condi- tions whereby the entire MuseGNN bilevel optimization pipeline converges to solution that jointly minimizes both lower- and upper-level energy functions 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:

---

Published as a conference paper at ICLR 2025
Definition 5.1. Assume ((X;W) = XW, ((y) =0, g(y: 0) y, and that D is Lipschitz continuous convex function. Given the above, we then define C (W) as 10 with K set to k. Analogously, also define (W") (10p with y,(k) replaced by Y arg miny, f (Y,) for all s.
Theorem 5.2. Let W be the optimal value of the loss misee (W) per Definition 5.1 while W de- notes the value of W aftert steps of stochastic gradient descent over Lmise (W) with diminishing step sizes Tit =0 -1)) Then provided we choose E (0,min, || AL 2 and Y! =fxx; W), for some constant we have that
I0 (1/ +o-ck)
E C(ke (W())
muce (W")
Given that Cimuse (W*) is the global minimum of the combined bilevel system, this result guaran- tees that we can converge arbitrarily close to with adequate upper- and lower-level iterations, i.e., and respectively. Note also that we have made assumption on the sampling method. In fact, as long as offline sampling is used, conver- gence is guaranteed, although the particular of- fline sampling approach can potentially impact the convergence rate; see Appendix D.2 for fur- ther details.
Algorithm MuseGNN Training Procedure
Require: {9 subgraphs, {X,, 1mm feat., K: # unfolded layers, E: # epochs
1: Randomly initialize W and 0; Initialize с I= and M E Raxd to zero
2: for ........ E do
do 4: ,a. f- M M(..).. 1.,...,
3: for s = = 1.2......
...,
(2 + f(xii W)
5:
6:
for k = 1--1, 1,2.... K do
Update Y:( using
7:
Lower-Level Convergence with arbitrary y We now address the more general scenario where ? 2 0s arbitrary. However, because of the added challenge involved in accounting for alternating minimization over both and M, is only pos- sible establish conditions whereby the lower- level energy is guaranteed to converge as fol- lows.
8: end for
9: Update M, e via 9
10: Update W, via SGD over loss
10
end for
12: end for
Theorem 5.3. Assume ((y) = O. Suppose that we have a series of (t) and M(k), k= 0, constructed following the updating rules (y(k) == arg miny Emuse (Y, (Y,M. (()))) and M(k) arg minM E ( (R) M). with y( and O()) initialized arbitrarily. Then
inf Emuse (Y, M). lim Emuse ( (x(()) ((0) Y.M 100
(11)
While technically this result does not account for minimization over the upper-level MuseGNN loss, still empirically find that the entire process outlined by Algorithml [ including the online mean update, is nonetheless able to converge in practice; see Appendix D.I example.
6 EXPERIMENTS
We now seek to show that MuseGNN serves as reliable unfolded GNN model that:
Preserves competitive accuracy across datasets of widely varying size, especially on the very largest publicly-availaae graph benchmarks, with single fixed architecture,
2. Operates with comparable computational complexity relative common alternatives that also capable of scaling to the largest graphs, and
3. Empirically converges across different values that modulate our proposed sampling- en- ergy. in accordance with theory.
For context though, note that graph benchmark leaderboards often include aopperfforiig tries based on complex compositions of existing models and training tricks, at times dependent on additional features external data not included in the original designated dataset. Although these
8

---

Published as a conference paper at ICLR 2025
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 architec- tures that have yet to be consistently validated across multiple, truly large-scale benchmarks.
Node classification accuracy (%) the test set, except for MAG24CM which only has labels for validation set. Bold numbers denote the highest accuracy. For the two largest datasets, MuseGNN is currently the top-performing homogeneous graph model on the relevant OGB-LSC and IGB leaderboards respectively, even while maintaining the attractive interpretability attributes of an unfolded GNN. We omit error bars for baseline results in the two largest datasets because of the high cost to run them all. Additionally, OOM refers to out- out-of-memory
Table 2:
Small
Medium
Large
Largest (V) > 220MM
(VI <0.5M)
(VI 22
(V) = 100M)
Model
arxiv
TGB-t
products papers100M
OMG240M
[GB-Ful1
GCN (NS)
69.71+00.25 69.77+0.11
78.49-0.553
65.83-00336
65.24
48.59
SAGE (NS)
70.49+00.20 72.42+0.09
78.29=0016
66.20-0.11
66.79
54.95
GAT (NS)
69.94110.2 69.70.10.09
79.4510.59
66.28110.08
67.15
55,51
SGFormer (NS)
70.55:10.24
73.3710.12
78.9410.25
66.01
65,29
N/A
MariusGNN
69.36110.09 73.06-10.06
76.95.10.24
62.97-0.113
63.17
54.99
FreshGNN
71.51 71.49++013
79.21=00377
66.22=0007
65.53
56.50
GCN (GAS)
71.6810.3
67.86110.20
76.6100.3
54.2.10.7
OOM
OOM
SAGE (GAS)
71.35:10.4
69.35
77.710.7
57.9.1..4
OOM
OOM
GAT (GAS)
00.89+0.1
69.23+0.17
76.9
OOM
OOM
OOM
UGNN (FG)
7..7
72.44410.09
OOM
OOM
OOM
OOM
LazyGNN!
72.30:10.18 72.
0..21
39.80=10009 N/A
OOM
OOM
USP
71.6+10.2
N/A
73.8
N/A
N/A
MuseGNN
0.19 73.42-00.3
81.23 0.339
6.082-10022
67.67 +
0.113 61.74-10.5
Datasets. We evaluate the performance of MuseGNN on node classification tasks from the Open Graph Benchmark (OGB) (Hu et al. 2020 2021) and the Illinois Graph Benchmark (IGB) (Khatua et al. 2023). which are based homogeneous graphs spanning a wide range of sizes. Table 4 Appendix A presents the relevant size-related details. Of particular note is IGB ful-full, currently the largest publicly-availabl
graph benchmark exceeding 1TB in size. We also point out that while MAG24 4 0M is originally a heterogenous graph, we follow the common practice of homogenizing (Hu et al. 2021); 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 yY; 8) are 3-layer MLPs, the number of unfolded layers K 8 (rationale discussed later below), and the embedding dimension d is 512, and the forgetting factor for the online mean estimation 0.9. For offline subgraph sampling, we choose variation on neighbor sampling called ShopowKHop (Zeng et al., 2021), which loosely approximates the conditions of Proposition 3.1 See Appendix for further details regarding the MuseGNN implementation and hyperparameters.
Baseline Models. With the stated objectives of this section in mind, we compare MuseGNN with very recent scalability-focused GNN frameworks MariusGNN (Waleffe et al. 2023). FreshGNN (Huang al. 2024), and SGFormer (Wu et al. 2024), the latter representing graph transformer equipped with neighbor sampling² and other adaptations explicitly for handling large graphs. Likewise we include GCN (Kipf & Wellingl 2017). GAT (Velickovic al., (2188), and GrephSAGE (Hamilton et al. 2017), in each case testing with both neighbor sampling (NS) (Hamil- ton al., 2017) and ScaletoS (GAS) (Fey al. 2021) for scaling the largest graphs. As a key point of reference, we note that GAT with neighbor sampling is currently the top cop-performing homogencous graph model on both the MAG24 4 OM and IGB- leaderboards. Another natural baseline we adopt is the analogous full-graph (FG) UGNN with the same architecture as MuseGNN
'The accuracy of LazyGNN on arxiv and products is slightly different from what is reported in Xue et al. (023) for reasons specified in Appendix|A]
In Wu et al. (2024), neighbor sampling is only used for larger graphs: however, for consistency across meth- ods and henchmarks in our experiments, apply sampling to SGFormer for all datasets regardless of size; likewise for MuseGNN and all other sampling-based models.

---

Published as a conference paper at ICLR 2025
but without scalable sampling. And finally, we compare with both recent approaches explicitly designed for scaling UGNNs, namely, USP (Li et al. 2022) and LazyGNN (Xue et al. (023) as dis- cussed in Sections and 2.3 For USP, the core solver applied to the energy from (2) for the most direct head-to-head comparisons, noting that MuseGNN, LazyGNN, and USP can all potentially be applied to alternative energy functions improve performance if needed (we the accuracy ported in the USP paper where available because there is no public code at this time). Further com- parisons against other scalable baselines Cluster-GCN (Chiang et al.) 2019), GraphFM (Yu al. 2022), SGC (Wu et al. 2019), and SIGN (Rossi et al. 2020) are deferred to Appendix B.1
Accuracy Comparisons. As shown in Table 2 MuseGNN achieves similar accuracy to a compara- ble 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 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 (NS) and FreshGNN, the cur- rent SOTA for homogeneous graph models on MAG24 0M and IGB IGB-full. Similarly, MuseGNN outperforms all of the additional scalable GNN frameworks presented Appendix [B.I In terms of prior work scaling UGNNs, public code for USP is presently not available; however, based on reported results from et al. 2022) using the same full- full-graph energy (2). MuseGNN performance significantly higher; see a rxx and products results in Table As for LazyGNN, results competitive on the small- and medium-sized datasets, but deteriorate rapidly on pape rs 1 0 OM and run OOM on the largest datasets. We note that all methods relying on GAS (including LazyGNN) experience degradation as the data sizes increase. See Appendix A for more USP/Lazy details.
Table 3: Training speed (epoch time) in seconds; hardware configurations in Appendix A
Timing Comparisons. Turning to model com- plexity, Table 3 displays the training speed of MuseGNN relative two common GNN baselines as well as LazyGNN. From these re- sults, we observe that MuseGNN executes with similar epoch time to the GNNs (satisfying our second objective from above), and much more efficiently relative to LazyGNN. Note that without public code, we are unable to compare timing with USP.
papers- 100M
Model
MAG240M
IGB- full
SAGE (NS)
102.48
795.19
17279.98
GAT (NS)
164.14
1111.30
19789.79
LazyGNN
6972.38
OOM
OOM
MuscGNN
158.85
1370.47
20413.69
Convergence Illustration and Ablations. In Figure we show the empirical convergence of (4) w.r.t. Y during the forward pass of MuseGNN models ogbn papers100f for differing values of y. In all cases the energy converges within 8 iterations, supporting our choice of K = for experiments (and satisfying our third objective). Please Appendix B for additional ablations involving and the MuseGNN offline sampler.
7 CONCLUSION
In this work, we have proposed MuseGNN, an unfolded GNN model that scales to large datasets by incorporating sampled subgraphs into the design of its lower-level, architecture-induinng energy function. In so doing, MuseGNN readily handles graphs with ~ 10 or more nodes and high- dimensional node features, exceeding 1TB in total size. Moreover, this is accomplished while main- taining interpretable layers with desirable inductive biases, concomitant convergence guarantees, and competitive (and at times SOTA) accuracy on the largest available benchmarks. The latter sig- nificantly exceeds the capabilities of prior efforts to scale UGNNs.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Hongjoon Ahn, Yongyi Yang, Quan Gan, Taesup Moon, and David P Wipf. Descent steps of relation-aware energy produce heterogeneous graph neural networks. Advances in Neural Infor- mation Processing Systems, 45:38436-384444 2022.
Sébastien Bubeck et al. Convex optimization: Algorithms and complexity. Foundations and Trends@ in Machine Learning, 8(3-4):31-337,
Chen Cai and Yusu Wang. note on over-smoothing for graph neural networks. arXiv preprint arXiv:2006. 13318, 2020.
Jianfei Chen, Jun Zhu, and Le Song. Stochastic training of graph convolutional networks with variance reduction. arXiv preprint arXiv:1710.1056 2017.
Jie Chen, Tengfei Ma, and Cao Xiao. FastGCN: fast learning with graph convolutional networks via importance sampling. arXiv preprint 1801 2018.
Jiuhai Chen, Jonas Muelerr, Vassilis N. Ioannidis, Soji Adeshina, Yangkun Wang, Tom Goldstein, and David Wipf. Does your graph need confidence boost? convergent boosted smoothing graphs with tabular node features. In International Conference on Learning Representations, 2022a. hRtpstts:///penete
net/ Leee:/oooummmGGnniipeeeggg/
Qi Chen, Yifei Wang. Yisen Wang, Jianlong Chang, Qi Tian, Jiansheng Yang, and Zhouchen Lin. Efficient and scalable implicit graph neural networks with virtual equilibrium. In 2022 IEEE International Conference Big Data (Big Data), pp. 864 873. IEEE, 2022b.
Siheng Chen, Yonina C Eldar, and Lingxiao Zhao. Graph unrolling networks: Interpretable neural networks for graph signal denoising. IEEE Transactions on Signal Processing, 69:3699-3713,, 2021.
Wei-Lin Chiang, Xuanqing Liu, Si Si, Yang Li, Samy Bengio, and Cho-Jui Hsieh. Cluster-GCN: An efficient algorithm for training deep and large graph convolutional networks. In Proceedings of the 25th ACM SIGKDD international conference on knowledge discovery & data mining, pp- 257-266, 2019.
Benoît Colson, Patrice Marcotte, and Gilles Savard. Bilevel programming: survey, 4or, 3:87-107, 2005.
Patrick L Combettes and -han-Chrittophe Pesquet. Proximal splitting methods in signal processing. Fixed-point algorithms for inverse problems science and engineering, pp. 182, 2011.
I. Csiszár and G. Tusnády. Information geometry and alternating minimization procedures. Statistics and Decisions, Supplemental Issue Number 1:205-237, 1984.
Francesco Di Giovanni. James Rowbottom, Benjamin Paul Chamberlain, Thomas Markovich, and Michael M Bronstein. Understanding convolution on graphs via energies. Transactions Ma- chine Learning Research, 2023.
Matthias Fey, Jan Lenssen, Frank Weichert, and Jure Leskovec. GNNAutoScale: Scalable and expressive graph neural networks via historical embeddings. In International Conference Machine Learning, pp. 3294-3304 PMLR, 2021.
Fabrizio Frasca, Beatrice Bevilacqua, Michael Bronstein, and Haggai Maron. Understanding and extending subgraph gnns by rethinking their symmetries. Advances in Neural Information Pro- cessing Systems, 35:31376-31390, 2022.
Guoji Fu, Mohammed Haroon Dupty, Yanfei Dong, and Lee Wee Sun. Implicit graph neural dif- fusion based on constrained dirichlet energy minimization. arXiv preprint arXiv:2308.03.3 2023.
Johannes Gasteiger, Aleksandar Bojchevski, and Stephan Günnemann. Predict then propagate: Graph neural networks meet personalized pagerank. In International Conference on Learning Representations, 2019. URL htR..tttsst... /ttttps:://peerrviiiee http:/wwwwweeeeteeteeee -I-2H111T-A999

---

Published as a conference paper at ICLR 2025
Johannes Gasteiger, Chendi Qian, and Stephan Günnemann. Influence- mini- for graph neural networks. In Learning on Graphs Conference, pp. 9-1. PMLR, 2022.
Justin Gilmer, Samuel S Schoenholz, Patrick F Riley, Oriol Vinyals, and George E Dahl. Neural message passing for quantum chemistry. In International Conference on Machine Learning, pp- 1263-1272. PMLR, 2017.
Fangda Gu, Heng Chang, Wenwu Zhu, Somayeh Sojoudi, and Laurent El Ghaoui. Implicit graph neural networks. In Advances in Neural Information Processing Systems, 2020.
Will Hamilton, Zhitao Ying, and Jure Leskovec. Inductive representation learning on large graphs. Advances in Neural Information Processing Systems, 30, 2017.
Weihua Hu, Matthias Fey, Marinka Zitnik, Yuxiao Dong, Hongyu Ren, Bowen Liu, Michele Catasta, and Jure Leskovec. Open graph benchmark: Datasets for machine learning on graphs. Advances in Neural Information Processing Systems, 33:221 18-22133, 2020.
Weihua Hu, Matthias Fey. Hongyu Ren, Maho Nakata, Yuxiao Dong, and Jure Leskovec. OGB- LSC: A large-scale challenge for machine learning on graphs. vrXiv22103.09:: 2021.
Kezhao Huang. Haitian Jiang. Minjie Wang. Guangxuan Xiao, David Wipf, Xiang Song, Quan Gan, Zengfeng Huang. Jidong Zhai, and Zheng Zhang. FreshGNN: Reducing memory access via stable historical embeddings for graph neural network training. Proc. VLDB Endow., 17(6):1433-1486
2024. ISSN 2150-8097. doi:
0777783364810.34488888 URL https: /tdp:
14778/3648188 3648184
Steven M. Kearnes, Kevin McCloskey, Mare Berndl, Vijay S. Pande, and Patrick Riley. Molecular graph convolutions: moving beyond fingerprints. 2016.
J. Comput. Aided Mol. Des., 30(8):595-608,
Arpandeep Khatua, Vikram Sharma Mailthody, Bhagyashree Taleka, Tengfei Ma, Xiang Song, and Wen-mei Hwu. Igb: Addressing the gaps in labeling. features, heterogeneity, and size of public graph datasets for deep learning research. arXiv preprint arXiv:2302. 13522, 2023.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint /412:.44/2.69900 2014.
Thomas N. Kipf and Max Welling. Semi-Supervised Classification with Graph Convolutional Net- works. In Proceedings of the 5th International Conference Learning Representations, ICLR 17, Palais des Congrès Neptune, Toulon, France, 2017. URLhttps::/ /tsttp:://peenrvview...
Corumpid-SJUAayYO
Mingjie Li, Yifei Wang, Yisen Wang. and Zhouchen Lin. Unbiased stochastic proximal solver for graph neural networks with equilibrium states. In The Eleventh International Conference Learning Representations, 2022.
Qimai Li, Zhichao Han, and Xiao-Ming Wu. Deeper insights into graph convolutional networks for semi-supervised learning. Proceedings of the AAAI Conference on Artificial Intelligence, volume 32, 2018
Yao Ma, Xiaorui Liu, Tong Zhao, Yozen Liu, Jiliang Tang, and Neil Shah. A unified view graph neural networks as graph signal denoising. In Proceedings of the 30th ACM International Conference on Information & Knowledge Management, pp. 1202-1211, 2021.
Arkadi Nemirovski, Anatoli Juditsky, Guanghui Lan, and Alexander Shapiro. Robust stochastic approximation approach to stochastic programming. SIAM Journal optimization, 19(4):5774- 2009.
Kenta Oono and Taiji Suzuki. Graph neural networks exponentially lose expressive power for node classification. In 8th International Conference on Learning Representations, ICLR, 2020.
Xuran Pan, Shiji Song, and Gao Huang. A unified framework for convolution- graph neural networks. arXiv preprint 2rXiv:2107.10234 2020.
12

---

Published as a conference paper at ICLR 2025
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, segh-performance deep learning library. Advances in Neural Information Processing Systems, 32, 2019.
Emanuele Rossi, Fabrizio Frasca, Ben Chamberlain, Davide Eynard, Michael Bronstein, and Fed- erico Monti. Sign: Scalable inception graph neural networks. arXiv preprint arXiv:2004. /11000004.11..... 7:15, 2020.
Nino Shervashidze, Pascal Schweitzer, Erik Jan Van Leeuwen, Kurt Mehlhorn, and Karsten M Borg- wardt. Weisfeiler-Iehman graph kernels. Journal of Machine Learning Research, 12(9), 2011.
Ke Sun, Zhanxing Zhu, and Zhouchen Lin. AdaGCN: Adaboosting graph convolutional networks into deep models. arXiv preprint /ttv:/990.050000 2019.
Petar Velickovic, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Liò, and Yoshua Bengio. Graph attention networks. In 6th International Conference on Learning Representations, ICLR 2018, Vancouver, BC, Canada, April 30 May 3, 2018, Conference Track Proceedings, Van- couver, BC, Canada, 2018. OpenReview.net. URL https: http///eeeeeeeeeeeeeeeeeeeeeeeem net forum?
id xJXMpikC
Roger Waleffe, Jason Mohoney, Theodoros Rekatsinas, and Shivaram Venkataraman. MariusGNN: Resource-efficient out-of-core training of graph neural networks. In Proceedings of the Eighteenth European Conference on Computer Systems, EuroSys '23, pp. 144-161, New York, NY. USA, 2023. Association for Computing Machinery. ISBN 7781440394771 doi: 10.1145/3552366. 3567501.
Minjie Yu Wang. Deep graph library: Towards efficient and scalable deep learning on graphs. In ICLR workshop on representation learning on graphs and manifolds, 2019.
Felix Wu, Amauri Souza, Tianyi Zhang, Christopher Fifty, Tao Yu, and Kilian Weinberger. Sim- plifying graph convolutional networks. In International conference on machine learning, pp- 6861-6871. PMLR, 2019.
Qitian WU, Wentao Zhao, Chenxiao Yang. Hengrui Zhang, Fan Nie, Haitian Jiang, Yatao Bian, and Junchi Yan. Simplifying and empowering transformers for large-graph representations. Advances in Neural Information Processing Systems, 36, 2024.
Rui Xue, Haoyu Han, MohamadAli Torkamani, Jian Pei, and Xiaorui Liu. Lazygnn: Large-scale graph neural networks via lazy propagation. arXiv preprint arXiv:2302.01503, 2023.
Yongyi Yang, Tang Liu, Yangkun Wang, Jinjing Zhou, Quan Gan, Zhewei Wei, Zheng Zhang, Zengfeng Huang, and David Wipf. Graph neural networks inspired by classical iterative algo- rithms. In International Conference Machine Learning. pp. 11773-11783, PMLR, 2021.
Yongyi Yang, Tang Liu, Yangkun Wang, Zengfeng Huang, and David Wipf. Implicit unfolded graph neural networks. arXiv preprint arXiv:2111.00
2022.
Rex Ying, Ruining He, Kaifeng Chen, Pong Eksombatchai, William L Hamilton, and Jure Leskovec.
Graph convolutional neural networks for web-scale recommender systems. In Proceedings of the 24th ACM SIGKDD international conference on knowledge discovery & data mining, pp. 974 983, 2018
Jacmin Yoo, Meng-Chin Chich Lee, Shubhranshu Shekhar, and Christos Faloutsos. Less is more: Slimg for accurate, robust, and interpretable graph mining. In Proceedings of the 29th ACM SIGKDD Conference Knowledge Discovery and Data Mining, pp.
2023.
Haiyang Yu, Limei Wang, Bokun Wang, Meng Liu, Tianbao Yang. and Shuiwang Ji. Graphfm: Im- proving large-scale gnn training via feature momentum. In International Conference on Machine Learning, pp.
PMLR, 2022.
Hanqing Zeng, Muhan Zhang, Yinglong Xia, Ajitesh Srivastava, Andrey Malevich, Rajgopal Kan- nan, Viktor Prasanna, Long Jin, and Ren Chen. Decoupling the depth and scope of graph neural networks. Advances in Neural Information Processing Systems, 34, 2021.
13

---

Published as a conference paper at ICLR 2025
Hongwei Zhang, Tijin Yan, Zenjun Xie, Yuanqing Xia, and Yuan Zhang. Revisiting graph convo- lutional network on sumi-supe node classification from an optimization perspective. arXiv preprint arX1v:2009.11669, 2020.
Amber Yijia Zheng, Tong He, Yixuan Qiu, Minjie Wang, and David Wipf. Graph machine learning through the lens of bilevel optimization. In International Conference on Artificial Intelligence and Statistics, pp. 982-990. PMLR. 2024.
Dengyong Zhou, Olivier Bousquet, Thomas Lal, Jason Weston, and Bernhard Schölkopf. Learning with local and global consisteney. Advances in neural information processing systems, 16, 2003,
Hao Zhu and Piotr Koniusz. Simple spectral graph convolution. In International conference learning representations, 2020.
Meiqi Zhu, Xiao Wang, Chuan Shi, Houye Ji, and Peng Cui. Interpreting and unifying graph neural networks with an optimization framework. In Proceedings of the Web Conference 2021, pp- 1215-1226, 2021.
Difan Zou, Ziniu Hu, Yewen Wang, Song Jiang, Yizhou Sun, and Quanquan Gu. Layer-dependent importance sampling for training deep and large graph convolutional networks. Advances in Neural Information Processing Systems, 32, 2019.
14

---

Published as a conference paper at ICLR 2025
A EXPERIMENT DETAILS
Dataset Statistics. Table 4 contains the details of the datasets we use for the experiments in Sec-
tion and the ablation study in Appendix B.2
Table 4: Dataset details, including node feature dimension (Dim.) and number of classes (# Class).
Dataset
[
[|]
Dim. # Class Dataset Size
ogbn-arxiv (Hu al. 2020
0.17M
1.2M
128
40
182MB
IG3-tiny (Khatua el al. 2023
0.1M
0.5M
1024
19
400MB
ogba-products (Hu et al.
(2200)
2.4M
123M
100
47
1.4GB
IG3-medium (Khatua al. 2023
10.0M
120M
1024
19
40,8GB
ogon-papersTU001 (Hu
2020
111.1M
1.6B
128
172
57GB
MAG240M (Hu al. 202T
244.2M
1.7B
768
153
377GB
IGB-full Khatua ct al.
2023)
269.3M
4.0B
1024
19
1.15TB
Model Details. For MuseGNN, we choose the base model f(X; W) and output function gY;; 0) to be two shallow 3-layer MLPs with standard residual skip connections (although for the somewhat differently-strueee strrutuuey ogbn sigon-poodee productss data we found that skip connections were not necessary). The node-wise constraint 17 in the energy function 4 is set be non-negativity. making the proximal operator prox,, in the model (8) to be ReLU. The number of unfolded layers, K, is selected to be 8, and the hidden dimension d is
512. The forgetring factor p for the online mean estimation is set to be 0.9. All the models and experiments are implemented in PyTorch (Paszke et al.) 2019)
using the Deep Graph Library (DGL) /Wang 2019).
Offline Samples. As supported by Proposition 3.1 subgraphs induced from nodes (where all the edges between sampled nodes are kept) pair naturally with the energy function. Therefore we use node-induced subgraphs for our experiments. For ogbn-arxiv, we use the 2-hop full neighbor because it is a relatively small graph. For IGB-t iny, peion-papers/ O0M, and MAG24 OM, we use Shadow>HHop with fanout [5,10,15]. For the rest of the datasets, we use Shadow Koop with fanout [10 [10, 15]. Note that Shadow KHop sampler performs node- pode-wise neighbor sampling and returns the subgraph induced by all the sampled nodes.
Training Hyperparameters. In the training process, we set the dropout rate of the MLP layers to be 0.2, and we do not have dropout between the propagation layers. The parameters are optimized by Adam optimizer (Kingma & Ba 2014) with the weight decay parameter rate being 0.001. For ogbn-arxiv v and ogbnnpaperss00k a = 0.05, A = 20. For MAG24 4 CM and the IGB series datasets, a 0.2, 4. And for ogbn-products, a À with precon- ditioning on the degree matrix for each unfolded step/layer. The batch size n', in the offline samples are all set to 1000. For large-scale datasets, an expensive hyperparameter grid search as commonly used for GNN tuning is not feasible. Hence merely applied simple heuristics informed from training smaller models to pick hyperparameters for the larger datasets. Also it is worth noticing that should not in principle affect accuracy if set small enough, since the model will eventually converge. In this regard, a can set dependent on X, since the latter determines the size of a needed for convergence.
0 and the learning
Evaluation Process. In the presented results, the training, validation and test set splits all follow the original splits from the datasets. The evaluation process for validation and test datasets uses the same pipeline the training process does: doing offline sampling first and then these fixed samples to do the calculation. The only difference is that in the evaluation process, the backward propagation is not performed. While used the same pipeline for training and testing, MuseGNN modular and we could optionally use online sampling for the evaluation or even multi-hop full- neighbor loader in the evaluation time.
Configurations in the Speed Experiments. We use single AWS EC2 24d.24xlarge instance to run the speed experiments. It comes with dual Intel Xeon Platinum 8275CL CPU (48 cores, 96 threads), 1. 1.I main memory and 8 A100 (40GB) GPUs. All the experiments run single
15

---

Published as a conference paper at ICLR 2025
GPU. When running MuseGNN, the pre-sampled graph structures are stored in the NVMe SSDs and the input feature is loaded into the main memory (for IGB-full, its feature exceeds the capacity of the main memory, so mmap is employed). In the baseline setting, ogbn-papers seneers0000 100M and MAG2 4 DM are paired with neighbor sampling of fanout [5, 10, 15] and TGB- is paired with neighbor sampling of fanout [10, 15]. So the fanout for the online neighhor sampling baselines and MuseGNN are the same so that they can have a fair comparison. In these experiments, the hidden size is set to 256 as is typical.
Additional Details Regarding USP and LazyGNN Results. Since the source code for USP is not available, we can only cite the published numbers that are directly applicable Table 2 For a fair comparison, we choose USP results listed as "IRLS" as the base model in (Li et al. 2022
This model corresponds to minimizing the energy function ( which is equivalent to what is used for our full-graph unfolded GNN baseline, UGNN (FG), and the original basis for MuseGNN. Note that MuseGNN and USP can both handle alternative energy function forms to potentially improve performance depending on the application.
As for LazyGNN, its training pipeline incorporates METIS partitioning similar to other GAS-based methods. The more partitions used for this, the smaller the subgraphs are (so that they can fit into the GPU memory), and the more iterations it takes for training one epoch. For obtaining oobnn and ogbn-proodee products results, the partitions are only 60 and 150. However, for oobnn OM, as this dataset is much larger in size, the number of partitions is set to 40000 so that the subgraphs can fit into GPU memory. But GAS- GAS-based LazyGNN also uses historical em- beddings and historical gradients, and the staleness of these variables can become very large due to higher number of iterations per epoch needed with so many partitions. Hence this staleness could partially account for the relatively low accuracy this aspect of Easebaaed models can be found in Huang et al. (2024). And finally, for obtaining the accuracy results of LazyGNN, ran the exact open-soureed code with the hyper-parameters provided by the authors for each overlapping dataset they considered (i.c., arxiv and product s), and found that it is actually maximum accuracy values (over 5 trials) that matches what is reported in their paper. However, for consistency within our experiments across all methods, we instead report the LazyGNN mean accuracy values with error bars in Table 2
ogbn-papers 1 00M. Related discussion of
B ADDITIONAL EXPERIMENTS AND DISCUSSION
B. ACCURACY COMPARISONS WITH ADDITIONAL BASELINES
For more general comparison beyond Table 2 in the main paper, this section introduces additional baselines, and in particular, some recent models explicitly designed for scalability (although none of these models are UGNNs). Specifically, we include SGC (Wu et al. 2019). SIGN (Rossi et al.) 2020), Cluster-GCN (Chiang et al., (019)) and GraphFM (Yu et al.) 2022). Note that GraphFM is built on top of GAS (Fey et al. 2021), while the others (i.c., SGC, SIGN, ClusterGCN) represent traditional GNN architectures that are naturally amenable to larger graphs. Results using these ad- ditional baselines, as presented in Table further solidify the competitive advantage of MuseGNN. Additionally, all baseline results are from OGB leaderboards or published papers (Huang et al. 2024 Zhu & Koniusz, 2020). In this regard, do not report performance on IGB datasets as, being quite new, most prior work does not contain such results and it is extremely expensive to run ourselves.
B.2 FURTHER DETAILS REGARDING THE ROLE OF y IN MUSEGNN
The incorporation of 0 serves two purposes within our MuseGNN framework and its supporting analysis. First, by varying 7 we are able to build a conceptual bridge between the decoupled 7 base scenario, and full-graph training as 7y co. provided the sampled subgraphs adhere to the conditions of Proposition 3.1 In this way, the generality of y > 0 has value in terms of elucidating connections between modeling regimes, independently of empirical performance.
http:::/gitthbb.com//////////
16

---

Published as a conference paper at ICLR 2025
Table 5: Node classification accuracy (%) results using additional non- non-UGNN baselines. Reported numbers are from the test set, except for MAG24000 which only has labels for the validation set.
Model
arxiv products |apers1000|MAAG24000 MAG240M
SGC
68.78
68.87
63.29
65.82
SIGN
71.95
80.52
65.68
66.64
Cluster-GCN
68.11
78.97
53.35
OOM
GraphFM
71.53
70.76
48.03
OOM
MuseGNN
72.50
81.23
66.82
67.67
That being said, the second role is more pragmatic, as y > 0 indeed improve predictive accuracy. As shown in the ablation in Table 6 (where all hyperparameters except remain fixed), y = 0 already achieves good performance; however, increasing Y leads to further improvements. This likely because larger enables the model to learn similar embeddings for the same node in different subgraphs, which may be more robust to the sampling process.
Table 6: Ablation study of the penalty factor y. Results shown represent the accuracy (%) the test set. Bold numbers denote the best performing method.
Y 7 0.1
7 0.5
7
7
papers100M
66.29
66.31
66.53
66.45
66.64
66.82
ogbn-arxiv
72.51
72.67
72.47
72.21
72.15
72.08
IGB-tiny
72.66
72.78
73.42
73.02
72.70
72.69
IGB-medium
75.18
75.80
75.83
N/A
N/A
N/A
ogbn-procucts
80.42
80.92
81.23
N/A
N/A
N/A
That being said, while memory cost for storing the required mean vectors. However, this memory complexity for the mean vectors only (nnd), which is preferable the memory complexity ( of GAS (Fey al., 2021), which depends on the number of layers K. Importantly, this more modest ((nd) storage only an upper bound for the MuseGNN memory complexity, as the 7 case is already very effective even without requiring additional storage. As an additional point of reference, the GAS- related approach from Xue et al. (2023) also operates with only single-layer ( of additional storage; however, unlike MuseGNN, this storage is mandatory for large graphs (not amenable to full-graph training) and therefore serves as a lower bound for memory complexity. Even so, the approach from Xue al. (2023) is valuable and complementary.
improve accuracy with y O, this generality comes with additional
B.3 SAMPLING ABLATION ON THE BASELINE MODELS
We choose neighbor sampling for the baselines hecause is most commonly used, but since Shad- ow owKHop is paired with MuseGNN, use Table show that the improvement is not coming from the change the sampling method but the usage of mergy-based scalable unfolded model. In Ta- 7 the bascline models are trained samples used in the training of MuseGNN, but they suffer from decrease in the accuracy compared with the online neighbor sampling counterparts, so the change sampling method cannot account for the enhancement in the accuracy.
gebbn-ppeers OM with the same offline ShadowKHop
Table 7: Ablation study of the accuracy improvement compared with baselines. Shown results are the accuracy (%) the test set of oobn-paperane papers100M.
Baseline Model
Sampling Method
GCN GAT SAGE
64.89 63.84 65.57 Neighbor Sampling (online) 65.83 66.28 66.2
ShadowKHop (offline)
17

---

Published as a conference paper at ICLR 2025
B.4 ALTERNATIVE SAMPLING METHOD FOR MUSEGNN
We also paired MuseGNN with online neighbor sampling, the most popular choice for sampling, on ogbn-papers to account for our design choice of sampling methods coupled with MuseGNN. Theoretically, the convergence guarantee no longer holds true for online sampling. Ad- ditionally, neighbor sampling removes terms regarding edge information from the energy function, so expect it to be worse than the offline ShadowKHoop sampling that usc. This is also sup- ported by Proposition becauee neighbor sampling does not provide subgraphs induced from the nodes with all the edges. Empirically, with online neighbor sampling, the test accuracy is 66. 111% %;; it is indeed lower than the 66.82% from our offline ShadowKHop results.
B.5 DIFFERENCES BETWEEN MUSEGNN AND USP
Despite both being methods for scaling UGNN training, the USP approach (Li et al. 2022 and MuseGNN differ multiple important respects. First, USP is based on an unbiased stochastic proximal solver and partitioned subgraphs to approximate the gradient steps used for minimizing an original full- full-graph energy as in ( Quite differently, MuseGNN actually predicated on defining new subgraph- based energy function given by (4) that is distinct from its full-graph counterpart. See Section
C below for discussion of the potentially greater expressiveness of 4 relative to (
Secondly, USP convergence has only been established for the forward pass (i.e., lower-level energy function minimization), with no guarantees provided for full convergence in conjunction with the backward pass (i.e., minimizing the upper-level training loss w.r.t model parameters). In contrast, we establish MuseGNN convergence under certain conditions across the full bilevel optimization process, simultaneously covering forward and backward passes.
And lastly, there does not as of yet exist any published demonstration that USP is empirically ef- fective large graph datasets, In fact, sgbn-prrodctts. is the largest benchmark adopted in Li et al. (2022), which is viewed as medium-sized relative the other datasets tively trained on ogbn-proodue -products s using a single GPU, so this dataset does not generally require intensive scalability measures. This is unlike MuseGNN, which we have validated on the largest publicly-available benchmarks. We also note that et al. (2022) presents a compelling acceleration method reduce training iterations; however, the per-iteration computational complexity of this step is comparable to full-graph training and therefore is not suitable for the largest graphs.
public code available for implementing such demonstration.
consider. Note that full -ill-graph GNNs can be effec-
C MUSEGNN EXPRESSIVENESS RELATIVE TO FULL-GRAPH UNFOLDED GNNS
In this section we consider the expressiveness of MuseGNN versus full graph unfolded GNN models. As suggested Section 3 there exists non-isomorphic graphs such that the induced full-graph energies from 2 have effectively equivalent minima, and yet the corresponding { (Y, } === from (6). as a special case of the MuseGNN energy (4) with y = 0, do not. This would imply that the subgraph-based energy of MuseGNN can actually be more expressive in some sense. We first formalize this notion and then subsequently provide illustrative examples.
C.I FORMAL ANALYSIS
The following proposition demonstrates that for two non-isomorphic graphs that are considered as isomorphic by the Seisfeiler-Lehmann (WL) test (Shervashidze et al., (2011) minimizing the full- graph energy will yield equivalent embeddings, even when initialized in such a way that the actual descent iterations need not follow the WL criteria which underpin the test itself.
Proposition C.1. Suppose two attributed graphs G1 = E1,, and G2 = (V2, X2), al- though possibly non-isomorphic, cannot be distinguished non-isomorphic by the WL test. Fur- thermore, w.L.o.g. assume the node indices of 91 and 92 are aligned such that each node with the same index has the same final hash value produced by the WL test, i.e., the multi set of their neigh- borhoods are identical (if this were not the case, it can be trivially achieved via permutation of node indices, an operation that does not impact the energy (2)) Now consider optimizing ( using either
18

---

Published as a conference paper at ICLR 2025
graph with a convex Ç. Beginning from arbitrary initial embeddings y1 and Y respectively, not necessarily equal to one another nor adhering to WL test stipulations, the unfolded GNN descent iterations from (3) will converge to minimizing solutions Yi and Y2 satisfying Yi =\2.
Proof. Since 91 and G2 are considered isomorphic by WL test, we have X1 X2. L,t = X2 X. The energy functions for G1 and G2 are then & (1( || [(X: Atr Y + E <(X, and (2 = || - f(x; W)ll + A tr(YY + ri ((Y)) respectively, where L and Lg denote the graph Laplacians associated with 91 and G2.
We first prove that within one graph. the nodes with the same final WL hash value will have the same embedding. Since here the energy function is strongly convex. the initialization can be arbi- trary and the forward pass will finally reach the same global minimum. Now suppose we initialize the embeddings as [(X; W). In this case, the nodes with same WT. hash values have the same initial embeddings. Now consider each proximal gradient descent step on the energy function, i.e., the forward propagation layer of the unfolded GNN. As the neighborhood multi-sets of the nodes with same WL hash are exactly the same, for these nodes in every step, the aggregation and update functions take in exactly the same inputs. Therefore in every step these nodes have the same em- beddings, thus at the final minimum, the embeddings are still the same. And since any initialization leads to the same minimum, this will be true regardless of the descent steps that led there.
Now we prove the proposition regarding two graphs. Considering the first order optimality condition of the energy & (Y),, we have - -YY - f(X:W) + ALY?) E @c(YT), where a((Y,) denotes the subdifferential of Ç evaluated at Y = Yr- With the assumption on the alignment of node indices, we want to show that Y = Yg. which is equivalent to Y satisfying the first order condition of (Y). Therefore, it is sufficient to show when -(Yi* - f(x: W) + X E o<(Y;') + E "((Y"). Since multiplying the Laplacian matrix and the embeddings
fX"- - essentially doing one step of neighbor aggregation, we have L = LgY,". This because the 1-hop neighborhood of each node in 9i and G2 have the same hash value combinations as guaranteed by the WI. test, and each node with same hash value has the same embedding in YT as we proved above. Therefore - -(Y" f(X; W) + AL)Y in (Y+)). Thus, Y is also the minimum for (Y) so Y Ygg
(Yr ((X;W)+++ \ they both exist
Interestingly though, something equivalent to Proposition C.I need not hold for the MuseGNN energy based sampled subgraphs. To see this, we rely analysis from (Frasca et al. 2022 which reveals that the expressive power of GNNs trained on subgraphs can actually surpass the expressive power of the WL test. In setting, it can be shown that this translates into more expressive minimizers of the MuseGNN energy 4 particularly when = 0. To intuitively convey this phenomena, we next lean on illustrative examples whereby minimizers of 4 can distinguish two graphs that were not originally distinguishable by the WL test or minimizers of
C.2 ILLUSTRATIVE EXAMPLES
To begin, we consider pairs of non-isomorphic graphs that have been initially assigned identical input features across all nodes in both graphs For simplicity of presenting useful illustrative ex- amples, we also base our assumed energy functions on symmetric normalized graph Laplacians
D = 1 - D AD- Note that Proposition C.1 still holds equally well under this revision.
Starting with uniform initial node features, the WL test will ultimately identify the two pairs of non-isomorphic graphs in Figure 2a as isomorphic. And with the WL test result, we can match each node in the left graph with corresponding node in the right graph if they have the same final hash value in the WL test. A bijection can be established for such corresponding nodes because these two graphs are considered isomorphic by the WL test.
If we use the left-right pair of graphs in Fadd uniform initial features to form the energy (2) subsequent optimization using (3) will produce the final embedding minimizers indicated by the node colors in Figure 2a Here the same color indicates the same embedding vector of the minimizer. As the corresponding nodes from the legt- pair have the same colors, the minimizers of the corresponding full-graph energies are equivalent. Thus the left-right pair cannot be distinguished. Such cases resonate with the implications of Proposition C.1
19

---

Published as a conference paper at ICLR 2025
(a) Same colors indicate these nodes have the same embedding vector upon min- imizing the full- bull-graph energy from (2) us- ing the corresponding graphs. The non- isomorphic graphs within each left-right pair cannor be differentiated.
(h) Based on subgraphs graphs on the left, minimizers of re- vised energy akin (6) and by extension fact differentiate each non- isomorphic graph. noted by distinet color patterns within each left-right pair.
original
Building on analysis from Frasca al. (2022), it is possible achieve increased expres- siveness via energy functions based on sampled subgraphs (as incorporated by MuseGNN).
Figure 2:
However, we sample subgraphs from the full graphs, we are now able to distinguish such pairs. By taking the subgraphs of the left-right pair in Figure 2a obtain the graph structures in Figure 2b But now, even with the uniform node features, the final embeddings minimizing the graph energy (6) will yield different distributions for the left-right pair, thus differentiating the two graphs. The final embeddings are illustrated as colors in Figure 2b As such, the MuseGNN energy function 44) based on sampled subgraphs in this sense more expressive than the full full-grap graph energy alternative from 2 as minimizers of the former can differentiate some non -sommoopppicc graph pairs that are not distinguishable by the minimizers of the latter.
D ADDITIONAL CONVERGENCE DETAILS
D.1 EMPERICAL CONVERGENCE RESULTS
Here we show the empirical results of the previous convergence analysis. Though we do not have the full bilevel convergence results for the 7 case like Theorem [.22] Figure 3 shows that in real-world dataset, the bilevel optimization system will still converge with the imprecise estimation of optimal embeddings y, (K) Y and the online mean estimation of M.
3.5
3.0
2.5
8
2.0
1.5
1.0
250
500
750
1000
1250
1500
1750
Iteration
Figure 3: Convergence of the upper-level loss ogbn arxiv dataset for 20 epochs with penalty
factor 7 = L.
20

---

Published as a conference paper at ICLR 2025
D.2 A NOTE ON SAMPLING METHODS IMPACTING CONVERGENCE RATES
As stated in the main paper, offline sampling allows conduct formal convergence analysis that is agnostic to the particular sampling operator. So our convergence guarantee is valid for all sampling methods. However, different sampling methods can still actually have an impact on the convergence rate of the bilevel optimization. This is because different sampling methods will produce subgraphs with different graph Laplacian matrices, and the latter may have differing condition numbers as mentioned later in Lemma [E.4] In this regard, a larger such condition number will generally yield slower convergence rate. Further details are elaborated in the proof of Theorem 5.2 in Appendix E.2
E PROOFS
E. CONNECTION WITH FULL GRAPH TRAINING
Proposition 3.1. Suppose we have m subgraphs (V,,E))... (Vm, Em) constructed independently such that Vs 1,.... m,, Vu, V, PrPr(v E V.] = Priu E V, u E V.l pi (i, E EE, > i E V.. E Vz, (i.j) E E. Then when y = oc, we have Elemase (M) = mp {(M) with the x in ((M) rescaled to pd.
Proof. The probability that an edge is sampled is Pr[(i.)) E & (1,5) E El = p². Ve E V. E V2, V. then Yew Ye,- M.. We have
m Σ
mass (M)
Σ (IMM, f(x; (Xx:W) -(xx + Σ (KJ)EE,
I|M M,l3
NEVi
E musse (M) =
[IM - -(x; ((NN Pr/u V.l
s=1 WE V.
+
Σ
|M, M3l12 Pr(((..) EC
(i,))
(|AM/- - f(x; f(x; +(1 + +mp?
Σ
| - M
=mp
-EV
(Ai))E
=mp
||
pAtr LM) Σ ((Mi)
=mp ((M)
E.2 FULL CONVERGENCE ANALYSIS
According to Definition 5.1 (,) = (I AL.) (L+ == P.f(X,; W) is the optimal
embedding for the subgraph energy Li (Y.). Additionally, we initialize () f(X.; W), so (*()) can be written as (((() f(X; W) where (((k) also a matrix. We have a < ||4 >Lal so p(() will converge to P- as k grows. The output of D over matrix input is the of D over each row vectors, as typical discriminator function like squared error or cross-entropy do.
Plugging the y: (W) and Y into the loss function, we have
(k) muse (W)= = Σ D D ( , (III)xw.))
(12)
21

---

Published as a conference paper at ICLR 2025
where 1 : ni, means the first ni, rows, = (Ti, (Ti,T2, .2..... T Ip(k) = diag (p!
limg
Similarly, for lw (W), we change all the Yd and k!k) to Y and P. That is,
p = diag ( (
s=1 muse (W) D (P*XWTT)
Stochastic Gradient Descent
The updating rule by gradient descent for the parameter is
w w(()) v wto C/b (Wi(t)
0D Y. (k) (Wit),
W/
OW
In reality, we use stochastic gradient descent to minimize C (W), ) and the updating rule becomes
(++++) W(0 0D (T..9 (Y. (w()))
(13)
OW())
Here is picked random from (1, .... m). the gradient is unbiased estimator of the true gradient.
Lemma E.1. As defined in 112 L (W) is a convex function of W. Furthermore, there exist the global optimal W(k)), such that ((k) (W(())) I Cinise (W) for all W.
(k)
Proof. C(k) W) is a composition of convex function D with an affine function, so the convexity is retained. By the convexity, the global optimal (*k)) exists.
Lemma E.2. When the loss function C is defined as equation 12 and the parameter W
updated by 13 with diminishing step sizes te ((t)), 1E Camis (miic(((()) ((k) (W(k)))
O(1/VI).
Proof. Since the loss function
mmise is convex by Lemma B.I and the step size is diminishing in
(k)
( by (Nemirovski et al. 2009 the convergence rate of the difference in expected function
value and optimal function value
0 (
Lemma E.3. The subgraph energy function 6 with ((v) = 0 (as defined in Definition 5.1 is dg- smooth and St-Sonyyyly convex with respect to Y., with XLs), where dmax and Jmin are the maximum and minimum singular value of the matrix.
mmax + Xba) and T. T +
Proof. The proof is simple by computing the Hessian of the energy function. Additionally, since the graph Laplacian matrix L, is positive-semidefniit. we have a 2 T. 0.
Lemma E.4. Let G arg max... ox 5, = 1,2, m, where gives the worst condition
number of all subgraph energy functions. In the descent iterations from (8) that minimizes with step size a a' we can establish the bound on the propagation matrix [[P(()) P
m exp(
20 Em ==1 (P()) P
Proof. By Bubeck al. (0155) ( Theorem 3.10], we have
P f(X; P: P: Se I p ( f(::W) - P. (X)
Since this bound holds true for any ((X:W), by choosing f(x; W) I to be the identity matrix,
we have
p!k)) - P:l Se -z
p(u)
p*
Se -
p(0) - P:
22

---

Published as a conference paper at ICLR 2025
Adding up all the subgraphs, we have
p( p
S me 25 Σ p(0) P.
Theorem 5.2. Let W be the optimal value of the loss Limuse (W) per Definition 5.1 while W de- notes the value of W after steps of stochastic gradient descent over ( (W) with diminishing step sizes O 7t Then provided we choose E min, I| + XL 112 and Y = f(xs: W),
for some constant C we have that
E
[Chdde (W())]
fimuse (W*) -- SO ( + + Vt
1
Proof. In Definition 5.1 we assume to Lipschitz continuous for the embedding variable, so for any input Y and Y we have [D(Y, T) D(Y", T)) S. LD Y'
((k)
(W(lk))
Limuse (W*)
(k) muse
(W*)
Liuso (W*)
xp= .T) D (P*XW",
=D
SLD
[P(k) XI" - -P''-'
pp(K) | |XW|||
LDD
<0 (e #
Therefore,
E
(()) (W(1))
Limuse (W²)
R
muse (W]l)))
L(kk) (W(())) + L (W(k))) Limuse (W*)
<0(
+ (
V
<0(
+ +e cck
Vt
where
E.3 ALTERNATING MINIMIZATION
The main reference for alternating minimization is Csiszár & Tusnády (1984). In this paper, they proved that the alternating minimization method will converge to the optimal value when five-point property both a three-point property and four-point property holds true. We will show that the global energy function and the corresponding updating rule satisfy the three- three-point point property and the four-point property. Thus, the alternating minimization method will converge to the optimal value.
For simplicity, we define ra,i as rx,i = Σm= Ej I{T(,, i) = ((s,,)) where I[..} is the indica- tor function. Therefore, Tai means the node in the full graph with index / ((s, i) appears ra times in all subgraphs. With some abuse of the notation,
let TI(a.) T
In the energy function (4) we still want to find set of {Y.)m and {p. \m to minimize the global energy function. We can use the alternating minimization method to solve this problem. In each step, we minimize (Y.) ym first and then minimize {m We want to show that when the
23

---

Published as a conference paper at ICLR 2025
parameter W is fixed, by alternatively minimizing {Y,)== and { the energy will converge to the optimal value.
We can easily get the updating rule for [Y,)m and {Ug } by taking the derivative of the energy function. Note that we always first update {Y, im and then update {1a]==1 For {Y,}, we have
y(k)
[[1 + 7) + - + +
(14)
For 1 we have the i-th row of ue is
(8)
[[( = I($
(15)
Namely, the updated (k) is the average of the embeddings of the same node in different subgraphs.
Note that here {Y,} with { and with M are used simultaneously. We will use (Y.) and {1e) when we want to emphasize the subgraphs and use Y and M when we want to emphasize them as the input of the global energy function. But essence they represent the same variables and can be constructed from each other.
The three-point property and four-point property call for non-negative valued helper function 8(Y, such that 8(Y, Y) O. We define the helper function as Y(Y, (1+ >) I' We can easily verify that 8(Y, Y) = and 0(Y, Y') 2 O.
Lemma E.5 (Three-point property). Suppose that we have series of y(k) and M(,) k
0, 1,,,.... constructed following the updating rule 14 and (15) and that they are initialized ar- bitrarily. For any Y and k, muso (Y, (x,H(()) (R) lmuae (Y
(R)
8(Y, V(V+)()1)) holds true.
Proof. We only need to show fiuse (Y.:1% (k)) - l (y,
(k)) 2 8(Y,, y((+D), By iterating from m and adding the m inequalities together, we can get the desired result.
Emsse (Y,, (k)) - m (x+1))), ,1(()))))) - 5(Y,, (((+t)))
= =||Y, f(xx: WI Atr (Y. LoYs) + ? |r -xxa)) 1 | f(x:: W) ||
-Atr (Y Loy (o(()))))TL - y (k+)) | (1 +1+ I y
\r. -y(++l) + y((++1) f(X,; W) 1/2 +Atr (Y I.Y.))
+ +7 +7||rs 1(h-1) (((+1) - ! 2
|||x -f(Xs: w) ?
-Atr ( La (((ttt) >7 | ( 1/ (1 + 7)) |
y((+1) -Y(
-2 YY, -2(v.
y(k+1), y(k+1) f(x,; W)) + 22 (Y, - -Y(kt)), y(k+1) - mS(1)
+ +att (Y. L.Y.)) - Atr ( (((K+I)T L (a+11)
=2 Y y(u+)), (1 + >)) (f(Xs: W) + + ) + + (Y, Yy, y,(k+1),, Lo (Y, y + y((+t)))
Y y((+1), 2 (eis + (1+ (1+7)) 7) y.(+2) -2 (f(X.: W) + +)x(tt) + >lxa (Y, y(lt2)))
Yo 1((+1), Aae (Y, - -Y! 00
Lemma E.6 (four- point property). Suppose that we have series of (k) and (K)), k k= 0, 1, 2... constructed following the updating rule (14) and (15) and that they are initialized arbitrarily. For any V. p and k 8(V, ((()) emae (M, 11 E Emase (7 p)) holds true.
24

---

Published as a conference paper at ICLR 2025
Proof. Expanding all the functions we can get our target is equivalent to
Ea+ni - y 2 v Σ
(1xs-
m1
|
-118s Mollz
We can actually show a stronger result, that is
E (A)
m Σ
(1%s-
-1111, 1e ma1
By the updating rule for Ma we have
m
m
n2
1 11's (K) y,(1)
M/. 1 p, y(())
(k) (k)
Us
Ma,i
[I(s,i) = 1) (ML -M!h) , Y1!h))
(16)
8=1 1=1
M M ,M 1
[I(s, 1) w)x((()
To
Tv
1=1
Therefore,
((b) Y, - x
2 = Σ
1Y, - Mal? + | -2,00 1|2 r
+2 (Yo Ho.Ho Y230))
|10 - u 11- F
+ |11150 - 1 + +2(Xo Y Us Hb - x(4ii)
||Ys - ue +
m
2 44 m (1) (N) -Y!1))
16
2 2||Y. - mal + ||4s- - 12. 1/- + ||1 - x 13
m
(LYs - Ms 14x - y(hi) + (1rs Vs, Y, - 11 (1))
+ 2
s=1
m
21 - -Hll + ||1r- - |/20 || + |1(50) y- ||²
Σ
(-2111 mall? + 2 22YY= ( - -v((+))
Y1 - (91 + (n( y.uus
+ 2 (Ys us. (k) - y(ai))
p
Σ
(116- [1)) 2
1|3 ++
|u(+) Y((k)
+ Y2
uo
2
I| Hell?
F
[||s- - -1(h)) | I|Y, - poll2
Theorem 5.3. Assume ((v) = O. Suppose we have a series of ((k)) and M(A), == 0, 2, .... constructed following the updating rule (14) and (15), with y(o) and M initialized arbitrarily. Then
lim Emcie ( (B(K),,()))
inf Emure (Y, M).
k k0000
Y.M
25

---

Published as a conference paper at ICLR 2025
Proof. By Lemma E.5 and Lemma E.6 where both the three-point and four cour-point property hold true, the theorem obtained according to Csiszár & Tusnády (1984).
F LIMITATIONS
While MuseGNN was primarily designed for scaling the most common UGNNs on homogeneous graphs there do exist prior UGNNs models based energy functions sensitive heterogeneous graph structure. The basic idea to update the trace term in Equation (4 to include additional trainable weight matrices that serve align embeddings of nodes of different types and relation- ships. HALO (Ahn et al. 2022) is one such example of this. But there is nothing to prevent us from extending the core techniques and analysis that undergird MuseGNN to scale such heterogeneous cases.
G IMPACT STATEMENT
While we do not envision that there are any noteworthy risks introduced by our work, it is of course always possible that large-scale model like MuseGNN could be deployed, either intentionally or unintentionally. in such way as to cause societal harm. For example, a graph neural network designed to screen fraudulent credit applications could be biased against certain minority groups
26


Reviews:
[{"review_id": "cQkqP9yrSw", "reviewer": "Reviewer_y5VQ", "summary": "The paper proposes a GNN model  that scales effectively to large datasets by incorporating sampled subgraphs into its energy function design. This approach allows the model to handle graphs with around 100 million nodes and high-dimensional features, achieving competitive accuracy on benchmarks exceeding 1TB in size. Additionally, it maintains desirable inductive biases and convergence guarantees.", "questions": "Refer to the content in the Weaknesses.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1). The organization and writing of this paper are excellent.\n\n2). The model enhances UGNN framework by incorporating sampling into its energy function design and demonstrates solid convergence properties.\n\n3).  Empirical results suggest that MuseGNN maintains competitive accuracy and scalability across various task sizes, performing well on large node classification datasets exceeding 1TB.", "weaknesses": "1). How to determine \\alpha and \\lambda for different datasets?\n\n2). Does the model's improvement differ for homogeneous and heterogeneous graphs? Can you provide some additional explanation?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "iRp8nSLBxa", "reviewer": "Reviewer_RrLb", "summary": "This paper introduces MuseGNN, a novel Graph Neural Network (GNN) architecture designed to address the challenges of scaling and convergence in GNNs. The core idea behind MuseGNN is to iteratively minimize a sampling-based energy function during the forward pass, which allows the node embeddings to serve dual purposes: as predictive features for downstream tasks and as minimizers of the energy function. The authors present a scalable GNN that is able to deal with large-scale node classification benchmarks, including datasets exceeding 1TB in size.", "questions": "- How is the method compared to SGC and its variants in terms of efficiency and accuracy?\n- How does the method perform on heterophilic graphs? The energy-related loss seems to highly rely on the homophilic assumption.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper provides an analysis of the convergence properties of the proposed energy function and the iterative reduction process. \n- MuseGNN is designed to handle very large graphs, as evidenced by its performance on the node classification benchmark exceeding 1TB in size. \n- The experimental results show that MuseGNN achieves competitive accuracy compared to state-of-the-art GNN.", "weaknesses": "- The paper focuses on comparing MuseGNN with a few specific GNN architectures and frameworks. A more comprehensive comparison with a wider range of state-of-the-art GNNs. In particular, I want to see the comparison to SGC [1] and its variants in terms of efficiency and accuracy.\n- The training speed of the proposed method seems to be slower than GAT as shown in Table 1, while GAT is a well-known slow GNN.\n\n[1] Simplifying Graph Convolutional Networks", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "3C6zZk41xE", "reviewer": "Reviewer_pWkD", "summary": "This paper developed one subgraph sampling strategy induced from the energy regularization process over the GNN learning process. The author also discussed the convergence and expressive power of their proposed model under some conditions (i.e., $\\gamma$). The newly proposed model is advanced with a shared weight matrix over all subgraphs and by the consideration of replicated nodes in multiple subgraphs. Empirical studies show promising results via large-scale graph datasets.", "questions": "At the current stage, my questions are as follows: \n\n1. The organization of the paper needs to be improved. For example, figure 1 on page 2 is referred to in both the introduction section and the experiment section (ablation study), which is rarely to be observed in other studies. Also, the ablation result for $\\gamma$, located in Table 6, in Appendix B2, can be put into the current main page, which only contains 9.5 pages. \n\n2. In addition to this, I think the ablation on $\\gamma$ is very important as this shows how the degree of dependence between subgraphs affects the final training. Therefore, it is recommended that a more detailed analysis of this aspect be input. \n\n3. Another important empirical verification is the performance of the model via different sampling strategies since this directly affects the power of the term $\\|Y_s - \\mu_s\\|$. I found the related content in the Appendix, but it is recommended that the author put some convincing evidence to the main content.\n\n4. Furthermore, the energy regularizer $\\lambda \\mathrm{Tr}(\\mathbf Y_s \\mathbf L_s \\mathbf Y_x )$ controls the degree of smoothing via each subgraph, and the quantity of $\\lambda$ is shared over all the subgraph. However, in real practice, this may not be ideal as subgraphs shall have different degrees of smoothing.\n\n5. The expressive power of the GNN is usually discussed under the graph pooling tasks (e.g., graph level classification), it would be better and interesting to show some pooling results from the proposed model.\n\nMinor modification. \n\n1. In row 183, table 4 is missing, and it is allocated in the appendix.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The paper resolved the scalability problem on the GNNs induced by their energy optimization form. Empirical studies show the proposed model outperforms many baselines via different graph sizes up to 1TB. \n\n2. The paper provided theoretical analysis on the convergence of the proposed model under different values of $\\gamma$, which controls the degree of dependency between subgraphs.", "weaknesses": "1. Some important analysis and empirical study results are missing and shall be presented via the main text, please see the question part. \n\n2. The organization of the paper needs to be carefully adjusted.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "WFovzu3YrM", "reviewer": "Reviewer_ddwX", "summary": "This paper introduces a scalable graph neural network (GNN) architecture that minimizes an energy function through a sampling-based approach, allowing for interpretable node embeddings optimized for tasks like node classification.", "questions": "1. The ablation study of hyperparameter gamma shows that MuseGNN preserves high accuracy when gamma=0. In this case, the embedding of the same node in different subgraphs are not aligned. Can the node embeddings in subgraphs capture the global information of the original graph? Are there any theoretical insights about why MuseGNN still shows comparative performance even when gamma = 0?\n\n\n2. What is the sampling strategy when sampling subgraphs from the large original graph?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper discusses a critical issue of scalability when adding energy regularization to GCN, where the cost of energy calculation becomes enormous for large-scale graphs and previous methods fail to work on large graphs.\n\n2. The paper proposes a new method, MuseGNN, to enable energy-regularized GNN to efficiently handle large-scale graphs. The paper also comprehensively discusses the optimization procedure and convergence analysis of their method.\n\n\n3. The paper validates the performance and efficiency of MuseGNN by experiments on multiple-datasets, including large graph datasets where previous methods need huge computational cost.", "weaknesses": "1. The annotations and formulation in 2.1 may be confusing for new readers. The authors could adopt similar annotations in Descent Steps of a Relation-Aware Energy Produce Heterogeneous Graph Neural Networks by Ahn et al, which distinguishes the node embedding by basic model and the embedding by energy optimization by different annotations y and y*, and specify the procedure of UGCN to make the problem formulation clearer to a broader audience.\n2. The authors could compare the time and memory cost of MuseGNN with more baselines that also uses energy regularization without sampling to strengthen the superiority of MuseGNN.\n3. In the introduction, the authors could provide more specified discussions about the advantages of applying energy regularization to GNN to show the importance of this work.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)