Bundle Neural Network for message diffusion on graphs
Authors: Jacob Bamberger
Federico Barbero
Xiaowen Dong
Michael M. Bronstein
Venue: ICLR 2025 Spotlight (2025)
Keywords: graph neural network
sheaf neural network
geometric deep learning
algebraic topology
vector bundles
expressivity
Primary area: learning on graphs and other geometries & topologies
TL;DR: We propose Bundle Neural Networks, a new type of Graph Neural Network that operates via message diffusion, a continuous version of message-passing that allows to mitigate over-smoothing, over-squashing, and is provably expressive.

Abstract: The dominant paradigm for learning on graphs is message passing. Despite being a strong inductive bias, the local message passing mechanism faces challenges such as over-smoothing, over-squashing, and limited expressivity. To address these issues, we introduce Bundle Neural Networks (BuNNs), a novel graph neural network architecture that operates via *message diffusion* on *flat vector bundles* — geometrically inspired structures that assign to each node a vector space and an orthogonal map. A BuNN layer evolves node features through a diffusion-type partial differential equation, where its discrete form acts as a special case of the recently introduced Sheaf Neural Network (SNN), effectively alleviating over-smoothing. The continuous nature of message diffusion enables BuNNs to operate at larger scales, reducing over-squashing. We establish the universality of BuNNs in approximating feature transformations on infinite families of graphs with injective positional encodings, marking the first positive expressivity result of its kind. We support our claims with formal analysis and synthetic experiments. Empirically, BuNNs perform strongly on heterophilic and long-range tasks, which demonstrates their robustness on a diverse range of challenging real-world tasks.

Published as a conference paper at ICLR 2025
BUNDLE NEURAL NETWORKS FOR MESSAGE DIFFUSION ON GRAPHS
Jacob Bamberger Federico Barbero* Xiaowen Dong Michael Bronstein 1.2
University of Oxford PAITHYRA
ABSTRACT
The dominant paradigm for learning graphs is message passing. Despite being a strong inductive bias, the local message passing mechanism faces challenges such over-smoothing, over-squashing, and limited expressivity. To address these issues, we introduce Bundle Neural Networks (BuNNs), novel graph neu- ral network architecture that operates via message diffusion on flat vector bundles geometrically inspired structures that assign to each node a vector space and orthogonal map. A BuNN layer evolves node features through diffusion- type partial differential equation, where its discrete form acts as special case of the recently introduced Sheaf Neural Network (SNN), effectively alleviating over- smoothing. The continuous nature of message diffusion enables BuNNs to operate at larger scales, reducing over-squashing. We establish the universality of BuNNs in approximating feature transformations on infinite families of graphs with injec- tive positional encodings, marking the first positive uniform expressivity result its kind. We support our claims with formal analysis and synthetic experiments. Empirically, BuNNs perform strongly on heterophilic and long-range tasks, which demonstrates their robustness diverse range of challenging real-world tasks.
INTRODUCTION
Graph Neural Networks (GNNs) (Sperduti, 1993; Scarselli et al., 2009; Defferrard et al., 2016) are widely adopted machine learning models designed to operate over graph structures, with successes in diverse applications such as drug discovery (Stokes et al., 2020), traffic forecasting (Derrow- Pinion et al., 2021), and recommender systems (Fan al., 2019). Most GNNs are Message Passing Neural Networks (MPNNs) (Gilmer et al., 2017), where nodes exchange messages with immediate neighbours. While effective, MPNNs face critical challenges such over-smoothing (Li et al., 2018: Oono & Suzuki, 2020; Cai & Wang, 2020), over-squashing (Alon & Yahav, 2021; Topping al., 2022: Di Giovanni al., 2023), and limited expressivity (Xu et al.. 2019b; Morris al., 2019).
Over-smoothing occurs when node features become indistinguishable the depth of the MPNN in- creases, problem linked to the stable states of the heat equation on graphs (Cai & Wang, 2020). While Sheaf Neural Networks (Bodnar et al., 2022) address this by enriching the graph with a sheaf structure that assigns linear maps to edges and results in richer sta- ble states of the corresponding heat equation, they remain MPNNs and inherit other limitations such as over-squashing, which restricts the amount of infor- mation that can be transmitted between distant nodes.
Message Passing
Message Diffusion
A
Figure 1:
Local message passing graphs versus global message diffusion on bundles.
We propose Bundle Neural Networks (BuNNs), a new type of global GNN that operates over flat vector bundles structures analogous connections on flat Riemannian manifolds that augment the graph by assigning to each node a vector space and an orthogonal map. BuNNs do not perform 'explicit' message passing through multiple steps of information exchange between neighboring nodes, but instead operate via message diffusion. Each layer involves a node update step, and a
"jacoh. wawherger@@sss.oo.o.....

---

