Bundle Neural Networks for message diffusion on graphs
Jacob Bamberger, Federico Barbero, Xiaowen Dong, Michael M. Bronstein
TL;DR
BuNNs address fundamental limits of local message passing in GNNs by replacing neighbor-centric messaging with global diffusion of features over flat vector bundles via the bundle heat kernel $\mathcal{H}_\mathcal{B}(t) = \exp(-t\,\mathbf{\mathcal{L}}_\mathcal{B})$. The model learns per-node orthogonal maps and propagates information through diffusion, mitigating over-smoothing and over-squashing while enabling long-range interactions. The authors prove a compact uniform universal approximation guarantee under injective positional encodings and demonstrate state-of-the-art results on heterophilic and long-range graph benchmarks, including Peptides-func and LRBG, establishing BuNNs as a scalable, theory-backed alternative to standard MPNNs and prior diffusion-based approaches. Overall, BuNNs offer a principled diffusion-based framework that enhances expressivity and scalability for graph learning tasks.
Abstract
The dominant paradigm for learning on graph-structured data is message passing. Despite being a strong inductive bias, the local message passing mechanism suffers from pathological issues such as over-smoothing, over-squashing, and limited node-level expressivity. To address these limitations we propose Bundle Neural Networks (BuNN), a new type of GNN that operates via message diffusion over flat vector bundles - structures analogous to connections on Riemannian manifolds that augment the graph by assigning to each node a vector space and an orthogonal map. A BuNN layer evolves the features according to a diffusion-type partial differential equation. When discretized, BuNNs are a special case of Sheaf Neural Networks (SNNs), a recently proposed MPNN capable of mitigating over-smoothing. The continuous nature of message diffusion enables BuNNs to operate on larger scales of the graph and, therefore, to mitigate over-squashing. Finally, we prove that BuNN can approximate any feature transformation over nodes on any (potentially infinite) family of graphs given injective positional encodings, resulting in universal node-level expressivity. We support our theory via synthetic experiments and showcase the strong empirical performance of BuNNs over a range of real-world tasks, achieving state-of-the-art results on several standard benchmarks in transductive and inductive settings.