Published as a conference paper at ICLR 2025
diffusion step evolving the features according vector diffusion PDE in Singer & Wu (2012). The resulting architecture enjoys the desirable properties of Sheaf Neural Networks, in that they can avoid over- smothhng,,,, but are global models that can operate at larger scales of the graph to tackle over-squashing. and it achieves better performance on a range of benchmark datasets. Addi- tionally, we prove that equipped with injective positional encodings, BuNNs are compact uniform approximators, a new type of universality result for feature transformation approximation.
In summary, our contributions are the following:
We derive BuNNs from heat equations over flat vector bundles, and show that flat vector bundles are more amenable to computation than general vector bundles (Section 3).
We prove that the diffusion process can mitigate over-smoothing and squashing (Sec- tion 4), and support these claims with novel synthetic experiments (Section 6.1).
We prove that, with injective positional encodings. BuNNs are compact uniform universal approximators. To the best of our knowledge, this is the first of such results (Section 5).
We show that BuNNs perform well on heterophilic and long-range tasks, for instance, achieving a new state-of-the-ar result on the Peptides func dataset (Section 6.2).
2 BACKGROUND
Graphs. Let G (V, be an undirected graph on |V nodes with edges E. We represent the edges via an adjacency matrix A E Ra*= where the entry Aue for u, U E V is the edge (u,v) u) E and 0 otherwise. Let D be the diagonal degree matrix with entry D = d, equal to the degree of v. The graph Laplacian is defined as L == D - A and the random walk normalized graph Laplacian is defined as L = I D D-1D--1- A. We assume that at each node v V we are given c-dimensional signal (or node feature) x, E IR and group such signals into a matrix X E Raxe,
GNNs and feature transformations. A feature transformation a graph G is a permutation- cquivariant map fc Ruxes + Ruxcz that transforms the node signals. A GNNe is (continuous) map parameterized by O that takes as input a graph alongside node signals G (V, X. and outputs transformed signal (V, E, X'). A GNN on a graph G is therefore a feature transformation GNNę R=xe_ + R=xe2 Given a collection of graphs G. a feature transformation F on 9 is an assignment of every graph G E G to a feature transformation Fa R"C xm + I x The set of continuous feature transformations over a collection of graphs G denoted C (G, R Ro).
Cellular sheaves. A cellular sheaf (Curry, 2014) (F, G) over an undirected graph G (V, augments G by attaching to each node v and edge e a vector space space called stalks and denoted by F(v) and F(e), usually the stalks are copies of R for some d. Additionally, every incident node-edge pair v < e gets assigned a linear map between stalks called restriction maps and denoted Fese:: FFu) + F(c). Given two nodes and connected by edge (v. u), we can transport a vector X2 E F(v) from e to u by first mapping it to the stalk at e = (v. u) using Fuse. and mapping it to F(u) using the transpose Firee- As a generalization of the graph adjacency matrix, the sheaf adjacency matrix A, E [Rndxnd is defined as a block matrix in which each d x d block (Az)uv is Ful Fuge if there is an edge between u and v and Oaxa otherwise. Similary, we define the block diagonal degree matrix Dy Randxdd as (Dr) == d Laxe, and the sheaf Laplacian is Ly Dr Ay. These matrices act as bundle generalizations of their well-known standard graph counterparts and we recover such matrices when F(v) = R and F.g.v
for all e E V and E E.
Vector bundles. When restriction maps are orthogonal, we call the sheaf a vector bundle, a struc- ture analogous to connections on Riemannian manifolds. For this reason, the sheaf Laplacian also takes the name connection Laplacian (Singer & Wu, 2012). The product FT Fusge is then also or- thogonal and is denoted O... referring to the transformation a vector undergoes when moved across a manifold via parallel transport. In this case we denote the node-stalk at v by B(u), the bundle- adjacency by Ag and the bundle Laplacian Lg, and its normalized version Lg == Idaxd DD Ag-
Consider d-dimensional vector field over the graph, i.e. a d-dimensional feature vector at each node denoted X E Rnd in which the signals are stacked column- wise. Similarly to the graph case, the operation D D AggXi is averaging over the vector field, and Lg measure of smoothness, since:
(Dg B X)X) du u:(e,))EE
Our E ERd and (LxX), = d, ai(U,
Σ (x,, O OwxX,) E IRd
2

---

Published as a conference paper at ICLR 2025
Graph
General bundle
Flat bundle
Xu O,X,
desyne
syne
O,3,
(Lgx), = O (0 oxx,)
desyne (L sync (X))
(LX), Σ
(xo x.)
(Lx)),
Σ (x--
Omxx))
Figure 2: Comparison of different Laplacian and their actions signals.
BUNDLE NEURAL NETWORKS
In this section, we derive BuNN from heat diffusion equations over flat vector bundles. We then discuss its relationship GCN (Kipf & Welling, 2017) and NSD (Bodnar al., 2022). We provide algorithmic implementations and additional practical details in the Appendix (Section E).
Heat diffusion over bundles. The bundle Dirichlet energy &g (X) of vector field X Rnd is:
L8X Ex( xX
1 2
Σ (0,u)EE
-
- ||x d
020x.112
The gradient of the Dirichlet energy (X) is the random-walk Laplacian Ls. We write down a heal diffusion equation over a vector bundle as a gradient flow, whose evolution equation with initial condition X(0) X satisfies 0, (() written using matrix exponentiation as X(t) = exp(-t())X(0) (e.g. Hansen & Gebhart (2020)). We call the operator Ht((t) == exp(-tCa E Rnddknd the bundle heat kernel, which is defined as:
(1). The solution to this equation be
K
(-tCg)* k!
HB(t) = lim
k=
Computing the heat kernel is necessary to solve the heat equation. An exact solution can be com- puted using spectral methods. For small t, one can instead consider the truncated Taylor expansion centered at O, which amounts to fixing K in the above equation. Bodnar et al. (2022) instead approximates the solution using the Euler discretization of the heat equation with unit time step.
However, all these methods pose challenge when the bundle structure is learned in Bodnar al. (2022), since the heat kernel has to be recomputed after every gradient update the bundle structure. This high computational overhead limits the usability of Sheaf Neural Networks in applications.
Flat vector bundles. To address the scalability issues of general sheaves and vector bundles, we con- sider the special case of flat vector bundles in which every node gets assigned orthogonal map O and every connection factorizes as O uu = O O..
a. Consequently, the bundle Laplacian factors:
===0² O (
where E Rnn is block diagonal with v-th block being O.. We call the matrices O and O syn- chronization and desynchronization, respectively. We compare different Laplacians in Figure 2. This factorization avoids the O (d (d)E) cost of computing the restriction map over each edge. Addition- ally, Lemma 3.1 shows that it allows cast the bundle heat equation into standard graph heat equa- tion. This reduces the computation of the bundle heat kernel to that of the cheaper graph heat kernel, an operator that does not change depending the bundle and can, therefore, also be pre-computed.
Lemma 3.1. For every node the solution at time l of the heat equation connected bundle (V, E, O) with input node features X satisfies:
(Hb(t))), = Σ H(tt, v, T.
NEV
where H()) is the standard graph heat kernel, and H(1, u, u) R its the entry at (v, u).
3

---

Published as a conference paper at ICLR 2025
(1) Embed
(2) Update
(3) Diffuse
(4) Non-lin.
Input
Output
0 (G, P,
he O. WO.
Z (t(
X oZ)
Figure 3: Example of the message diffusion framework on graph with 4 nodes and 4 edges. From left to right: The imput is a simple graph embedding with each color representing the feature vector at that node. (1) An orthogonal map is computed for each node in the graph by embedding the nodes in a continuous manifold with local reference frames (represented as a torus for visual aid), the features represented
colored vectors do not change. (2) The features are updated using learnable parameters W. (3) The features are diffused for some time t according to the heat equation on the manifold: a larger value of leads to higher synchronization between all nodes as illustrated by the alignment of node features with respect to their local coordinates. (4) The output embedding is obtained by discarding the local coordinates and applying non-linearity.
The model. The BuNN layer occurs in four steps, as illustrated in Figure 3. First, the bundle maps O are computed using neural network Φ, the graph G. positional encodings P E Ty and the use of Householder reflections (Householder, 1958) or direct parameterization of the orthogonal axxd when d 2. Second, an encoder step updates the node signals via learnable matrix W E and bias b E Rd. Next, the features are diffused over the learned vector bundle using the heat kernel. Finally, non-linearity is applied. We summarize the steps in the following equations:
0 == (G, v) VU EV
h == O T w O 10 +b(0 WEE
(2)
( H)(()
(3)
x ((+1)) = (z(+
The diffusion time f in Equation is a hyperparameter determining the scale at which messages are diffused. For the case of small t, approximate the heat kernel via its truncated Taylor series of degree K, and for large 1, we use spectral methods. For simplicity of exposition, the steps above describe an update given single bundle (i.e., c d), meaning that x, Rdd In general, allow multiple bundles and vector field channels (Appendix E). Note that Equations 1, 2, and 3 are linear (or affine), and the non-linearities lie in 4. Equation 2 may be interpreted encoder, while Equation 3 is the message diffusion step guided by the heat kernel.
bundle-aware
Without the bundle structure, Equation 3 would converge exponentially fast to constant node repre- sentations over the graph (c.g. Theorem 1 in Li et al. (2018)), potentially leading to wer-smmothiin. This is a limitation of existing diffusion-based GNNs (Xu et al., 2019a; Zhao et al., 2021). Accord- ingly, the bundle is crucial in this formulation to prevent node features from collapsing.
Link Graph Convolutional Networks. It possible derive Graph Convolutional Networks (GCNs) (Kipf& & Welling, 2017) as an approximation of BuNNs operating over a trivial bundle. Setting t = 1 Equation becomes Z (1+1) = exp(- exp(-Lo)HH), The approximation exp(-Lc) P
I- LG gives the update (1+1) = (1 Lc) = (GH recovering the GCN update.
Comparison with Sheaf Neural Networks. Flat vector bundles are special case of cellular sheaves (Curry, 2014; Bodnar et al., 2022), meaning that our model has close connections to Sheaf Neural Networks (SNNs) (Hansen & Gebhart, 2020; Bodnar al., 2022: Barbero et al., 2022b:a). While most SNNs operate on fixed sheaves (Hansen & Gebhart, 2020: Barbero et al., 2022a; Bat- tiloro al., 2023), we focus on learning sheaves as in Neural Sheaf Diffusion (NSD) from Bodnar et al. (2022). BuNNs distinguish themselves from NSD in several ways. First, NSD approximates the heat equation using a time-discretized solution to the heat equation, which results in a standard message passing algorithm. In contrast, the direct use of the heat kernel allows BuNNs to break away from the explicit message-passing paradigm. This is in line with Battiloro et al. (2024) who define convolutions manifolds using the heat kernel of fixed sheaves. Secondly, the of flat bundles increases scalability since the bundle maps are computed at the node level. Additionally, flat bundles guarantees path independence, requirement for the theory on the long time limit of NSD)
4

---

Published as a conference paper at ICLR 2025
in Bodnar al. (2022) to hold, often not satisfied for general sheaf constructions such ones used in NSD. Thirdly, we allow ó to be any GNN while NSD restricts & to be an MLP. We found that in- corporating the graph structure in é improved the experimental results. The update in Equation 2 is also different to NSD how the W and b are applied, and is necessary to prove our main theoretical result, Theorem 5.3. Additionally, Bodnar al. (2022) experiments with general restriction maps as opposed to restricting them to orthogonal maps, and find that in practice orthogonal restrictions per- form better, hence we consider orthogonal restriction maps. We provide experimental comparisons to NSDs in the experimental section and show that BuNNs significantly outperform their sheaf counterparts. We provide summarized comparison between GCNs, SNNs, and BuNNs in Table 1.
Comparison with other methods. Another paradigm for learning on graphs is using graph trans- formers (GT), where every nodes communicate to each other through the use of self-attention. When the graph transformer is fully connected, all nodes communicate and therefore GT should not suffer from under-reaching or over-squashing. but might suffer from over-smoothing (Dovonon et al., 2024). When not fully connected, they might suffer from over-squashing (Barbero et al., 2024). Other paradigms such as Implicit Graph Neural Networks Fu et al. (2023) are also designed to tackle under-reaching and long- gangrange dependencies while empirically not suffering from over-smoothing.
4 PROPERTIES OF BUNDLE NEURAL NETWORKS
We now give a formal analysis of the BuNN model. In Section 4.1, we derive the fixed points of the bundle heat diffusion, which the subspace of signals towards which solutions converges, and show that even in the limiting case BuNNs can retain information at the node level and therefore avoid over-smoothing. Section 4.2 discusses how our model capture long-range interactions and mitigate over-squashing.
Table I:
Comparison between models in terms of message type and capability of mitigating issues related GNNs.
GCN
SNN
BuNN
Propagation
Ac
C
Ht(()
Standard (Local)
Standard (Local)
Diffusive (Global)
Message type
No under-reaching
x
x
Alleviates
X
X
over-squashing
4.1 FIXED
Alleviates
POINTS AND NNER-SMOOTHING.
over-smoothing
major limitation of MPNNs is over-smoothing, where node features become indistinguishable as MPNN depth increases. This phenomenon is a major challenge for training deep GNNs. arises because the diffu- sion in MPNNs resembles heat diffusion graphs (Di Giovanni et al.,, which converges to uninformative fixed points', leading to a loss of information at the node level. Bodnar et al. (2022) show that the richer bundle structure, however, gives rise to richer limiting behavior. Indeed, by Lemma 3.1, since lim,-+20 H(t, v, u) Z the limit over time of solution is EE Euev d., O O.xy. To understand the space of fixed points, notice that any Y E Raxd ex- pressible as y. = Σ Evev de O,,,, for some X is a fixed point, and for any two nodes u, v:
1
O.y, 2
Σ de O = O.yu. WEV
Consequently, the fixed points of vector diffusion have a global geometric dependency where all nodes relate to each other by some orthogonal transformation, e.g. the output in Figure 3. Equation 5 provides insight into the smoothing properties of BuNNs. When the bundle is trivial, the equation reduces to y. yu, with no information at the node level, i.e. over-smoothing, When it is non-trivial, the output signal can vary across the graph (e.g., Ye * yu for two nodes u and v). We summarize this in Proposition 4.1 and prove an implication in terms of over-smoothing in Proposition 4.2. Similar results showing that the signal can survive in deep layers as opposed to resulting in the constant signal have previously been done empirically and theoretically for other architectures, e.g. Chamberlain et al. (2021); Fu et al. (2023).
Up to degree scaling if using the symmetric-normallize Laplacian, see Li al. (2018).
5

---

Published as a conference paper at ICLR 2025
Proposition 4.1. Let Y be the output of a BuNN layer with t oc, where G is connected graph, and the bundle maps are not all equal. Then, there exists u, v E V connected such that y. # Yu- Proposition 4.2. Let X be 2 2-dimensional features on a connected and assume that there are two nodes with Pu. + Pv- Then there is an 5 >> O such that for every K, there exist K-layer deep BuNN with dimensional stalks, = oc, ReLU activation, df is a node-level MLP, and |[ S 1, with output Y such that | Yull 8.
4.2 WVER-SQUASHING AND LONG RANGE INTERACTIONS.
While message-passing in MPNNs constitutes a strong inductive bias, it is problematic when the task requires the MPNN capture interactions between distant nodes. These issues have been attributed mainly to the over-squashing problem. Topping et al. (2022) and Di Giovanni et al. (2023) formalize over-squashing through a sensitivity analysis, giving upper bounds on the sensitivity of the output at node with respect to the input at an other. In particular, they show that under weak assumptions on the message function, the Jacobian of an MPNN satisfies the following inequality
ja (X) /0x.1 S d' (A')
(6)
for any nodes u, v, where is the depth of the network and is constant. Given two nodes u, at a distance r, message-passing will require at least r layers for the two nodes to communicate and overcome under-reaching. If é is large, distant nodes can communicate, but Di Giovanni et al. (2023) show that over-squashing becomes dominated by vanishing gradients. Building on top of such sensitivity analysis, we compute the Jacobian for BuNN layer in Lemma 4.3.
Lemma 4.3. Let BuNN be a linear layer defined by Equations 1, 2 & 3 with hyperparameter t. Then, for any connected graph and nodes u, U. we have
① (BuNN (X)) H(tt, u, ")O
The form of the Jacobian in Lemma 4.3 differs significantly from the usual form in Equation First, all nodes communicate in single BuNN layer since H(t, u, v) > for all u, v, and allowing for direct pair-wise commumication between nodes, making a BuNN layer operate globally similarly to Transformer model, and therefore overcome under-reaching. Secondly, taking & be large allows BuNNs to operate on a larger scale, allowing stronger communication between distant nodes and overcoming over-squashing without the vanishing gradient problem.
Further, Lemma 4.3 gives finer picture of the capabilities of BuNNs. For example, to mitigate over- squashing a node may decide to ignore information received from certain nodes while keeping in- formation received from others. This allows the model to reduce the receptive field of certain nodes:
Corollary 4.4. Consider n nodes u, v, and W4, for i = 1,...n 2, of a connected graph with 2 dimensional bundle such that Pe # Pw, and Pu + Pw, Vi. Then in a BuNN layer with MLP ó, at given channel, the node can learn to ignore the information from all w;s while keeping information from
EXPRESSIVITY OF THE MODEL
We now characterize the expressive power of BuNNs from a feature transformation perspective. This analysis extends that of Bodnar et al. (2022), who show that heat diffusion on sheaves can be expres- sive enough linearly separate nodes in the infinite time limit. Instead, our results concern the more challenging problem of parameterizing arbitrary feature transformations and hold for finite time.
Most work on GNN expressivity characterize the ability of GNNs to distinguish isomorphism classes of graphs nodes (Xu et al., 2019b; Morris et al., 2019; Azizian & Lelarge, 2021; Geerts & Reutter, 2022) or equivalently to approximate functions on them (Chen et al., 2019). These results typically rely on two major assumptions: !) node features are fixed for each graph or are ignored completely: and 2) apply to a single graph or finitely many graphs of bounded size. Instead, our setting 1) includes features ranging in an infinite uncountable domain, and 2) includes infinite families of graphs. To this end, we define the notion of compact uniform approximation as a modification to that of uniform approximation from Rosenbluth al. (2023) and we discuss our choice in Appendix C.

---

Published as a conference paper at ICLR 2025
Definition 5.1. Let F C C,, Re') be set of feature transformations over family of graphs 9, and let H E C(G, IRe Re) feature transformation over G. We say that F compactly uniformly approximates H, if for all finite subsets K C G. for all compact K C R. and for all E > 0, there exists an F E F such that for all G K. and X E K we have that (X/'(X) Hx(X) S e.
Injective Positional Encodings (PEs). In the case of finite collection of graphs G and fixed finite feature space K, the above definition reduces to the setting of Theorem 2 in Morris et al. (2019), since the node features are fixed to the finitely many values in K. However, when K is not finite, the arguments in Morris et al. (2019) no longer work, as detailed in Appendix D. Consequently, Definition 5.1 subsumes graph-isomorphism testing, and it is, therefore, too strong for volynomial-time GNN to satisfy. Hence, we will assume that the GNN has access to injective positional encodings, which we formally define in Definition D.I. This allows us to bypass the graph-isomorphist problem, while not trivializing the problem as we show next.
Negative results with injective PE. Characterizing the expressive power of GNNs in the uniform setting is an active area of research, with mostly negative results. Rosenbluth et al. (2024) prove neg- ative results in the con-compact setting for both MPNNs with node and graph transformers, even with injective positional encodings. While extending their result to the compact setting is out- side the scope of our work, we prove Proposition 5.2 negative result for bounded-depth MPNNs with injective PEs. Indeed, fixing the depth of the MPNNs to f, can take single compactly fea- tured graph G E G with a diameter larger than l. As there is a node whose receptive field does not include all nodes in such G, the architecture cannot uniformly approximate every function on G.
Proposition 5.2. There exists a family 9 consisting of connected graphs such that bounded- depth MPNNs are not compact uniform approximators, even if enriched with unique positional encoding.
Universality of BuNN. In contrast to the negative results above, we now show that when equipped with injective PEs, BuNNs are universal with respect to compact uniform approximation. To the best of our knowledge, Theorem 5.3 is the first such positive feature approximation result for a GNN architecture, which demonstrates the remarkable modelling capabilities of BuNNs.
Theorem 5.3. Let G be a possibly infinite set of connected graphs equipped with injective positional encodings. Then 2-layer deep BuNNs with encoder/decoder at each layer and Φ ((2) being 2- layer deep MLP have compact uniform upproximation over G.
In particular, let e > 0 and H be a feature transformation on a finite subset K. S G and K S Rdd a compact set, then there 2-layer deep BuNN with width of order O (EGER (VGl) that approximates H over LaEE km C LIgeg Rued_ In other words, the required hidden dimension of BuNN is only linearly dependent the number of nodes in the family of graphs.
EXPERIMENTS
In this section, we evaluate BuNNs through a range of synthetic and real-world experiments. We first validate the theory from Section 4 with two synthetic tasks. We then evaluate BuNNs on pop- ular real-world benchmarks. We use truncated Taylor approximations for small values of t, while for larger
the truncated spectral solution. We provide supplementary information the implementation and precise experimental details in the Appendix (Sections E and F respectively)."
6.1 SYNTHETIC EXPERIMENTS: SVER-SQUASHING AND TWER-SMOOTHIE
Tasks. In this experiment, we propose two new node-regression tasks in which nodes must average the input features of subset of nodes. The input graph contains two types of nodes, whose features are sampled from disjoint distributions. The target for nodes is to output the average input feature over nodes of the other type and vice-versa, as illustrated in Figure to mitigate over-squashing. We consider the case where the underlying graph is a barbell graph consisting of two fully connected graphs each being type and bridged by single edge. This bottleneck makes it hard to transfer information from one cluster to another. Second, we test the capacity to mitigate over-smoothing. We consider the fully connected graph, in which all nodes are connected. The fully connected graph is a worst-case scenario for over-smoothing since after one step of message passing, the features are fully averaged over the graph and hence over-smoothed.
First, we test the capacity
"This phenomenon in MPNNs is often called 'under-reaching'.
'lll code can he found at https: / //tp///tttttttt/ ccm http:////.accccaaaaeeereerr/eeee

---

Published as a conference paper at ICLR 2025
Setup. As first baseline, consider a constant predictor always predicting 0, the expected mean over the whole graph. As a second baseline, the cluster-specific constant predictor predicting the expected mean over the opposite cluster, that is, Ły3 depending the cluster.
2
Additionally, we consider GNN base- lines to be node-level MLP, GCN (Kipf & Welling, 2017). Graph- SAGE (Hamilton et al., 2017), GAT (Veličković et al., 2018), NSD (Bodnar al., 2022), and fully connected GraphGPS (Rampášek et al. (2022)). The depth of MPNNs is fixed to the minimal depth to avoid under-reaching, namely 3 for the barbell and 1 for the fully connected graph, and ensure the width is large (> 128) considering the task. The depth of the transformer is tuned between 1 and with heads. We compare these to BuNN with MLP learning the bundle maps a comparable number of parameters and the same depth. We use Adam optimizer with 10-3 learning rate, batch size 1, and train for 500 epochs. We use 100 samples for both training and testing.
u (-v3N, 0) x 2 102 v3N
y.
N Σ
N
Figure 4: smoothing (right). In both cases, blue nodes output the average over the red nodes and vice-versa.
Synthetic over-squashing (left) and over-
The results for N 10 reported in Table 2. All MPNNs perform poorly the over-squashing task. All MPNNs perform comparably to baseline 2, showing their incapa- bility to transfer information between clusters. This is explained by the fact that nodes from different clusters have high commute time and effective resistance, which tightly connects to over-squashing Di Giovanni et al. (2023); Black et al. (2023); Dong al. (2024). On the other hand, BuNN achieves almost perfect accuracy on this task, which supports the claim that BuNN mitigates squashin.. We note, however, that to solve the task perfectly, we need t 2 10, allowing BuNN to operate on larger scale more adapted to the task. To solve this task, BuNN can assign the orthogonal maps to separate the two types of nodes, making each node listen only to the nodes of the other type, a behavior proved to be possible in Corollary 4.4.
Results.
Similarly, the over-smoothing task on the clique graph is also challenging. Indeed, GCN and GAT perform exceptionally poorly, com- parably to Baseline 1 which only has access to global information. Indeed, to the best of our knowledge, these are the only models with formal proofs of over-smoothing (Cai & Wang, 2020; Wu et al., 2023). GraphSAGE performs slightly better because processes neighbors differently than processes nodes themselves. Interestingly, GPS fails to gener- alize providing evidence for over-smoothing and over-squashing in graph transformers. Moreover, NSD and BuNN solve the task due to their capability to mitigate over-smoothing. Indeed, BuNN can learn to ignore nodes from given cluster as proved in Corollary 4.4.
Table 2: BuNN mitigates over-smoothing and over-squashing. Mean squared error (MSE) of different models on the two synthetic tasks.
Barbell
Clique
(over-squashing)
(over-smoothingg)
Base. ]
30.97 + 0.42
30.94 + 0.42
Base. 2
1.00 + 0.07
0.99 + 0.08
MLP
1.08 = 0.07
1.10 0.08
GCN
1.05 = 0.08
29.65 155::0.0 0.34
SAGE
0.90 0.29
0.100
GAT
+
20.97
GPS
1.06 = 0.13
1.06
NSD
1.09 + 0.15
0.08 + 0.02
BuNN
0.01 + 0.07
0.03 + 0.01
6.2 REAL-WORLD TASKS
We evaluate BuNNs the Long Range Graph Benchmark (Dwivedi et al., 2022) and the het- erophilic tasks from Platonov et al. (2023). We provide the implementation details Appendix E.
Heterophilic datasets. As have shown in Section 4, BuNNs provably capable of avoiding over-smoothing. It is, therefore, natural to test how BuNN performs on heterophilic graphs where over-smoothing is recognized important limitation (e.g. Yan al. (2022)). We follow their
8

---

Published as a conference paper at ICLR 2025
Table 3: Results for the heterophilic tasks. Accuracy is reported for roman-empire and amazon-raa ratings, and ROC AUC is reported for minesweeper, tolokers, and quest ions. Best results are denoted by bold. Asterisk" denotes that some runs ran out of memory on an NVIDIA A10 GPU (24 GB).
roman-empire
amazon-rattigg
minesweeper
tolokers
questions
GCN
73.69 + 0.74
48.70 0.63
89.75 + 0.52
83.64 0.67
76.09 = 1.27
SAGE
85.74 = 0.67
53.63 = 0.39
93.51 + 0.57
82.43 = 0.44
76.44 0.62
GAT
80.87 = 0.30
49.09 = 0.63
92.01 0.68
33.00
77.43 = 1.20
GAT-sep
88.75 0.41
52.70 0.62
93.91 = 0.35
83.78 0.43
76.79 = 0.71
GT
86.51 + 0.73
51.17=0.. 0.66
91.85 + 0.76
83.23 0.64
77.95 + 0.68
GT-sep
87.32 + 0.39
52.18 + 0.80
92.29 0.47
82.52 0.92
78.05 + 0.93
NSD
80.41 + 0.72
42.76 + 0.54
+
78.83 + 0.76'
69.69 + 1.46*
BuNN
+ 11.5500.33
53574
98.99 0.16
84.78 0.80
78.75 + 1.09
methodology to evaluate BuNN on the 5 heterophilic tasks proposed in Platonov et al. (2023). We run the models with 10 different seeds and report the mean and standard deviation of the test ac- curacy for roman-empe re and amazon-rati ngs, and mean and standard deviation test ROC AUC for minesweeper, tolokers, and questions We use the classical baselines from Platonov et al. (2023) and NSD, and provide the syper-parameter in the Appendix (Section F.2).
Results. We report the results Table BuNN achieves the best score on all tasks, with an average relative improvement of 4.4%. Its score minesweeper is particularly impressive, which is sig- nificantly ahead of the rest and for which BuNN solves the task perfectly. We found that the optimal value of over grid search varies across datasets, being 1 for amazon- -att ings and 100 for roman-emp: re. BuNN consistently outperforms the sheaf-based model NSD by a large margin, which we believe is due to the fact that NSD learns map for every node-edge pairs making NSD) more prone to overfitting, while BuNNs can act as a stronger regularizer. Such strong performance confirms our theory on over-smoothing from Section 4.1 and showcase the strong modeling capacity of BuNN in heterophilic settings.
Table 4: Results for the Pept sentiidesssuucte
and Pascal VOC-SP tasks
from the Long Range Graph Benchmark (results are 100 for clarity). The best result is bold.
Pept deiiidediiii des-funn Test AP
Pept septides-strrce
Pascal VoccsP Test F11
Model
Test MAE +
GCN
68.60 + 0.50
24.60 0.07
20.78 0.31
GINE
66.21 = 0.67
24.73 + 0.17
27.18 0.54
GatedGCN
67.65 = 0.47
24.77 0.09
38.80 = 0.40
DReW
71.50 + 0.44
25.36 0.15
33.14 0.24
SAN
64.39 + 0.75
25.45
32.30 0.39
GPS
65.34 + 0.91
25.09 0.14
44.40 + 0.65
GAPH VIT
69.42 0.75
24.49 0.16
Exphormer
65.27 + 0.43
24.81 + 0.07
39.75 + 0.37
BuNN
72.76 + 0.65
24.63 + 0.12
40.49 + 0.46
Long Range Graph Benchmark. In Section 4.2, we showed that BuNNs have desirable properties when it comes over-squahhing and modeling long-range interactions. To verify such claims em- pirically, we evaluate BuNN on tasks from the Long Range Graph Benchmark (LRGB) (Dwivedi al., 2022). We consider the Pept i des dataset consisting of 15 535 graphs which come with two associated graph-level tasks, Pept i des- and Peptides -es-structtsttsse strutt, and the node classifi- cation on the PascalVOs woc VOC-SP dataset with 11 355 graphs and 5.4 million nodes where each graph corresponds to image in Pascal VOC 2011 and each node to superpixel in that image.
The graph classification task in Peptides- 10 classes, while the regression task in peptides. In both cases, we follow the standard experimental setup detailed by Dwivedi et al. (2022) alongside the updated suggestions from Tönshoff et al. (2023). The performance metric is Average
is to predict the function of the peptide from is inferring the 3D properties of the

---

Published as a conference paper at ICLR 2025
Precision (AP) for The node classification task in PascalVOC-SP consists of predicting a semantic segmentation label for each node out of 21 classes, with performance metric being the macro Fl score. We run each experiment on distinct seeds and report mean and standard deviation over these runs. Baseline models are taken from Tönshoff al. (2023) and include MPNNs, transformer models, and the current SOTA models (Gutteridge et al., 2023; Hc et al., 2023; Rampášek et al., 2022).
and Mean Absolute Error (MAE) for Peptides-struct.
Results. We report the results in Table 4. BuNNs achieve, the best of our knowledge, new state-of-the-art result on Peptides func. BuNNs also perform strongly on both other tasks, being second overall on PascalVOC-SF and clearly outperforming all MPNN models, and third within one standard deviation of the second model on Pept sespridds-stuss suctt.. The overall strong performance on the LRGB benchmarks confirms our theory on oversquashing from Section 4.2 and provides further evidence of the long-range capabilities of BuNNs.
CONCLUSION
In this work, we proposed Bundle Neural Networks new type of GNN that operates via message diffusion on graphs. We gave formal analysis of BuNNs showing that message diffusion can mitigate issues such as over-smoothing since the heat equation over vector bundles admits richer set of fixed points and over-squashing since BuNNs can operate MPNNs. We also prove compact uniform approximation of BuNNs, a first expressivity result of its kind, characterizing their expressive power and establishing their superiority over MPNNs. We then confirmed our theory with carefully designed synthetic experiments. Finally, we showed that BuNNs perform well heterophilic and long range tasks which are known be challenging for MPNNs.
larger scale than standard
ACKNOWLEDGEMENTS
This research is partially supported by EPSRC Turing Al World-Leading Research Fellowship No. EP//040062/1 and EPSRC AI Hub on Mathematical Foundations of Intelligence: An "Erlangen Programme" for AI No. EP/Y028872/1. X.D. acknowledges support from the Oxford-Man Institute of Quantitative Finance and the EPSRC (EP/T023333/I).
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Uri Alon and Eran Yahav, On the bottleneck of graph neural networks and its practical implications. In International Conference on Learning Representations, 2021.
Waiss Azizian and Marc Lelarge. Expressive power of invariant and equivariant graph neural net- works. In International Conference on Learning Representations, 2021.
Federico Barbero, Cristian Bodnar, Haitz Sáez de Ocáriz Borde, Michael Bronstein, Petar Veličković, and Pietro Liò. Sheaf Neural Networks with Connection Laplacians. In Proceed- ings of Topological, Algebraic, and Geometric Learning Workshops 2022, pp. 28-36. PMLR, November 2022a.
Federico Barbero, Cristian Bodnar, Haitz Sáez de Ocáriz Borde, and Pietro Lio'. Sheaf Attention Networks. In NeurlPS 2022 Workshop on Symmetry and Geometry in Neural Representations, 2022b.
Federico Barbero, Andrea Banino, Steven Kapturowski, Dharshan Kumaran, João GM Araújo, Alex Vitvitskyi, Razvan Pascanu, and Petar Veličković. Transformers need glasses! information over- squashing in language tasks. arXiv preprint arXiv:2406. 042677 2024.
C. Battiloro, Z. Wang, H. Riess, P. Di Lorenzo, and A. Ribeiro. Tangent bundle filters and neural net- works: From manifolds to cellular sheaves and back. In ICASSP 2023 2023 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 1-5, 2023.
Claudio Battiloro, Zhiyang Wang, Hans Riess, Paolo Di Lorenzo, and Alejandro Ribeiro. Tangent bundle convolutional learning: from manifolds to cellular sheaves and back. IEEE Transactions on Signal Processing, 2024.
Mitchell Black, Zhengchao Wan, Amir Nayyeri, and Yusu Wang. Understanding Oversquashing in GNNs through the Lens of Effective Resistance, June 2023. arXiv:2302.06835 [es].
Cristian Bodnar, Francesco Di Giovanni, Benjamin Paul Chamberlain, Pietro Lio, and Michael M. Bronstein. Neural Sheaf Diffusion: A Topological Perspective on Heterophily and Oversmoothing in GNNs. Advances in Neural Information Processing Systems, 2022.
Chen Cai and Yusu Wang. A Note on Ovor-Smoothing for Graph Neural Networks. ICML Graph Representation Learning workshop. arXiv, 2020.
Benjamin Paul Chamberlain, James Rowbottom, Davide Eynard, Francesco Di Giovanni, Xiaowen Dong, and Michael M. Bronstein. Beltrami Flow and Neural Diffusion on Graphs, October 2021. arXiv:2110.09443 Jes, stat].
Zhengdao Chen, Soledad Villar, Lei Chen, and Joan Bruna. On the equivalence between graph isomorphism testing and function approximation with GNNs. In Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
Justin Michael Curry. Sheaves, cosheaves and applications. University of Pennsylvania, 2014.
G. Cybenko. Approximation by superpositions of a sigmoidal function. Mathematics of Control, Signals and Systems, 2(4):303-314, December 1989.
Michaël Defferrard, Xavier Bresson, and Pierre Vandergheynst. Convolutional Neural Networks Graphs with Fast Localized Spectral Filtering. In Advances in Neural Information Processing Systems, volume 29. Curran Associates, Inc., 2016.
Austin Derrow-Pinion Jennifer She, David Wong, Oliver Lange, Todd Hester, Luis Perez, Mare Nunkesser, Seongjae Lee, Xueying Guo, Brett Wiltshire, et al. Eta prediction with graph neural networks google maps. Proceedings of the 30th ACM international conference on informa- fion & knowledge management, pp. 3767-3776, 2021.
Francesco Di Giovanni, James Rowbottom, Benjamin Paul Chamberlain, Thomas Markovich, and Michael Bronstein. Understanding convolution on graphs via energies. Transactions on Ma- chine Learning Research.
11

---

Published as a conference paper at ICLR 2025
Francesco Di Giovanni, Lorenzo Giusti, Federico Barbero, Giulia Luise, Pietro Lio, and Michael M Bronstein. On over-squashing in message passing neural networks: The impact of width, depth, and topology. In International Conference Machine Learning, pp. 1865-7885, PMLR, 2023.
Yanfei Dong, Mohammed Haroon Dupty, Lambert Deng, Zhuanghua Liu, Yong Liang Goh, and Wee Sun Lee. Differentiable cluster graph neural network. arXiv preprint arXiv:2405.16188 16185, 2024.
Gbetondji IS Dovonon, Michael M Bronstein, and Matt J Kusner. Setting the record straight on transformer oversmoothing. arXiv preprint urXiv:2401.0430) 2024.
Vijay Prakash Dwivedi, Ladislav Rampášek, Mikhail Galkin, Ali Parviz, Guy Wolf, Anh Tuan Luu, and Dominique Beaini. Long range graph benchmark. In Thirty-sixth Conference on Neural Information Processing Systems Datasets and Benchmarks Track, 2022.
Wenqi Fan, Yao Ma, Qing Li, Yuan He, Eric Zhao, Jiliang Tang, and Dawei Yin. Graph neural networks for social recommendation. In The world wide web conference, pp. 417-426, 2019.
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. 033066 2023.
Floris Geerts and Juan Reutter. Expressiveness and approximation properties of graph neural networks. In International Conference on Learning Representations, 2022.
Justin Gilmer, Samuel S. Schoenholz, Patrick Riley, Oriol Vinyals, and George E. Dahl. Neural Message Passing for Quantum Chemistry. In Proceedings of the 34th International Conference on Machine Learning, pp. 1263-1272. PMLR, July 2017.
Benjamin Gutteridge, Xiaowen Dong, Michael Bronstein, and Francesco Di Giovanni. Drew: dy- namically rewired message passing with delay. In Proceedings of the 40th International Confer- ence on Machine Learning, ICML 23. JMLR.org, 2023.
William L. Hamilton, Rex Ying, and Jure Leskovee. Inductive representation learning on large graphs. In Proceedings of the 31st International Conference on Neural Information Processing Systems, NIPS' 17, pp. 1025-1035, Red Hook, NY, USA, 2017. Curran Associates Inc.
Jakob Hansen and Thomas Gebhart. Sheaf Neural Networks. In NeurIPS Workshop TDA and Beyond, 2020.
Xiaoxin He, Bryan Hooi, Thomas Laurent, Adam Perold, Yann LeCun, and Xavier Bresson. A generalization of vit/mlp-mixer to graphs. In Proceedings of the 40th International Conference on Machine Learning, ICML'23. JMLR.org. 2023.
Alston Householder. Unitary triangularization of a nonsymmetric matrix. Journal of the ACM (JACM), 1958.
Thomas N. Kipf and Max Welling, Semi-Supervised Classification with Graph Convolutional Net- works. In International Conference on Learning Representations, 2017.
Qimai Li, Zhichao Han, and Xiao-Ming Wu. Deeper Insights into Graph Convolutional Networks for Semi-Supervised Learning. In Proceedings of the Thirty-Second AAAI Conference Ar- tificial Intelligence (AAAI-18), pp. 3538-3545. Association for the Advancement of Artificial Intelligence, February 2018. ISBN 978 978-1-57735-800-
Christopher Morris, Martin Ritzert, Matthias Fey, William L. Hamilton, Jan Eric Lenssen, Gaurav Rattan, and Martin Grohe. Weisfeiler and Leman Go Neural: Higher-Order Graph Neural Net- works. Proceedings of the AAAI Conference on Artificial Intelligence, 33(01):4602-4600 July 2019.
Anton Obukhov. Efficient Householder transformation in PyTorch, 2021.
URL https://
12

---

Published as a conference paper at ICLR 2025
Kenta Oono and Taiji Suzuki. Graph Neural Networks Exponentially Lose Expressive Power for Node Classification. In International Conference on Learning Representations, 2020.
Oleg Platonov, Denis Kuznedelev, Michael Diskin, Artem Babenko, and Liudmila Prokhorenkova. A critical look the evaluation of GNNs under heterophily: Are really making progress? In The Eleventh International Conference on Learning Representations, 2023.
Ladislav Rampášek, Michael Galkin, Vijay Prakash Dwivedi, Anh Tuan Luu, Guy Wolf, and Do- minique Beaini. Recipe for general, powerful, scalable graph transformer. Advances in Neural Information Processing Systems, 15:14501-14515, 2022.
Eran Rosenbluth, Jan Tönshoff, and Martin Grohe. Some Might Say All You Need Is Sum. In Proceedings of the Thirty-Second International Joint Conference on Artificial Intelligence, pp- 4172-4179, Macau, SAR China, August 2023. International Joint Conferences on Artificial In- telligence Organization.
Eran Rosenbluth, Jan Tönshoff, Martin Ritzert, Berke Kisin, and Martin Grohe Distinguished in uniform: Self-attention vs. virtual nodes. In The Twelfth International Conference on Learning Representations, 2024.
Franco Scarselli, Marco Gori, Ah Tsoi, Markus Hagenbuchner, and Gabriele Monfardini. The Graph Neural Network Model. IEEE transactions on neural networks / a publication of the IEEE Neural Networks Council, 20:61-80, January 2009.
A. Singer and H.-T. Wu. Vector diffusion maps and the connection Laplacian. Communications on Pure and Applied Mathematics, 65(8):1067-1 144, 2012.
Alessandro Sperduti. Encoding Labeled Graphs by Labeling RAAM. In Advances in Neural Infor- mation Processing Systems, volume Morgan-Kauffanne 1993.
Jonathan M. Stokes, Kevin Yang, Kyle Swanson, Wengong Jin, Andres Cubillos-Ruiz, Nina M. Donghia, Craig R. MacNair, Shawn French, Lindsey A. Carfrae, Zohar Bloom-Ackermann, Vic- toria M. Tran, Anush Chiappino-Pepe, Ahmed H. Badran, Ian W. Andrews, Emma J. Chory, George M. Church, Eric D. Brown, Tommi S. Jaakkola, Regina Barzilay, and James J. Collins. A Deep Learning Approach to Antibiotic Discovery. Cell, 180(4):688-700) -00.8-337111, February 2020. ISSN 0092-8674. doi: 10. 10.1016/j.ccll 2020.00..0.1..01....1
Jan Tönshoff, Martin Ritzert, Eran Rosenbluth, and Martin Grohe Where did the gap go? reassess- ing the long-range graph benchmark. In The Second Learning on Graphs Conference, 2023.
Jake Topping. Francesco Di Giovanni, Benjamin Paul Chamberlain, Xiaowen Dong, and Michael M. Bronstein. Understanding over-squashing and bottlenecks on graphs via curvature. In Interna- fional Conference on Learning Representations, 2022.
Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Liò, and Yoshua Bengio. Graph attention networks. In International Conference Learning Representations, 2018.
Xinyi Wu, Amir Ajorlou, Zihui Wu, and Ali Jadbabaie. Demystifying Oversmoothing in Attention- Based Graph Neural Networks. Advances in Neural Information Processing Systems, 36:35084- 35106, December 2023.
Bingbing Xu, Huawei Shen, Qi Cao, Keting Cen, and Xueqi Cheng. Graph Convolutional Networks using Heat Kernel for Semi- Semi-supervised Learning. In Proceedings of the Twenty-Eighth Interna- tional Joint Conference on Artificial Intelligence, pp. 1928-1934, Macao, China, August 2019a. International Joint Conferences Artificial Intelligence Organization.
Keyulu Xu, Weihua Hu, Jure Leskovec, and Stefanic Jegelka. How Powerful are Graph Neural Networks? In International Conference on Learning Representations, 2019b.
Y. Yan, M. Hashemi, K. Swersky, Y. Yang, and D. Koutra. Two sides of the same coin: Het- erophily and oversmoothing in graph convolutional neural networks. In 2022 IEEE International Conference on Data Mining (ICDM), pp. 1287-1292, Los Alamitos, CA, USA, dec 2022. IEEE Computer Society.
13

---

Published as a conference paper at ICLR 2025
Jialin Zhao, Yuxiao Dong, Ming Ding, Evgeny Kharlamov, and Jie Tang. Adaptive Diffusion in Graph Neural Networks. In Advances in Neural Information Processing Systems, 2021.
14

---

Published as a conference paper at ICLR 2025
A LIMITATION AND FUTURE WORK
A limitation of our framework is that while message diffusion allows to operate on different scales of the graph, the computation of the heat kernel for large requires spectral methods and is therefore computationally expensive. An exciting research direction consists of using existing computational methods approximate the heat kernel efficiently for large values of A limitation of our exper- iments is that we consider inductive graph regression/clasifcctiinn and transductive node regres- sion/classificttion tasks but no link prediction task. limitation our theory is that Theorem 5.3 assumes injective positional encodings at the node level, which might only sometimes be available; future work could characterize the expressiveness when these are unavailable.
B PROOFS
In this section, we provide proof of the theoretical results from the main text. Namely Lemma 3.1, Proposition 4.1, Lemma 4.3, Corollary 4.4, Proposition 5.2, and finally Theorem 5.3.
Lemma 3.1. For every node v, the solution time t of heat diffusion on connected bundle G (V, E, O) with imput node features satisfies:
(Hx(())), = E H(l,, u
WEV
where H((t) is the standard graph heat kernel, and H(t, u) E R its the entry at (v, u).
Since LB = OgLOB we get Hz(t, u, H v)0g by the definition of the heat kernel.
Proof.
B.1 OVER-SMOOTHING: PROOFS OF SECTION 4.1.
In this section, we prove the results on the stable states and on over-smoothing. The first result follows straightforwardly from Lemma 3.1.
Proposition 4.1. Let Y be the output of a BuNN layer with t = co, where G is connected graph, and the bundle maps are not all equal. Then, u, E is connected such that y. * Ya almost always.
Proof. Consider stable signal Y E ker La and pick v E V such that O,, * O. As Y is stable, must have bundle Dirichlet energy, so we must have that Ouy, = O,y,. but as O * Og we have that y. A Yu. which holds except in degenerate cases such as when the matrices O. reducible, or when the original signal the zero vector.
The idea of the second result is that due to the injectivity of the positional encodings on the two nodes, we can set their restriction maps that the first layer zeroes the first channel of one node and the second channel of the second. The next K layers simply keep this signal fixed by realizing them as a fixed point of the bundle heat equation.
Proposition 4.2. Let X be 2-dimensional features nodes u, e with Pa + Pa. Then there is an > such that for every K, there exist a K -llayer deep BuNN with 2 dimensional stalks, gi()) a node-level MLP, ReLU activation, and [[(]] 1, with output Y such that ly yull 5.
connected G, and assume that there are two
(1 0 0
Proof. We will show the result for the limit co. Let O,, Id for * and O,
These restriction maps can be learned by assumption since Pu * Pv and MLPs are universal. Define
d
d
0
1
h
Xy +
X E R
2
2|E|
0
and denote ho the entry in the first dimension and h its second entry.
15

---

Published as a conference paper at ICLR 2025
Case 1: (hu)² 0 or (h1)² > O. We start with the case (ho)² 0. Set 8 (hu)². If ho 0,
0 Set the first bias to 0. The output before the activation will then be h()) (lhol 0)T for all o and h (0, hhol"" by Equation 5. Since is ReLU, it does not change the output. For the next layers, picking the same restriction maps O ww and W Id also do not change the output. Hence at layer K, the output at v is Ye = (phol. 0)* and at u Yu = (0 ( )" We conclude since Ily yn 2| 2 6. If (ha)² 2 0 and (hg 2 0 the argument is the same.
set the first weight matrix to be w = ( 0)
0 otherwise let W =
(-1 (-) 0
Case 2: (ha)² 0 and (h1)² O. Set 8 } and W = 0 and b (1, 0)T. Before the message diffusion step, the signal is constant on the nodes and equal to b. By Equation 5, the pre-activation output of the first layer is b (1, 0) at all nodes of and O (0, 1)2 u. Since all entries are positive, the ReLU activation leaves them fixed. The subsequent layer can be set as in the previous case keep the node embedding fixed. Consequently, lly, - yvll = 1 O
B.2 OVER- STURSSQQAHHING: PROOFS OF SECTION 4.2.
In this section we prove our results over-squashing and long-range interactions. The Jacobian re- sult follows straightforwardly from Lemma 3.1 and the definition of a BuNN layer, and the Corollary follows from the Jacobian result.
Lemma 4.3. Let BuNN be a linear layer defined by Equations 1, 2 & 3 with hyperparameter t. Then for any connected graph and nodes u, U, we have
① (BuNN (X)) Hl(t, u, w))""
Ox.
and therefore
a (BuNN (X)),
d, O W...... 2/E
lim 10000
8x,
Proof. The result follows from the closed-form solution of the heat kernel from Lemma 3.1. We start by applying the bundle encoder from Equation 2 that updates each node representation as h O wO.x, b. Since the O, do not depend on the signal X, we get
b (BuNN (X)) θ Õx,,
wO,x, Σ H(t, u. t), wO,OO, +b)
(8)
Ox,
H(t, u, u) v) wo
(9)
The second statement follows from the fact that H(t, v) + d
2T
To illustrate the flexibility of such result, we examine a setting in which want nodes to select which nodes they receive information from, therefore 'reducing' their receptive field.
Corollary 4.4. Consider n nodes u, v, and Wi, for i = I,...n 2, of a connected graph with 2 dimensional bundle such that pv: + Pug and Pa * Pw, Vi. Then in BuNN layer with MLP d, given channel, the node u can learn to ignore the information from all wis while keeping information from
Proof. We denote y the output of the layer, and index the two dimensions by super-scripts, i.e. (1))
() dy
@y By (0 0) all i. This would make
==
(2) Our goal is to have
# (0 0), while
the first channel of the output at insensitive to the input at all w,s while being sensitive to the input at node u.
Fix O = O, - -
and Ow. = (o 1
0 Such maps can always be learned by an MLP,
-
by the assumptions pa, Pu, and Pic and by the universality of ML.Ps. Let the weight matrix be
16

---

Published as a conference paper at ICLR 2025
Wi1
W12 W22
By Lemma 4.3 we get dy=
U,,, U,l, u)O w. .. - H((, 0,
W22
W12
W21
W21
W1
and
H(t D, w/ O WOO. = H((t U, wi)
W21
W22
Setting W21 and W22 to gives
8x
W11
w12
ay
0) and Qy Ux
H(l, U, u) (O W12) + O, as desired.
B.3 EXPRESSIVITY OF BUNNS: PROOFS OF SECTION 5.
We now turn to BuNN's expressivity. Before proving that BuNNs have compact uniform approxi- mation, we prove that MPNNs fail to have this property. This proves BuNNs' superiority and shows that uniform expressivity a good theoretical framework for comparing GNN architectures. Proposition 5.2. There exists a family 9 consisting of connected graphs such that bounded-depth MPNNs are not compact uniform approximators, even if enriched with unique positional encoding.
Proof. Let & be any family of connected graphs with an unbounded diameter (for example, the n x n grids with - oo). Let the depth of the MPNN be L. Let G E 9 be graph with diameter >> T., and let u and v be two nodes in VG at distance > Note that the output at u will be insensitive the input at u, and therefore, the MPNN cannot capture feature transformations where the output at depends on the input at u. This argument holds even when nodes are given unique positional encodings.
We now turn to main thcoretical contribution. The proof of Theorem 5.3 is split into two parts. The first proves that 1-layer BuNNs have compact uniform approximation over linear feature trans- formations. The second part is extending to continuous feature transformation, which is applica- tion of classical results.
We start by recalling the definition of a linear feature transformation over family of graphs G:
Definition B.1. A linear feature transformation E C(G, R" "" over a family of graphs 9 is an assignment of each graph G E 9 to a linear map LG : Rece -> Redd Here, linearity means that for any two node-signals X E "R and X2 E R Ru' and any real number a E IR, holds that LG (aX1)= = aLg (X1), and LG (X1 + X2) = LG (X1) + LG (X2).
We will need the following Theorem, which adapts classical results the universality of MLPs.
Theorem B.2. If a class of neural networks has compact uniform approximation over G with re- spect to linear functions and contains non-pon non-polynomial activations, then it has compact universal approximation over G with respect to continuous functions.
Proof. Classical theorems such as Theorem 1 in (Cybenko, 1989) allow us to approximate any continuous function over compact set in finite dimensional vector space by composing linear map C, an activation a, and an affine map A- +b. Given finite family of graph G. the space of node features on all graphs is a finite dimensional vector space. By assumption, we can implement the linear map, the activation, and the affine map. Hence, by composing them, we can approximate any continuous function over the compact set.
We ready to prove the paper's main result: that, given injective positional encodings, BuNNs are compact universal approximators of feature transformations.
Theorem 5.3. Let G be set of connected graphs with injective positional encodings. Then 2-layer deep BuNNs with encoder/decoder at each layer and 6 being 2-layer MLP have compact uniform approximation over G.
In particular, let > 0 and h be feature transformation supported Llee Kss C Liceg Road with K. C G finite and K C. IRd a compact set, then there 2-layer deep BuNN with width O (Eee (Va)) that approximates h with uniform error <
Proof. Reducing to linear approximation. suffices to show that a BuNN layer can approximate any linear feature transformation T. because can apply classical results such Theorem B.2
17

---

Published as a conference paper at ICLR 2025
to get universal approximation of 2-layer deep networks with activation. Following Definition 5.1, we aim to show that we can approximate a linear feature transformation L on any compact subset. For this, fix
O, the finite subset K. C G, and compact feature space K S R. In fact, assume that K = R since approximating a linear map on any compact feature space is equivalent to approximating it on the whole space because linear map defined on a neighborhood of the 0 vector can be extended uniquely to the whole vector space. Our goal is therefore to find a parameterization of a single BuNN layer such that for any graph G K and for any input feature X E g' we have ||LG (X) - BuNNG (X) 10 < We will show that L can be parameterized exactly. Since is linear, suffices to find a linear BuNN layer that satisfies for any G E K. and any X E RR a(L(X)) 0(BuNNX), X),, By Lemma 4.3, we have BuNN(X)), = H(t, u, o)O, WO" Hence, since MLPs are universal and the positional encodings are injective, it suffices to find bundle maps 8x.
8x
8x.
O Leex VG O (k)) and W such that RGdu EvEy OOW ((Ix), for every u, G and ax.
every E K.
Defining the encoder and decoder: In order to find such a BuNN, we first need a linear encoder lift: IRc 222c which will be applied at every node before applying 2ck dimensional BuNN layer. The lifting transformation maps each node vector X, to the concatenation of l vectors X. interleaved with vectors 0 R. This equivalent the linear transformation given by left multiplication by (Iexe; O,...; Lexes 0)" E RR2ckxe After the 2ck dimensional BuNN network, we will also need linear decoder pool R2ck - IK applied to every node individually, which is the sum of the & different o-dimensional vectors that are at even indices. This is equivalent to left multiplication by the matrix (Lexe: Oexes.... ................................ Lexes Oexe) Rex2ck These two can be seen as a linear encoder and linear decoder, often used in practical GNN implementations. We prove the result by adding the lifting and pooling layers and using the higher dimensional BuNN layer, i.e. we prove that BuNN pool cBuNN lift can approximate any linear maps which satisfy the encoder and decoder assumption of the Theorem statement.
Defining the 'universal bundle': We fix & = EGEK (Va,, so we can interpret our embedding space as a lookup table where each index corresponds node LGeex VG. In turn, we can as lookup table where each entry corresponds to a pair of nodes in our dataset K. Still thinking of the indices of the 2ck dimensions as 2c-dimensional vectors indexed by the nodes in our dataset, we define Og E E (2ck) as a block diagonal matrix with & different 2c- 2e-dimensional blocks, where the k;th block is denoted O% These are all set to the identity except for the block at the index corresponding to node u, which is defined
think of the parameter matrix W E (VEE=e<VVe)x(x<c>>>
Oexe Lexc which is 2c x 2c matrix that acts by permuting the first e dimensions with the Loxe Ocxe
second e dimensions.
Computing the partial derivatives. Since our model BuNN is a composition of linear maps, and since the maps pool and lift are applied node-wise, we get
d (BuNN (X))
Ox,
a BuNN (lift (X))
pool
lift
lift (X,))
-Texe, Oaxe:
Texes Oexe) H(t, u, v) O (Texe, Oexes Texes Ocxe)
= Hlt,, u, v)
Σ
(Texer Oxxe) O o kz (Texe; Oxxc)"
1Sk1, Mask
We proceed by partitioning the indexing by (k1, k2) into four cases. The first case is C1 {(k), l2z) such that (k) A u, v and k2 * u, }}} for which both Oh and O act like the identity. The second case is C2 {Kk), k2) such that k = u and K2 A b} where 0% flips the first rows with the second e rows and 0 acts like the identity. C3- = {(A,, kz)) such that b2 = e and k # u} where 0 flips the first columns with the second e columns, and O acts like the identity on the rows. Finally, the last case is when b = u and h = v in which 0h flips the rows, and O flips the columns. Note that by the injectivity assumption defined in Definition D. each node has
18

---

Published as a conference paper at ICLR 2025
unique positional encoding and therefore there such matrices be parameterized for all nodes simultaneously.
...-((4 H,l,,,, Σ
(Lexes Ocxe) O 2 Wa O fh (Texas Oaxe)* )
...-
ISKi, K2Sk
Who l2 Wkikk loa
Whi Whi
W l2 Whikk
Whi Whi
=H(t, u, (Toxes 0oxe)
Σ (AE
Σ
(Ar Ka)ECC
Σ
Whikk
Whikk Whi
Way Wor
Wue W 00
(Texes Oexc) T
Whi l2
Wgp 01
u)t,u,v)
Σ
Whi
Σ (KI
Whi 22
Σ (21,s
Wki + W
10
(2)Cs
" Where the last line is obtained by applying (Lexe: Oexe) on the left and (Lexe: 0 Oexe) on the right, an operation that selects the upper left c x block. We observe that setting all W whi b2 W to Oexe and setting Wir in the same graph and Oexe otherwise. This allows conclude that any linear layer can be parameterized, completing the proof of the theorem.
8(LX), if the nodes corresponding to and v lie [[[,u.v] Ox.
C DISCUSSION ON COMPACT UNIFORM APPROXIMATION VERSUS UNIFORM APPROXIMATION
A strong definition of expressivity that deals with infinite collections of graphs proposed in Rosenbluth et al. (2023). This definition subsumes graph-isomorphism testing (where the input feature on graphs constant). Furthermore, also deals with infinite families of graphs, opposed to most mainstream theorems of GNN expressivity, which are proved for graphs of bounded size (e.g. Azizian & Lelarge (2021); Geerts & Reutter (2022)). See Section for the notation and definition of features transformations.
Definition C.1 (From Rosenbluth al. (2023)). Let E and take R as feature space. Consider collection of graphs G. Let Ω S C (G, be a set of feature transformations over G, and let H C (G, 125, IR=') a feature transformation over We say that Q uniformly additively approximates H. notated 52 ss H if Ve > 0 V compact K" C Ruce 3F Ω such that:, VG E VX E Re ||Fc - HG (X)ll S where the sup norm aloc is taken over all nodes and dimensions of ngd! dimensional output.
Note that this definition differs from our Definition 5. in that it requires uniform approximation over all graphs in G simultaneously, while we allow the width to vary with the finite subset I. C G. similar to how classical results allow the width to vary with the compact set over which to approximate the function. Such a definition has proven useful in Rosenbluth al. (2023) to distinguish different aggregation functions and in Rosenbluth et al. (2024) to distinguish MPNNs with virtual nodes from Transformers. However, we argue that the definition above too strong for a finite parameter GNN. This is because it requires uniform approximation over a non-compact set, which contrasts with traditional work on expressivity and is generally unfeasible and impractical. Indeed, finite- parameters MLPs are not universal over the whole domain I under the for-norm On infinite collection of featured graphs, the topology is the disjoint union topology on Lkee Read a compact subset consists of finite set of graphs, and for each graph G only non-zero compact subset of and For these reasons, we introduce Definition 5.1, which is still rich enough to distinguish between BuNNs and MPNNs.
19

---

Published as a conference paper at ICLR 2025
D WHY CLASSICAL ARGUMENTS DO NOT APPLY TO COMPACT UNIFORM APPROXIMATION
A seminal result in GNN expressivity is Theorem 2 in Morris al. (2019). In this section, we discuss why the arguments do not hold for compact uniform approximation. even when enriched with injective positional encodings. We start by formally defining injective positional encodings.
Definition D.1. A positional encoding æ a graph G is map π (T) V - R which assigns k-dimensional feature to every node... A positional encoding я on a family of graphs g is a positional encoding on all graphs G E G. A positional encoding क is injective on & for any graph G g and every node E G there is no other node v E H for any H E 9 with π (u) = T()). In other words, each positional encoding corresponds to a unique node in the dataset.* 4
The setting in Morris et al. (2019) Theorem 2 considers single graph with fixed node features (or colors/labels). In contrast, our Theorem is for family of graphs, and more importantly, for each graph G, the node features are not fixed but vary in any compact subspace of feature space. This means that for a single graph of size n, our statement holds, for example, on the unit cube [0, 1)", while the result in Morris et al. (2019) only holds for a specific point in E
Fixing the node features is precisely what makes the construction in Morris et al. (2019) possible. In- deed the proof starts by assuming that the initial coloring is "linearly independent modulo equality". denoted by ((O) in their proof. This property on node features is indeed central to the construction. is used several times, for example "Observe that colors are represented by linearly independent row vectors" and "F((+1) is linearly independent modulo equality". Such an assumption is possible when dealing with fixed node features: since there are at most n colors, it suffices to take a one-hot encoding of those colors, which fits in a feature space of dimension at most n. This assumption is also crucial for the injectivity of the sum aggregation.
In our setting, the node features can take any value a continuum of features, each belonging to a compact subspace K of R. Encoding such continuum as a one-hot encoding cannot be done in a finitediimensional vector space (since there are {0, 13 possibilities, which can be uncountable, and each need to be linearly independent). Hence, their construction fails in the setting where the node features are not fixed but can vary on any compact subspace of R
E ALGORITHMIC AND IMPLEMENTATION DETAILS
In this section, we provide more details on the implementation of BuNNs. We start by discussing how to use several vector-field channels when the input dimension is greater than the bundle dimen- sion. We then discuss how to use several bundles at once when single bundle is insufficient. We then combine both views, namely having several vector-field channels on several bundles once. Finally, we describe how we compute our bundle maps in the experiments.
Extending to several vector-field channels. When the signal dimension exceeds the bundle dimen- sion, i.e. c >> d, we cannot directly apply BuNNs to the input signal. In that case, we first transform the signal into a hidden dimension, a multiple of the bundle dimension, i.e.c dp. We reshape the input signal into p channels of d-dimensional vector fields, where we apply the diffusion step (Equation 3) on each p channels simultaneously, and we apply the weight matrix W E Rdpxxdp by first flattening the node signals into dp dimensions, then multiplying by W, and then reshaping into p channels of d dimensional vector fields.
Extending to several bundles. Learning high dimensional orthogonal matrix O(d) becomes ex- pensive since the manifold of orthogonal matrices is d(d-1) dimensional. However, we can compute many fow-dimensional bundles in parallel. In practice, we found that using several 2-dimensional bundles was enough. Computing b different 2-dimensional bundles requires only b-parameters since the manifold O(2) is 1-dimensional. We, therefore, also use different 'bundle channels' given by an additional syper-pprameter- - the number of bundles, which we denote b. Given an input signal of
2
'Note that on symmetric graphs such as the cycle. it is not possible to have injectivity if the positional encoding is equivariant. However, in practice many steps of the 1-WL. color refinement will be equivariant and satisfy such assumption on most graphs.
20

---

Published as a conference paper at ICLR 2025
decompose the signal into b bundle channels of dimension d. We can compute the diffusion step (Equation 3) for each bundle in parallel. For the update step (Equation 2), apply the weight matrix W Rddxdd by first flattening the node signals into bd dimensions, then multiplying by W, and then reshaping it into bundle channels of d dimensional vector fields over b different bundle structures.
dimension db,
Remark E.I. We note that using b different d dimensional bundles is equivalent to parameterizing subset of one bdddimensional structure, consisting of the orthogonal map O E O C Rddddb that are block diagonal matrices O = ① 1=1...b O, with each O, E (dd).
Extending to several bundles and vector-field channels. We can combine the above two obser- vations. Given an input signal of dimension bdp, can subdivide this into b different bundle structures dimension and p channels for each bundle. diffuse on the appropriate bundle structure and flatten the vector fields into
vector before applying the learnable parameters.
Computing the bundle maps. In our experiments, we noticed that having several bundles di- mension 2 was more efficient than one bundle of large dimensions, while there was clear per- formance gain when using higher dimensional bundles. To compute the bundle maps O we therefore only need b rotation angles 6., one per bundle. In experiments, reflections using the python package Obukhov (2021) or direct parameterization. For direct param- eterization, we do the following: since the matrix group ((2) is disconnected, we always take b
Housholder
to be and parameterize half the bundles rotation matrices r(0) cos (0) sin
sin (0) COS (0)
and the other half to correspond to matrices with determinant -1, which can be parameterized by (0) cos (0) sin (0)
sin (0) cos (0) We compute the angles in Equation where the network (1()
is either an MLP or a GNN. The network & is either shared across layers or differing at every layer.
Taylor approximation algorithm. We now provide pseudo-code on how we implement Equa- tions 2, and 3. We then proceed with complexity analysis. The key idea of the algorithm is that the bundle heat kernel can approximated efficiently using the standard graph heat kernel.
Algorithm Taylor expansion implementation of a BuNN layer
Input: Normalized graph Laplacian L, Orthogonal maps O Ve E G. Node features ( E gaxd, Time t, Maximum degree K, Channel mixing matrix W()), bias (())
Output: Updated node features (
1: h 0 X Vo V
> Sync.: Go to global representation
2: HO + HW +b(*)
b Update features with parameters b approximation of degree
3: ((++1) < (9)
4:
for k = 1..... K do
5: H + (H
Term of degree Approximation of degree k
6:
((+1)) x (++1) + (
7: end for
8: x
(e+1)
+ 0
10) X
((+1) V
b Deyne.: Return to local representation
9: return (++1)
The complexity of the algorithms is as follows. There 3 metrix-vector vector multiplications done at each node in lines 1, 2, and 8, which are done in O (3d(V)), The for loops consist of matrix-matrix multiplications, which are done in O (d|El) with sparse matrix-vector multiplication. The memory complexity is O ((d + ) Vi) since we need to store d dimensional vectors and the orthogonal maps for each node. The exact implementation is described in Algorithm
Spectral method. We now describe how to implement a BuNN layer using the eigenvectors and eigenvalues of the Laplacian.
21

---

Published as a conference paper at ICLR 2025
Algorithm 2 Spectral implementation of a BuNN layer
Input: Eigenvectors and eigenvalues graph Laplacian (Ø,, Ax); Orthogonal maps Og Ve G,
Node features ( E gaxd Time t, Maximum degree K, Channel mixing matrix (()), bias b(2)
Output: Updated node features X(e)
1: h
0 x!!") Ve
Sync.: Go global representation
2: HO) HW +b())
Update features with parameters Spectral solution to heat equation
3: X E, E "i
(e+1)
+ 0 (0) (e+1) VVEV
4: X
X 5: return ((+1))
P Desyne.: Return to local representation
E.1 HOUSEHOLDER REFLECTIONS.
Many different parameterizations of the group O(n) exist, While direct parameterizations are pos- sible for 2, becomes increasingly complex to do so for larger n, and general method working for all n is a desirata. While there are several methods to do so, we use Householder re- flection since is used in related methods (Bodnar al., 2022). We use the Pytorch package from (?). Given given & vectors 0. Rd, define the Householder matrices H, I 2 wan." and ||ev
define = 11 H₁ All orthogonal matrices may be obtained using the product of d such matri- ces. Hence the map R ddd + (dd) mapping V (vt) to a parameterization the orthogonal group. We use pytorch implementations allowing autograd provided in (?).
F EXPERIMENT DETAILS
In this section provide additional information about the experiments on the heterophilic graph benchmarks, the LRGB benchmarks, and the synthetic experiments. All experiments were ran on a cluster using NVIDIA A10 (24 GB) GPUs, each experiment using most 1 GPU. Each machine in the cluster has 64 cores of Intel(R) Xeon(R) Gold 6326 CPU at 2.90GHz, and 500GG of RAM available. The synthetic experiments from Section 6.1 were run on CPU and each run took roughly 20 minutes. The heterophilic experiments from Section 6 were run GPU and varied between minutes 1.5 hours. The LRGB experiments were run on GPU and varied between 0.5 hours and 4 hours.
F.1 LRGB: TRAINING AND TUNING.
For pept i fune- and Dwivedi et al. (2022). We fix ppper-parameters to be the best GCN hyper-parametest from Tönshoff et al. (2023), and tune only BuNN-specific parameters as well as the use of BatchNorm. In Table 5, report the grid of hyper-parmmeters that searched, and denote in bold the best combinations of hyper-parameters. The parameters fixed from Tönshoff et al. (2023) are the following:
we use a fixed parameter budget of ~ 500k as in
Dropout 0.1
Head depth 3
Positional Encoding: LapPE for struct and RWSE for func
Optimizer: AdamW with a cosine annealing learning rate schedule and linear warmup.
Batch size 200
We use skip connection implemented in Dwivedi et al. (2022) and not in Tönshoff et al. (2023). That is, the skip connection does not skip the non-lincarity.
For the BuNN specific parameters, we use dimensional bundles, whose angles we compute with the help of small SumGNN architecture using sum aggregation as defined by g! (Wxx + W, EweN()) x where the input dimension is the hidden dimension, the hidden
22

---

Published as a conference paper at ICLR 2025
dimension is twice the number of bundles and the output is the number of bundles. The number of SumGNN layers a hyper--arameter we tune. When it is we use a layer MLP with hidden dimension also twice number of bundles. For each hhper--ppramaee parameter configuration, set the hidden dimension in order to respect to the parameter budget.
All Values
Best Values
Parameters
func
struct
PascalVOC-SP
Num bundles b Number of BuNN layers Number of SumGNN layer
4, 8, 16
16
16
16
20
6
20
0 - 3
1
0
0
Weight decay
0, 0.1, 0.2, 0.3
0.2
0.2
0.2
Time t.
0.1, 10, 100
Table 5: Grid of hyper-parameters for peptides-func, paptides-struct, and PascalVOC-SP.
F.2 HETEROPHILIC GRAPHS: TRAINING AND TUNING.
For the heterophilic graphs we use the source code from Platonov et al. (2023) which we add layer definition. We report all training parameters that have tuned. Namely, activation functions, the Adam optimizer with learning rate 3 x 10 and train all models for 2000 To compute the bundle maps,
GELU
epochs and select the best epoch based the validation compute the parameters / with GraphSAGE architecture shared across layers (ø method = shared) or different each layer method not shared), with hidden dimension dimension the number of bundles. The number of layers of this GNN is hyper-parameter we tuned, which when set to we use 2 layer MLP. For each task we manually tuned parameters, which are subsets of the combinations of parameters in the grid from Table 6. The implementation of the heat kernel used either truncated Taylor series with degree 8, or the spectral implementation. We report the best performing combination of parameters in Table 7. For the NSD baseline
code from Bodnar et al. (2022) and tune equivalent parameters. We report the grid of hyperparameters in Table 8 and best values in Table 9
Parameters
All Values 256, 512
Hidden dim Num bundles
2, 4, 8, 16, 32, 64, 128, 256
Bundle dimension
Number of BuNN layers
1
Number of GNN layer ó
0 8
Time t
0.1, 10, X.V
shared Ø
Dropout
0.0, 0.2
Learning rate
3 x 10-4 3 10
Table 6: Parameters searched when tuning on the heterophilic graph benchmark datasets.
Parameters
Best Values
roman- eaaan-empiree
amazon - ratingg
minesweeper
tolokers
quest ions
Hidden dim
512
512
512
512
256
Num bundles
4
128
256
Bundle dim
2
2
Number of BuNN layers
2
6
Number of GNN layer
0
7
Time
100
1.5
shared 6
X
X
Dropout
0.2
0.2
0.2
0.2
0.2
Learning
x 10-4
3 10
10
3
10-5
10
Table 7: Best parameter for each dataset in the heterophilic graph benchmarks.
23

---

Published as a conference paper at ICLR 2025
Parameters
All Values
Hidden dim Sheaf dimension d Number of layers Dropout
64, 128, 256", ,512* ,56",512 2, 4, 8
1, 2,
4", 5*, 6*, 7",8" 0.0, 0.2
Learning rate
10-1.3 .33x 10-5
Table 8: Parameters searched when tuning NSD on the heterophilic graph benchmark datasets. Parameters marked by ran out of memory on some datasets.
Parameters
Best Values
roman-empil
amazon-raa rat ings
mi minesweeper
tolokers
quest ions
Hidden dim Sheaf dim
256
256
512
64
64
2
2
Number of layers
8
Dropout
0.2
0.2
0.2
0.2
0.2
Learning rate
3 x 10-4
3 10
x 10
3 10
3 x 10
Table 9: Best parameter for NSD on each dataset in the heterophilic graph benchmarks.
G EMPIRICAL RUNTIME
Table 10: Average training time, over runs, for different architectures with 5 layers and hidden dimension of 512 on the Heterophilic datasets (averaged over 3 runs). BuNN and NSD single dimensional sheaf. BuNN implemented spatially. All experiments were performed on an NVIDIA A10 (24GB) GPU.
avg. training time roman-empi i re amazon-ratinmg minesweeper tolokers quest ons
# num nodes
22,662
24,492
10,000
11,758
48,921
num edges
32,927
93,050
39,402
519,000
153,540
SAGR
1:45
1:43
0:48
1:15
3:01
GAT
2:33
2:42
1:11
3:24
5:22
GT
3:31
4:12
1:52
4:20
7:57
NSD
9:16
7:13
BuNN
4:21
5:18
2:48
3:51
9:32
H ABLATION: POSITIONAL. ENCODING ABLATION
We perform an ablation on the use of PE on the peptides func and peptides-sttruee datasets. We retrain model using the best hyperparameter where change the used PEs. We compare the two main PE used in graph machin learning, namely Laplacian Positional encoding (LPE) and Random Walk Structural Encodings (RWSE), and we consider using PE baseline. Results can be found in Table 12. We observe that using PE is always beneficial to not using PE, however each task seem admit preferred PE, since LPE is better for Pepti strutt and RWSE for reptides-func
I ABLATION: IMPORTANCE OF W AND b
We run an ablation on the importance of the W and parameters in Equation 2. We retrain model removing them and compare it to a model trained using them. We report results Table 13. The results suggest that these parameters help. However, the result suggest that they are not essential to achieve good results, as even without them, the model performs well and beats all but one of the baselines pept des-func.
24

---

Published as a conference paper at ICLR 2025
Table 11: Average time per epoch, 5 epochs, for different architectures with 6 layers and 500k parameters on the LRGB datascts. BuNN is implemented spectrally and uses single 2 dimensional bundle. All experiments were performed on an NVIDIA A10 (24GB) GPU.
avg. time / epoch Pept Septides-funee func Peptides-stzuer PascalVoc-SE
# graphs
15,535
15,535
11,355
nodes
2,344,859
2,344,859
15,955,687
edges
4.773,974
4.773.974
32.341,644
GCN
5.6s
5.3s
15.2s
GatedGCN
8.5s
8.3s
22.5s
GPS
17.2s
17.2s
38.2s
BuNN
12.7s
12.5s
28.7s
Table 12: Positional encoding ablation on peptides-fumm and peptides-struct
RWSE
LPE
No PE
Pept idee-furn Test A
72.76 0.65
72.25 + 0.51
71.76
0.68
Pepti saptiides-stucce Test MAE +
25.02 + 0.15
24.63 + 0.12
25.32 0.19
J TREE-NEIGHRORSMATCH TASK
As additional synthetic task, we evaluate BuNNs on the Tree-Neiighorssatthh task, pro- posed in Alon Yahav (2021) to show that MPNNs suffer from over-squashing. We use their code and setup to evaluate the capacity fo BuNNs alleviate over-squashing. We use their reported re- our own results, ran using their experimental setup, for a 2-layer deep BuNN with
sults and t and 32 bundles. We report the results in Figure 5. We observe that BuNN beat all MPNNs by a large margin, with perfect until r = 6. As the task gets harder with a larger depth, the accuracy for BuNNs drops slower than the accuracy of MPNNs. These results confirm that BuNNs alleviate over-squashing.
Table 13: W and b importance ablation pept des- and peptidee des-structt
wits
without
Pept desiidesfiin func Test AP
72.76
0.65
70.75 4 0.36
Peptides -eriiddssssstutteeeeoee struct Test MAE
24.63
0.12
25.28 + 0.32
25

---

Published as a conference paper at ICLR 2025
Accuracy v r
1.0
BuNN
0.8
GCN
0.72
3.50
0.21
0.2
634
0.00
620
0.0
Figure 5: Results for Tree -ree-MeighborsMaee task.
26


Reviews:
[{"review_id": "D3K6v8SW8c", "reviewer": "Reviewer_9bAu", "summary": "This paper extends previous work on graph neural networks equipped with cellular sheaves by introducing a more efficient way of computing the heat diffusion over bundles. The key idea of the authors is to assume a simpler form for the sheaves, namely a flat vector bundle, which can be thought of as an orthogonal map associated with each node (as opposed to one associated with each node-edge pair). This assumption significantly simplifies the computations necessary to learn the bundles from data. In particular, the heat diffusion kernel can now be computed efficiently during the learning of the bundles. This enables the authors to use spectral methods to run the diffusion process for long periods of time addressing the over-squashing problem on graphs. This flat vector bundles also bring the strength of the cellular sheaves in that the fixed point of diffusion is no longer a trivial state where all nodes share the same features. Rather, the long time limit now has richer structure where variation is allowed across the nodes. This addresses the over-smoothing problem. The authors propose GNN architecture using flat vector bundles and heat kernels (computed both using Taylor expansions and spectral methods) and demonstrate improved performance on synthetic data sets and real data sets.", "questions": "What price is paid for going rom the full cellular sheaves to the flat vector bundle in expressivity? In practice, in what settings is the flat vector bundle approach good enough? For example, it seems like that for sparse graphs the orthogonal maps associated with each node would be good enough.\n\nSimilarly, can the authors provide a comparison between the structure of the fixed points for cellular sheaves and flat vector bundles? It seems like some of the \"richness\" of the long time limit of diffusion is lost with the flat vector bundles. How does this reflect on expressivity of flat vector bundles compared with cellular sheaves?\n\nDoes the transformation learned for each node need to be orthogonal in practice? It seems that the approach is somewhat limited because the learned matrices must be orthogonal. It is difficult to learn matrices that are constrained to be orthogonal. Can the authors relax this requirement in practice and use invertible matrices or a fully unconstrained matrix?\n\nIt seems like a key aspect of the proposed approach is how the orthogonal transformations at each node are learned? The authors mention MLPs or even GNNs that take in the graph structure and positional encodings of the nodes. What type of positional encodings are best for learning the orthogonal transformations? Why are the node features not taken into account when computing the orthogonal map for each node? Are the authors trying to keep the graph geometry distinct from the feature information? \n\nIn Table 3, how does the BuNN approach outperform the NSD for the data sets in which memory is not a limit for NSD? Is the depth of the models in BuNN larger helped by spectral methods? It seems like depth was 1 for the amazon-ratings data set. Does NSD overfit the data because of its use of cellular sheaves?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "This reviewer generally liked the paper. The idea of simplifying cellular sheaves using flat vector bundles is clever. This idea seems to solve the key problem with the computational complexity challenges faced by neural sheaf diffusion. Another strength of the paper is the clever synthetic toy models which demonstrate the gain from the proposed bundle architecture very clearly, especially when it comes to the over-squashing problem. The empirical performance of the proposed method is also impressive, especially on heterophilic graphs like minesweeper. The novel notion of expressivity introduced by the authors is also intriguing.", "weaknesses": "A weakness of this paper is that the central idea is arguably a simple extension of the neural sheaf diffusion paper. It is not clear that this constitute significant advance. This is balanced by the impressive empirical results in the paper.\n\nThe authors highlight well the advantages of replacing the cellular sheaves with flat vector bundles. However, it is not clear what is the cost of doing so. How does expressivity suffer with this assumption? Are there circumstances under which this simplification works better? For example, does the sparsity of the graph play role. More discussion on this would have helped the reader. The authors nicely show the existence of fixed points for the diffusion process under flat vector bundles and that features across nodes are related by orthogonal transformations. How does this compare with the structure of fixed points under general cellular sheaves? How does this impact the expressivity of the bundle neural networks compared with their sheave counterpart?\n\nThis reviewer struggled at times to understand how much the principled geometric approach brought over from topology and differential geometry contributed to formulation of the proposed GNN. Although interesting, can the author's approach be though of as learning a transformation on each node. With the formality relaxed, it is not clear if this transformation needs to be orthogonal. The authors seems to go to quite a bit of trouble to ensure that the learned matrices at each node are orthogonal. Of course, Riemannian manifolds require this, but is it really needed? Would invertible matrices done just as good of a job? It seems like in practice they are limited to using only 2-dimensional vector spaces on each node, with the transformation matrices parameterized explicitly to be orthogonal.\n\nThis reviewer was bothered at times by the strong use of language by the authors. The claim that the proposed bundle approach solves the over-squashing problem is over-stated. Any spectral method will address over-squashing. The authors have proposed an efficient way to compute the heat kernel more efficiently by replacing cellular sheaves with flat vector bundles. This allows them to run diffusion for longer times. There is no conceptual new solution to over-squashing.", "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": "p5hTf9rg2J", "reviewer": "Reviewer_rcKT", "summary": "The paper proposes a new graph neural network architecture called Bundle Neural Networks (BuNNs). The architecture is based on message diffusion on flat vector bundles which are topological structures that assign to each node a vector space and an orthogonal map. They seem to be inspired from Sheaf neural networks and, in different forms, are claimed to address the problems of overs-moothing and over-squashing. Theoretical analysis on the feature transformation with injective positional encodings shows uniform expressivity. Further analysis on the properties of BuNNs and experimental results on both synthetic and real-world datasets, are provided for validating the proposed model.", "questions": "Please see weaknesses", "limitations": "", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. The paper addresses an important problem of oversmoothing and oversquashing seen in GNNs via vector bundles, building on cellular sheafs in GNNs.\n1. The use of flat vector bundles with orthogonal maps to reduce the computational complexity is interesting.\n1. The paper is well-organized.", "weaknesses": "1. The experiments to validate oversmoothing and oversquashing are not comprehensive. The authors can test the oversquashing on Neighbours match dataset [1], which has been used well in multiple papers. Further the authors can explore using effective resistance metric to measure oversquashing [3] and [4].\n1. Related works which are designed via dirichlet energy optimization like [2] are not discussed. May help to discuss the difference with [2] in terms of results on oversmoothing.\n\n### References:\n[1]$~$ Uri Alon and Eran Yahav. On the bottleneck of graph neural networks and its practical implications. arXiv preprint arXiv:2006.05205, 2020.\n\n[2]$~$ Fu, Guoji, et al. \"Implicit graph neural diffusion based on constrained Dirichlet energy minimization.\" arXiv preprint arXiv:2308.03306 (2023).\n\n[3]$~$ Mitchell Black, Zhengchao Wan, Amir Nayyeri, and Yusu Wang. Understanding oversquashing in gnns through the lens of effective resistance. In International Conference on Machine Learning, pages 2528–2547. PMLR, 2023.\n\n[4]$~$ Dong, Yanfei, et al. \"Differentiable Cluster Graph Neural Network.\" arXiv preprint arXiv:2405.16185 (2024).", "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": "a1wNr8oxDt", "reviewer": "Reviewer_uDEq", "summary": "This paper proposes to perform message-passing over a diffused graph using flat vector bundles. This is achieved by transforming each node state within its own vector space using a learned and distinct orthogonal map. This is the message that a node sends to each neighbor, where the inverse of the orthongonal map from the receiving node is used as a decoding mapping. The authors show that this can avoid over-smoothing, over-squashing, and can be a compact uniform approximation over a set of graphs. Experiments show the effectiveness of the proposed approach.", "questions": "* See weaknesses", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "The proposed method is sound and very general. BuNNs have many favorable and interesting properties. I particularly like the property regarding fixed points. The paper is nicely written and mostly understandable.", "weaknesses": "* Eq. 2: From the theoretical derivation, this step does not seem needed as $\\mathcal{H}_B$ already applies $O_v$ and $O_u$. The only provided reason is that it helps to prove Theorem 5.3. Are there other reasons why $W$ and $b$ are introduced? What happens when these are removed?\n* BuNNs and their properties are mostly compared to MPNNs applied to the original. It seems fairer to me to also compare its properties to Graph Transformers and infinite-depth MPNNs (e.g. [1]) (Table 1).\n* Theorem 5.3: \n    * This statement is quite interesting but I do not fully understand it. In which cases are injective positional encodings achievable? Would it be satisfied when $\\mathcal{G}$ is a set of WL distinguishable graphs and we use WL colors are positional encodings?\n    * How does it deal with two nodes $u,v$ in the same graph with the same WL unfolding tree? By my understanding of the injective positional encodings, they could still have the same initial state. The MLP (Eq. 1) would produce the same $\\mathbf{O}_v = \\mathbf{O}_u$ for both nodes.\n    * Which changes would be needed for graph transformers to also satisfy this universality?\n* Table 2: A comparison with Graph Transformers or infinite-depth MPNNs would be insightful\n* Table 3: It seems like the authors performed experiments for NSD and BuNN, but hyperparameter ranges and optimal hyperparameters are only presented for BuNN.\n* Details about runtimes are missing. It would be important to have execution times for at least one experiment for BuNN for a comprehensive evaluation.\n\n---\n[1] Gu et al., Implicit graph neural networks, NeurIPS (2020).", "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 (Spotlight)