Supercharging Graph Transformers with Advective Diffusion
Authors: Qitian Wu
Chenxiao Yang
Kaipeng Zeng
Michael M. Bronstein
Venue: ICML 2025 poster (2025)
Keywords: geometric deep learning
graph machine learning
topological shifts
transformers
graph neural networks
Primary area: deep_learning->everything_else
TL;DR: We propose a physics-principled graph Transformer for improving the generalization capability against topological/structural distribution shifts

Abstract: The capability of generalization is a cornerstone for the success of modern learning systems. For non-Euclidean data, e.g., graphs, that particularly involves topological structures, one important aspect neglected by prior studies is how machine learning models generalize under topological shifts. This paper proposes AdvDIFFormer, a physics-inspired graph Transformer model designed to address this challenge. The model is derived from advective diffusion equations which describe a class of continuous message passing process with observed and latent topological structures. We show that AdvDIFFormer has provable capability for controlling generalization error with topological shifts, which in contrast cannot be guaranteed by graph diffusion models, i.e., the generalization of common graph neural networks in continuous space. Empirically, the model demonstrates superiority in various predictive tasks across information networks, molecular screening and protein interactions

Supercharging Graph Transformers with Advective Diffusion
Qitian Wu! Chenxiao Yang?
Kaipeng Zengà Michael Bronstein
Abstract
and graph Transformers (Ying et al., 2021; Wu et al., 2021; Rampásek et al., 2022; Wu et al., 2022b), which propagate information among arbitrary node pairs through global at- tention. GNNs can be seen as discretized versions of local diffusion equations on graphs (Atwood & Towsley, 2016; Klicpera et al., 2019; Chamberlain et al., 2021a), while graph Transformers can be considered as the counterparts of non-local diffusion (Wu et al., 2023; 2024c). Linking graph learning models with diffusion equations offers pow- erful tools from the domain of partial differential equations (PDEs), allowing us to study the expressive power (Bod- nar et al., 2022), behaviors such as over-smoothing (Rusch et al., 2023) and over-squashing (Topping et al., 2022), the settings of missing features (Rossi et al., 2022), and guide architectural choices (Di Giovanni et al., 2022).
The capability of generalization is a cornerstone for the success of modern learning systems. For Euclideid data, e.g., graphs, that particu- larly involves topological structures, one impor- tant aspect neglected by prior studies is how ma- chine learning models generalize under topologi- cal shifts. This paper proposes ADYDIFFORMER physies-inspired graph Transformer model de- signed to address this challenge. The model is derived from advective diffusion equations which describe a class of continuous message passing process with observed and latent topological struc- tures. We show that ADYDIFFORMER has prov- able capability for controlling generalization er- ror with topological shifts, which in contrast can- not be guaranteed by graph diffusion models, i.e., the generalization of common graph neural networks in continuous space. Empirically, the model demonstrates superiority in various predic- tive tasks across information networks, molecular screening and protein interactions
While significant efforts have been devoted to understanding the expressive power of graph learning models, the gener- alization capabilities of such methods are largely an open question. Recent works attempt to study generalization of graph learning models (particularly GNNs) from various perspectives such as extrapolation in feature space (Xu et al., 2021), subgroup faimess (Ma et al., 2021), invariance prin- ciple (Wu et al., 2022a), feature propagation (Yang et al., 2023), causality (Wu et al., 2024a), and training dynam- ies (Yang et al., 2024). However, most of these works focus on the distribution shifts of features and labels. In many critical real-world settings, the training and testing graph topologies can be generated from different distributions (e.g., molecular structures with diverse drug likeness) (Koh et al., 2021; Hu et al., 2021; Bazhenov et al., 2023; Zhang et al., 2023), a phenomenon we refer to as "topological distribution shift". This can be a predominant nature of non- Euclidean data in contrast with commonly studied feature and label shifts in Euclidean space. Despite its practical sig- nificance, how to enable graph learning models to generalize under topological shifts still remains unclear.
1. Introduction
Learning representations for non-Euclidean data is essential for geometric deep learning. Graph-structured data in partic- ular has attracted increasing attention, as graphs are a very popular mathematical abstraction for systems of relations and interactions that can be applied from microscopic scales (e.g. molecules) to macroscopic ones (social networks). Common frameworks for learning on graphs include graph neural networks (GNNs) (Scarselli et al., 2008; Gilmer et al., 2017; Kipf& & Welling, 2017), which operate by propagating information between adjacent nodes of the graph networks,
'Eric and Wendy Schmidt Center, Broad Institute of MIT and Harvard TToyota Technological Institute Chicago 'Shanghai Jiao Tong University "University of Oxford 'Aithyra. Correspondence to: Qitian Wu <wuqitian @ @rrraddintttuutoorrrr
In this paper, aim to address generalization under topo- logical shifts through the lens of a physics-inspired graph learning model dubbed as Advective Diffusion Transformer (ADyDIFFORMER) The model is derived from advective diffusion equations which describe a class of continuous message passing process with observed and latent topo- logical structures. On top of this, we connect advective diffusion with graph Transformer architecture for gener-
Proceedings of the 4nnd International Conference on Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
!Codes available at qitiannn/AAvVDIFFono
com/

---

Supercharging Graph Transformers with Advective Diffusion
(Riemannian meiric).
Advective Diftusion Transformer
(Sii ((6))-5V((Ve Z(ri)
Now consider some quantity (e.g., temperature) func- tion of the form S + R, which we refer to as a scalar field. Similarly, we can define a (tangent) vector field Q: + TΩ, associating to every point u on a manifold a tangent vector Q(u) E T.S. which can be thought of as a local infinitesimal displacement. We use 2(22) and Q(T9) to denote the functional spaces of scalar and vector fields, respectively. The gradient operator V Q(2) + Q(T)) takes scalar fields into vector fields representing the local direction of the steepest change of the field. The diver- gence operator is the adjoint of the gradient and maps in the opposite direction, V : Q(TΩ) + @(22).
diffusion: global attentive propagation
BEE
t (t)
advection: local meesage passing
Z(O)
Z(T)
G A)
Genc
dee
observed space
latent space
output space
Figure I: Illustration of Advective Diffusion Transformers.
A manifold diffusion process models the evolution of a quantity (e.g., chemical concentration) due to its differ- ence across spatial locations on Ω. Denoting by q(u, t) Ω x [0, oo) - R the quantity over time l, the process is described by a PDE (diffusion equation) (Romeny, 2013):
alization against topologiaal shifts (Fig. 1): the non-local diffusion term (instantiated as global attention) aims to cap- ture latent interactions learned from data; the advection term (instantiated as local message passing) accommodates the topological features pertaining to observed graphs.
Ōq(u, t) 0t
= V (5(u, t) o Vg(u,t)), 2 0, u Q
To justify the model designs, we show that the closed-form solution of this advective diffusion system possesses the capability to control the generalization error caused by topo- logical shifts, which further guarantees the desired level of generalization. In contrast, commonly used local diffusion models that can be considered as a simplified variant of our model leads to the generalization error whose upper bound exponentially grows with topological shifts.
with initial conditions q(u, 0) = qo(u) and possibly ad- ditional boundary conditions if 2 has a boundary. $ de- notes the diffusivity of the domain. It is typical to dis- tinguish between an isotropic rncctiin-inndeendenn diffu- diffusivity
sivity), non-hoomgeneous doccation-depnneee S su) E RR,, and anisotropic (location- and direction- dependent S(u) E [²*=)) settings. In the cases studied be- low, we assume the dependence of diffusivity on locations is via a function of the quantity itself, i.c., S = S(g(u, t)).
For implementation, we resort to numerical scheme based on the Padé-Chebyshev theory (Golub & Van Loan, 1989) for efficiently computing the diffusion equation's is closed-form solution. Experiments show that our model offers supe- rior generalization performance across various downstream predictive tasks in diverse domains, including information networks, molecular screening, and protein interactions.
Diffusion on Graphs. Recent works adopt diffusion equa- tions as foundation principle for graph representation learn- ing (Chamberlain al., 30;la;;; Thorpe et al., 2022: Bod- nar et al., 2022; Choi et al., 2023; Rusch et al., 2023; Wu et al., 2024b), employing analogies between calculus on manifolds and graphs. Let G = (V,E) be a graph with nodes V and edges E, represented by the |V] x [V] ad- jacency matrix A. Let X = [xuluev denote a [V x D matrix of node features, analogous to scalar fields on manifolds. The graph gradient (VX) X, X2 de- fines edge features for (u, v) E E, analogous to vector fields manifolds. Similarly, the graph divergence of edge features E = [Sao](,,)e defined as the adjoint (V' E),, = > Le:u,,)e Sum, produces node features. Dif- fusion models replace discrete GNN layers with continu- ous time-evolving node embeddings Z(t) = [z.. (t)), where zu (t) [0, oo) + Ed driven by the diffusion equation:
2. Background and Preliminaries
We recapitulate diffusion equations manifolds (Freidlin & Wentzell, 1993; Medvedev, 2014) and their connection with graph learning.
Diffusion on Riemannian manifolds. Let Ω denote an abstract domain, which we assume here to be a Riemannian manifold (Eells & Sampson, 1964). A key feature distin- guishing an n-dimensional Riemannian manifold from Euclidean space is the fact that it is only locally Euclidean, in the sense that at every point E Ω one can construct n-dimensional Euclidean tangent space T.Ω z R" that locally models the structure of Ω. The collection of such spaces (referred to as the tangent bundle and denoted by TQ) further equipped with a smoothly-van varying inner product
az(t) V (S(Z(t): A) 3 VZ(t). t $2 0, (1)
with initial conditions Z(0) = dene XX) where Penc is a node-wise MI.P encoder and w.l.o.g., the diffusivity
2

---

Supercharging Graph Transformers with Advective Diffusion
S(Z()); A) over the graph can be defined as a [V] x [V] matrix-valued function dependent on A, which measures the rate of information flows between node pairs. With the graph gradient and divergence, Eqn. becomes
Ug
4 E
V
B
u u
X
az(() (C( -)))() OStST, (2)
Ot
with initial conditions Z(0)= dene (X) where C(Z(t);A)
is a [V] x [V] coupling matrix associated with the diffusivity. Eqn. 2 yields a dynamics from t = 0 to an arbitrary given stopping time T, where the latter yields node representations for prediction, e.g., Y = deec ((()). The coupling matrix determines the interactions between different nodes in the graph, and its common instantiations include normalized graph adjacency (non- par-maametrrr and learnable attention matrix (parametric), in which cases the mrite-differenc numerical iterations for solving Eqn. 2 correspond to the discrete propagation layers of common GNNs (Chamberlain et al., 2021a) Ia) and graph Transformers (Wu et al., 2023; 2024c) (see Appendix A for details).
Figure 2: The data-generating mechanism with topological shifts caused by environment E. The solid (resp. dashed) nodes represents observed (resp. latent) random variables.
erated from each U.. through a certain node-wise function X_ y(U,; W). We denote by matrix X particular real- ization of the random variable X.
ii Similarly, the graph adjacency A = [Aaol is a ran- dom variable generated through a pairwise function Aue = h(U,. Uo: W, E) additionally dependent on the environment E. The change of E happens when it transfers from train- ing to testing, resulting in different distribution of A. We denote by A particular realization of the adjacency matrix.
is typical to tacitly make a closed-world assumption, i.e., the graph topologies of training and testing data gener- ated from the same distribution. However, the challenge of generalization arises when the testing topology is different from the training one. In such an open-world regime, it still remains unclear how graph diffusion equations and, more broadly, learning- based models on graphs (e.g., GNNs) ex- trapolate and generalize to new unseen structures.
iii) The label Y can be specified in certain forms. As we assume in below, Y is generated by a function over sets, Y = r(((vvv)),, A; W). Denote by Y a realization of Y.
The above process formalizes the data- getta-genratingg mecha- nism behind various data of inter-dependent nature, where the graph data (X, A, Y) is generated from the joint dis- tribution p(X, A, Y E) with a specific environment. The learning problem boils down to finding parameters O of a parametric function T, that establishes the predic- tive mapping from observed node features and graph adjacency A to the label Y. I's is typically implemented as GNN, which is expected possess sufficient expressive power (in the sense that 30 such that I' X) & Y) as well as generalization capability under topological shifts (i.c., when the observed graph topology varies from training to testing, which in our model amounts to the change in E). While significant attention in the literature has been devoted to the former property (Morris et al., 2019; Xu et al., 2019; Bouritsas et al., 2023; Papp et al., 2021; Balcilar et al., 2021; Bodnar et al., 2022); the latter is largely open question.
3. Generalization by Advective Diffusion
As prerequisite for analyzing the generalization behaviors of learning-basee models on graphs, we need to charac- terize how topological shifts occur in nature. In general sense, extrapolation is impossible without any exposure to the new data or prior knowledge about the fata-gene mechanism. In work, we assume testing data is strictly unknown during training, in which case structural assump- tions become necessary for authorizing generalization.
3.1. Problem Formulation: Data Generation Hypothesis
We present the causal mechanism of graph data genera- tion in Fig. 2 as a hypothesis, inspired by the graph lim- its (Lovász & Szegedy, 2006; Medvedev, 2014) and random graph models (Snijders& & Nowicki, 1997). In graph theory. the topology of a graph g = (V.E) can be assumed to be generated by a graphon (or continuous graph limit), a ran- dom symmetric measurable function W [0, 1J² - [0, 1], which is an unobserved latent variable. In our work, we gen- eralize this data-generating mechanism to include alongside graph adjacency also node features and labels:
3.2. Proposed Model: Advective Diffusion Transformers
To deal with the problem formulated in the previous subsec- tion, we next present new diffusion equation model that of- fers a provable level of generalization in the data- teterggnnratiig mechanism as stipulated in Sec. 3.1. The model is inspired by a particular class of diffusion equations, called advec- rive diffusion, that describe common physical phenomenons driven by both diffusion and advection effects.
i) Each node u E V has a latent i.idd variable U, ~ 10, 1].
The node features are a random variable X = [X,] gen-
3

---

Supercharging Graph Transformers with Advective Diffusion
Advective Diffusion Equations. The classic advective dif- fusion is commonly used for characterizing physical sys- tems with convoluted quantity transfers, where the term advection refers to the evolution caused by the movement of the diffused quantity (Chandrasekhhar, 1943). Consider the abstract domain of our interest defined in Sec. 2, and assume (u,,4 1) E T (a vector field in S2) to denote the velocity of the particle location u and time t. The advec- tive diffusion of the physical quantity q on l is governed by Qq(u.)) 0t
where vy is learnable pairwise similarity function.
Advection as local message passing. The advection pro- cess driven by the directional movement belongs to an exter- nal force, with the velocity depending on contexts (c.g. dif- ferent rivers). This is analogous to the environment-sensitive graph topology that is informative for prediction in specific enviromments. We instantiate the velocity as the normal- ized graph adjacency, i.c., V D A///A/-1/2 reflecting observed structural information, where D is the diagonal degree matrix of A. Then our advective diffusion model can be formulated as:
the PDE as (Leveque, 1992):
4 t) t) O Vq(u, + AV (V(u,)) - q(u,t)) (3)
where 2 O, E Ω and 3 2 0 is a weight for the advection term. For example, if we consider q(u, E) as the water salin- ity in a river, then Eqn. 3 describes the temporal evolution of salinity at each location that equals to the spatial trans- fers of both diffusion process (caused by the concentration difference of salt and S reflects the molecular diffusivity in the water) and advection process (caused by the movement of the water and V characterizes the flowing directions).
az()) at
[C + 3V -VI] 0 OSIST,
(6)
with initial conditions Z(0) = ( where 3 E [0, 1] is a hyper-parameter The integration of non-local diffu- sion (implemented through attention) and advection (imple- mented as MPNNs) give rise to a new architecture, which we call Advective Diffusion Transformer (ADYDIFFORMER)
Remark. Eqn. has closed-form solution Z(t) (I-C--V>*Z(0) A special case of B = 0 (no advec- tion) can be used in situations where the graph structure is not useful. Moreover, one can extend Eqn. 6 to a non-linear equation with time-dependent C()( in which case the equation has no closed-fore solution and needs numerical schemes for solving. Similarly to (Di Giovanni et al., 2022), we found in our experiments a simple linear diffusion to be sufficient to yield promising performance. We therefore leave the study of the non-linear variant for the future.
Advective Diffusion on Graphs. Similarly, on graph g = EV,E), we can define the velocity for each node as a | -il-eneenio vector-valued function Vtt = [t)] We thus have (V)(V(I) z())), = Σnev Vno (a (1) (f) and the advective diffusion equation on graphs:
oz(t) (C ( I(),, 0S1ST. (4) a
We next instantiate the coupling matrix and the velocity to endow the model with generalizability under topological shifts, by drawing inspirations from physical phenomenons.
3.3. Theoretical Justification
Non-local diffusion as global attention. The diffusion process led by the concentration gradient acts as an internal driving force, where the diffusivity keeps invariant across environments (e.g., the molecular diffusivity stays constant in different rivers). With this intuition in mind, we consider the non-local diffusion operator allowing instantaneous in- formation flows among arbitrary locations (Chasseigne et al., 2006). In the context of learning on graphs, the non-local diffusion can be seen as generalizing the feature propagation to a complete or fully- fully-connected (latent) graph (Wu al., 2023; 2024c), in contrast with common GNNs that allow message passing only between neighboring nodes. Formally speaking, we can define the gradient and divergence opera- tors on a complete graph: (u) X, Xx (u, v E V) and EV*E), m Evev Gue (u E V). This resonates with the latent interactions among nodes, determined by the under- lying data manifold, that induce palr information flows over a complete graph and stay invariant w.r.t. the change of F. We thus instantiate C as global attention matrix that computes the similarities between arbitrary node pairs:
We proceed to analyze the generalization capability of our proposed model w.r.t. topological shifts as defined in See 3.1. We are interested in the generalization error of To instan- tiated as the continuous diffusion model in Eqn. 6, when transferring from training data generated with the environ- ment Err to testing data generated with Eie- The latter causes varied graph topologies as stipulated in See. 3.1.
We denote by {(X(), A Y(0)), the training data set sized Ner generated from p(X, A, YJE Err), and ((:,) any bounded loss function. The training error (i.e., empirical risk) can be defined as Remp Fer)
E No. ((To(X ( (, Y("). NY. 1-1
1
(7)
Our target is to reduce the generalization error on testing data generated from o(X,A,YI = Ere): R(To; Ete) &
ExX.A.Y)) ((To(X A'), Y'). (8)
1/(,,, (0), 2. (0)
Particularly, if E,e- = Eir, the learning setting degrades to the standard one commonly studied in the closed-world
C = [CurJu,EV, Cav = EwEV 77 m (0), Zu (0))
(5)
4

---

Supercharging Graph Transformers with Advective Diffusion
assumption, wherein the in-distribution generalization error has an upper bound (Shalev-Shwartz. & Ben-David, 2014):
condition since g establishes a mapping from a smooth and compact latent space high-dimensional space. Apply- ing Theorem 3.1 we have the generalization error of the advective diffusion model.
RToo: Far) - Remp (To; FAr) - Din (To, Far, Ner) (9)
= 2H( + + (Vlox(100//No)
Corollary 3.3. On the same condition of Theorem 3.1 and 3.2, the model-dependent generalization error bound of Eqn 6 can be reduced to arbitrary polynomial orders W.E. topological shifts, i.e., Dood-model (To, Eir, Eiy) = (+1)A)))))))
where ((T)) denotes the Rademacher complexity of the function class induced by Te and Din (T', Etr, Nor) is de- termined by training data size and model complexity.
When Ete o Ear that occurs in the open-world regime, i.e., our focused learning setting, the analysis becomes more dif- ficult due to the topologiall shifts. In the diffusion equation (either Eqn. 6 or 2), the change of graph topologies leads to the change of node representations (solution of the diffusion equation (T)). Thereby, the output of the diffusion pro- cess can be expressed as Z(T; A) fZZ()), A). Our first result below decouples the out-of-distributinn generalization gap RTT 8; Exe) - Remp oT: E into three error terms.
This implies that the generalization error of the model in Eqn. 6 can be controlled within arbitrary rate w.r.t. | The model has provable capacity for achieving a desired level of generalization with topological shifts, To verify the efficacy of the model, we consider syn- thetic data that simulates the topological shifts as defined in Sec. 3.1 and investigate into three types of topologi- cal shifts as shown in Fig. 3 (see experimental details in Sec. 5.1). We found that our model (ADYDIFFORMEL) and ADVDIFFORMER-S whose implementation is pre- sented in Sec. 4) keeps the testing error nearly constant as A increases.
Theorem 3.1. Assume 1 and Qdec are Lipschitz continu- ous. For any graph data generated with the mechanism of Sec. 3.1, it holds with the probability - ô that the general- ization gap of Tg satisfies
3.4. Comparisons with Other Models
(R(To; - Remp To; E1-] S Din Too Ner)
+ O(EA-M(AIEL,) AI-A)) A(T;A)) Z(T; 1)22)
To further illuminate the effectiveness of the proposed model, we next compare with two related models that are commonly adopted and can be considered as the simplified variants of our model.
Dond Rond-made Ecr.
Y| Y (1))
Doud Poud-tabell Fic)
Local Diffusion Model. We first consider typical model instantiation, i.e., local diffusion equation on graphs, wherein the model discards the advection term in Eqn. 6 and degrades to Eqn. 2 with the coupling matrix dependent on A. In such a situation, the propagation of node signals is constrained within connected neighbored nodes. The common choice for the coupling matrix is the symmetric normalized graph adjacency matrix À D 101 A22 1/2 In this case, the finite-difference iteration for solving Equ. 2 would induce the discrete propagation layers akin to the mes- sage passing rule of SGC (Wu et al., 2019) and GCN (Kipf & Welling, 2017) if the feature transformation and non- linearity are neglected (see more illustration in Appendix A). Given the constant coupling matrix C, Eqn. 2 has a closed- form solution Z(t)= e-I--C (0). However, unlike the advective diffusion model, the change rate of Z(T; A) w.r.t. A = À' - A produced by the local diffusion model has an exponential upper bound.
Remark. Since Din. is independent of the testing data gener- ated with Ete # Etr, the impact of topological shifts on the -iittib-iitiitttoonnnn generalization error is largely dependent on Dood-mode and Pood-tabet' the former reflects the vari- ation magnitude of ZTT A) yielded by Т3 w.r.t. varying topologies; the latter measures the difference of labels gener- ated with different environments. Notice that Dood-iobet is fully determined by the lata-gene mechanism, while Dood-mode is mainly dependent on the model I's, particu- larly the sensitivity of node representations w.r.t. topological shifts. We thus next zoom in on the specific design of I' as Eqn. 6, and the next result shows the upper bound of the change rate of (T; A) w.r.t. variation of graph topologies.
(1-
Theorem 3.2. For any graph data generated with the mech- anism of Sec. 3.1, if g is injective, then the model Eqn. 6 can reduce the variation magnitude of the node representation any order ((v(I AA [2)) where b denotes an arbitrary polynomial functiiooonnoonn AA- = À' and A = D AD -1/2
[I -')- A)
Proposition 3.4. For local diffussio model (defined by Eqn. 2) with the coupling matrix C = D 111/2AD 1/2 or C = D-1 A, the yielded node representation satisfies Z Z - Z(T A) = O(1 ||]l exp (1) I|
This suggests that the advective diffusion model is capable of controlling the change rate of node representations to arbitrary rates w.r.t ||/ The injectivity of g is mild
The label prediction Y = dueec (Z(;; A)) could be highly sensitive to the change of the graph topology. Pushing fur-
5

---

Supercharging Graph Transformers with Advective Diffusion
Dif7 Linear
Dift MuttiLayer
DIY Time
DKY NonLocal
AdvCIFRomee--
AdVDIFFarmer--
(a) Homophily Shift
(b) Density Shift
(c) Block Shift
Mean Square Eroo
0.06
0.55
Mean Square Eroo
0.35
0.50
0.05
Mean Squane a
0.45
0.30
0.04
0.40
0.25
0.35
0.20
0.02
0.15
20
40
60
70
80
100
110
120
20
80
100
||AA|2
||AA||
||AA||2
Figure 3: Testing errors (y-axix) w.r.t. differences in graph topologies (x-axis) on synthetic datasets that simulate the topological distribution shifts according to the data generation hypothesis of Fig. 2.
ther, we have the following corollary the generalization capability of local diffusion models under topological shifts.
graph. The corresponding diffuusio equation still exhibits the form of Eqn. 2 but allows non-zero entries for arbitrary (u, ' in the coupling matrix to accommodate the all- all-pair information flows. Then we can easily derive a result that if Y is conditionally independent from with given {U,) JEE in the data generation hypothesis of Sec. 3.1, the non-local diffusion model (i.e., Eqn. 2 with the attention- -based cou- pling matrix) leads to the generalization gap
Corollary 3.5. Under the same condition as in Theorem 3, for diffusion models Eqn. 2 with the normalized graph adja- cency as the coupling matrix, the model-dependent gen- eralization error on testing data generated with Ese + Etr has an upper bound: Dood modes (To, Eir, Eir) ((AA) Ap(A E..) E., |ll2 exp (Il (I)
RI 9; Eie) Remp (T':: Evr) S Din Err Ner), (10)
By definition in Sec. 3.1, the graph adjacency is realization of a random variable A huU,,, Un; W, E) dependent on varying environment E. Topological shifts caused by different distributions of A's between training and testing environments may result in large Dood-model- 2 This result together with Theorem 3.1 suggests that local diffusion models may lead to undesirably poor generalization in cases where models are expected to be insensitive to the change of topologies. For example, for situations where the ground- truth labels do not dramatically change with topological shifts (i.e., D Pood- is small), local diffusion models may induce large Dood-mode, that prejudices generalization. The above conclusion can be extended to models with layer-wise feature transformations and non-linearity (see Appendix B.6 for illustration). While the upper bound result does not mean the exponentially growing error would definitely happen in practice, our empirical comparison in Fig. 3 shows that the generalization error of local diffusion models (Diff-Linear, Miff-MulriLayer and Diff-Time) on test data grows super- linearly w.r.t. | ll2 across three types of topological shifts (see experimental details in Sec. 5.1).
which holds with the probability - 6. The assumption of conditional independence between Y and A, however, can be violated in many situations where labels strongly corre- late with observed graph structures. Furthermore, the per- formance on testing data (i.e., what we care about) depends on both the model's expressiveness and generalization. The non-local diffusion alone, discarding any observed topology, has insufficient expressiveness for capturing the structural information. By contrast, the advective diffusion model pro- posed in Sec. 3.2 that accommodates the observed structures can provably generalize under topological shifts without the conditional independence assumption between Y and A. This makes the advective diffusion model more powerful in real cases. Furthermore, as empirically validated in Fig. 3, the non-local diffusion model (Diff- (offfNooLooca) indeed yiclds comparably stable yet obviously inferior performance to our models (ADYDIFFORMER- and ADYDIFFORMER-S5
4. Model Implementation
The remaining question concerning model implementation boils down to how to solve the advective diffusion equa- tion Eqn. 6. One straightforward solution to harness the scheme adopted by (Chen al., 2018) for back prek--rpppggtion through PDE dynamics. However, since it is known that the equation has closed-form solution e cI---C-3V) resort to a implementation-wue ing the solution instead of numerically solving the equation. Nevertheless, direct computation of the matrix exponential through eigendecomposition is computationally intractable for large matrices. As an alternative, we leverage numer-
Non-Local Difftusion Model. We next discuss the non- local diffusion model without the advection term, which as previously mentioned can be essentially treated as gener- alization of local diffusion to a complete or fully-connected
"The influence of topology variation is inherently associated with h. For example, if one considers as the stochastic block model (Snijders & Nowicki, 1997), then the change of E may lead to generated graph data with different edge probabilities. In the case of real-world data with intricate topological patterns, the functional forms of h can be more complex, consequently inducing different types of topologcaal shifts.
simpler method by comput-
6

---

Supercharging Graph Transformers with Advective Diffusion
ical techniques based on series expansion that produces two model versions. Due space limit, present the main ideas in this subsection and defer details on model implementation to Appendix D.I.
Table I: Results on Arxiv and Twitch, where we use time and spatial contexts for data splits, respectively. We report the Accuracy (1) for three testing sets of Arxiv v and average ROC-AUC (t) for all testing graphs of Twitch (results for each case are reported in Appendix F.1). Top
ADYDIFFORMER-I uses a numerical method based on the extension of Padé-Chebyshev theory to rational frac- tions (Golub& & Van Loan, 1989; Gallopoulos & Saad, 1992). which has shown empirical success in 3D shape analy- sis (Patané, 2014). The matrix exponential is approximated by solving multiple linear systems (see more details and derivations in Appendix C) and we generalize it as a multi- head network where each head propagates in parallel:
performing methods are marked as first/sccond/hird. 00M indicates out-of-memory error.
Arii (2018)
Arxiv (2019)
Arsiv (2020)
Twitch (arg)
MLP
49.91 : 0.59
47,30 :: 0.63
46.78 0.98
61.12 : 0.16
GCN
50,14 + 0.46
48.06
1.13
46,46 + 0.85
59.76 + 0.34
GAT
51,60 :: 0.43
48.60 1: 0.28
46.50 : 0.21
59.14 1: 0.72
SGC
51,40 0..00
49.15
0.16
46.94
+ 0.29
60.86 + 0.13
GDC
51,53 -
49.02
0.10
GRAND
52.45 =: 0.27
50.18
0.18
48.01 1: 0.24
61.65 4: 0.23
L (1 00I C₁ SV. h = -,,,,, ,H, Z(T) ~ Σ h=1 oFd (Linsolver 2(0))).
A-DGNs
50,91 1: 0.41
47.54
60.11 0.09
CDE
50.54 2: 0.21
47.31 2 0.52
45.32 1: 0.26
60.69
1: 0.10
(11)
GraphTrans
00M
0OM
00M
61.65 0.23
GraphGPS
$1.11 = 0.19
48.91 + 0.34
46.46 + 0.95
62.13
+ 0.34
DIFFormer
50.45 0.94
47.37
1.58
44.30 + 2.02
62.11
0.11
ADVDIF FORMER-S
53.41 0.48
51.53
0.60
49.64 +: 0.54
62.51
0.07
The linsolver computes the matrix inverse Z (Lx))-1(()) and can be efficiently implemented via torch. Barch.linalg.oon ferentiation. Each head contains propagation with the pre- computed attention C₁2 and node-wise transformation of FC
() that enables automated dif-
in Sec. 3.1. We instantiate h as a stochastic block model which generates edges Aue according to block numbers (b), intra-block edge probability (p1) and inter-block edge probability (p2). Then we study three types of topological distribution shifts: homophily shift (changing P2 with fixed Pı); density shift (changing P1 and p2); and block shift (varying b). The predictive task is node regression. More details on data generation are presented in Appendix E.I.I.
ADYDIFFORMER-S resorts to approximation by finite ge- ometric series (see Appendix C for derivations):
Р C₂ + AA h- 1,.., H,
(h) Z(T) ~ Σ h=1 FC (0). P,Z()) ... (Ph) (()))
(12)
Fig. 3 plots the testing error (i.e., Mean Square Er- ror) w.r.t. differences in graph topologies || ||2 (i.e., the gap between training and testing graphs) in three cases. We compare our model (ADYDIFFORMER-I and ADVDIFFORMER S) with other diffusion-basse models as competitors. The latter includes Diff-Linear (graph diffu- sion with constant C), Diff-MultiLayer (the extension of Diff-Linear with intermediate feature transformations), Difff- Time (graph diffusion with time-dependent C(Z(t)) and Diff-NonLocal (non-local diffusion with the global attention- based C(Z())). The results show that three local graph diffusion models exhibit clear performance degradation, i.c., the regression error grows super-linearly w.r.t. topological shifts, while our two models yield consistently low error across environments. In contrast, the non-local diffusion model produces comparably stable performance yet inferior to our models due to its ignorance of the useful information in input graphs. These empirical observations are consistent with our theoretical results presented in Sec. 3.3 and 3.4.
This model aggregates K-order propagated results with the propagation matrix PA, in each head. One advantage of this model version lies in its good scalability with linear complexity w.r.t. the number of nodes in the feed-forward computation (see detailed illustration in Appendix D.1.2).
5. Experiments
We apply our model to wide variety of downstream tasks of disparate scales and granularities that involve topolog- ical shifts led by distinct factors. Due to the diversity of datasets and tasks, the competing models that are applicable to specific cases can vary case by case, so the goal of our experiments is to showcase the wide applicability and supe- riority of ADVDIFFORMER against commonly used GNNs and graph Transformers well as bespoke methods tailored for specific tasks. In the following, we delve into each case separately with the overview of experimental setup and dis- cussions. More detailed dataset information is provided in Appendix 1. Details on basclines and are deferred to Appendix E.2 and E.3, respectively.
5.2. Real-World Datasets
We next evaluate ADVDIFFORMER on real-world datasets with more complex topological shifts concerning non- Euclidean data in a diverse set of applications.
5.1. Synthetic Datasets
To validate our model and theoretical analysis, we create synthetic datasets simulating the data generation hypothesis
Information Networks. We first consider citation net- works Arxiv (Hu et al., 2020) and social networks
7

---

Supercharging Graph Transformers with Advective Diffusion
Table 2: Results of RMSE ()) for node regression and edge regression on dynamic networks of protein-protein interactions.
Node Regression Test Average
Edge Regression Test Average
Valid
Test Worst
Valid
Test Worst
MLP
0.768 0.011
0.672 0.014
0.768 0.014
0.150 0.004
0.192 0.003
0.204 0.003
GCN
1.791 + 0.023
1.308 0.013
1.797 + 0.007
0.185 + 0.003
0.196 + 0.001
0.213 0.001
GAT
1.255 1: 0.022
1.057 + 0.030
1.708 0.067
0.210 + 0.010
0.204 1 0.006
0.216 0.010
SGC
1.622 0.004
1.154 0.006
1.616 + 0.002
0.193 0.000
0.191 + 0.001
0.209 0.001
GraphTrans
3.798 + 1.146
3,203 + 0.889
3.795 + 1.123
0.189 + 0.005
0.189 = 0.008
0.202 = 0.003
GraphGPS
0.713 0.050
0.671 0.040
0.803 0.081
0.168 0.004
0.182 0.007
0.216 0.019
DIFFormer
0.672 + 0.046
0.637 0.034
0.710 + 0.028
0.171 + 0.007
0.183 - 0.005
0.197 0.003
ADYDIFFORMER-I
0 = 0.010
0.643 = 0.019
0.679 + 0.021
0.159 + 0.002
0.166 + 0.006
0.184
0.011
ADYDIFFORMER-S
0.547 0.040
0.574 0.028
0.644 0.040
0.156 0.006
0.167 0.004
0.188 0.010
Ground Truth
AdvDIFFormer (0.697)
GCN (0.685)
GAT (0.664)
GraphGPS (0.694)
DIFFormer (0.674)
Figure 4: Testing cases for molecular mapping operators generated by different models with averaged testing Accuracy (1)
reported. The task is to generate subgraph-level partitions (marked by different colors) resembling the ground-truth.
Twitch (Rozemberczki et al., 2021) with graph sizes rang- ing from 2K to 0.2M, where we use the scalable version ADYDIFFORMER- To introduce topological shifts, we partition the data according to publication years and geo- graphic information for Arxiv and Twitch respectively. The predictive task is node classification, and we follow the common practice comparing Accuracy (resp. ROC- AUC) for Arxiv (resp. Twitch). We compare with three types of state-of-the-ar baselines: (i) classical GNNs (GCN (Kipf& & Welling, 2017), GAT (Velickovic et al., 2018) and SGC (Wu et al., 2019)); (ii) diffusion-based GNNs (GDC (Klicpera et al., 2019), GRAND (Chamberlain et al., 2021a), (Gravina et al., 2023) and CDE (Zhao et al., 2023)), and (iii) graph Transformers (GraphTrans (Wu et al.. 2021), GraphGPS (Rampásek et al., 2022), and the diffusion-based DIFFormer (Wu et al., 2023)). Ap- pendix E.2 presents detailed descriptions for these models. Table reports the results, showing that our model offers significantly superior generalization for node classification.
eralization of the model in different scenarios, we consider two predictive tasks: i) node regreesion for gene expression values (measured by RMSE), and 2) edge regression for predicting the co-expression correlation coefficients (mea- sured by RMSE). Table reports the averaged scores and worst-case scores across all testing graphs. The results show that ADYDIFFORMER-S and ADYDIFFORMER- are ranked in the top three (resp. two) models in terms of the averaged (resp. worst-case) testing performance. More- over, ADYDIFFORMER performs better in node regres- sion tasks, while ADYDIFFORMER-1 exhibits (slightly) better competitiveness for edge regression. The possible reason might be that ADVDIFFORMER- can better exploit high-order structural information as the matrix inverse can be treated ADY with K co.
Molecular Mapping Operator Generation. We next con- sider the generation of molecular coarse-grained mapping operators, an important step for molecular dynamics sim- ulation, aiming to find a representation of how atoms are grouped in molecule (Li et al., 2020). The task is graph segmentation problem which can be modeled as predicting edges that indicate where to partition the graph. We use the relative molecular mass to split the data and test how the model extrapolates to larger molecules. Fig. 4 compares the testing cases (with more cases shown in Appendix F.I) gen- erated by different models, which shows the more accurate estimation of our model (we use ADVDIFFORMER for experiments) that demonstrates desired generalization.
Protein Interactions. We then test on protein-protein inter- actions (Fu & He, 2022). Each node denotes a protein with time-aware gene expression value and the edges indicate co-expressed protein pairs at each time. The dataset con- sists of 12 dynamic networks each of which is obtained by one protein identification method and records the metabolic cycles of yeast cells. The networks have distinct topological features (e.g., distribution of cliques) (Fu & He, 2022), and 6/1/5 networks for train/valid/test. To test the gen-
8

---

Supercharging Graph Transformers with Advective Diffusion
Arxiv
DDPIN-nr
DDPIN-er
0.55
Accuracy
0.50
1.54
RMSE
0.45
Test (2018)
1.52
0.145
(2019)
0.40
Test (2020)
1.48
1.46
0.140
1.0 1.5 2.0"
1.44
0.8
0.5
1.0
B
B
Figure 5: Analysis of B on Arxiv and node regression (nr) and edge regression (er) tasks on DPP TN.
Hlpee-Paarameter Analysis. The hyper-parameter 3 con- trols the importance weight for the advection term. Fig. 5 shows the model performance of ADVDIFFORMER Arxiv and DPPIN with different B's. We found that the optimal settings for 3 can be different across datasets and tasks. For node classification Arxiv, the model gives the best performance with f E [0.7, 1.0]. The performance degrades when 3 is too small (<0.5) or too large (>~2.0). The reason could be that the graph structural information useful for the predictive task on Arxiv yet too much emphasis on the graph structure can lead undesired gen- eralization. Differently, for DPPIN, found that using smaller can bring up satisfactory performance across node regression and edge regression tasks. In particular, set- ting B = 0, in which case the advection term is completely dropped, can yield optimal performance for the node re- gression task. This is possibly because the graph structure uninformative and pure global attention can learn gen- eralizable topological patterns from latent interactions. To sum up. in practice, the model enables much flexibility for adjusting the weight on the advection effect (the importance of observed structural information) to accommodate the diversity of graph-structured data. More hyper-parane parameter analysis (w.r.t. 0 and K) is deferred Appendix F.2.
models. In general sense, improving the generalizability of the model is important for many aspects associated with Al's societal responsibilities, such as addressing the obser- vational bias in training data and promoting the fairness of the outcome the test set. Specifically speaking, paper explores the generalization capability of neural net- work models operated on varying graph topologies, which has implications on several significant applications such as social network analysis, drug discovery and healthcare. The results and methodology presented in this work can shed lights on enhancing the trustworthiness and reliability of machine learning models in the open world regime.
on
A Acknowledgement
QW is supported in part by funding from the Eric and Wendy Schmidt Center at the Broad Institute MIT and Harvard. MB is partially supported by the EPSRC Turing Al World- Leading Research Fellowship No. EP/X040062/1 and EP- SRC AI Hub No. EP/Y028872/1.
References
Atwood, J. and Towsley, D. Difftssion-conoluttinall neural networks. In Advances in Neural Information Processing Systems, pp. 1993-2001, 2016.
Ablation Studies. We defer ablation studies on the diffusion and advection terms of our model to Appendix F.2.
Balcilar, M., Renton, G., Héroux, P., Gaüzère, B., Adam, S., and Honeine, P. Analyzing the expressive power of graph neural networks in a spectral perspective. In International Conference on Learning Representations, 2021.
6. Conclusions
This paper studies generalization with topological shifts, a largely open question in machine learning, and the in- sights in this work open new possibilities of leveraging diffusion PDEs as principled guidance for navigating gener- alizable neural network models. As exemplified in this work, proposed Advective Diffussin Transformer, inspired by advective diffusion equations, has provable potentials for generalization and shows superior performance in various downstream tasks across different scenarios.
Bazhenov, G., Kuznedelev, D., Malinin, A., Babenko, A., and Prokhorenkova, L. Evaluating robustness and un- certainty of graph models under structural distributional shifts. arXiv preprint arXiv:2302.13877 2023.
Bodnar, C., Di Giovanni, F., Chamberlain, B., Liò, P., and Bronstein, M. Neural sheaf diffusion: A topological perspective on heterophily and oversmoothing in gnns. Advances in Neural Information Processing Systems, 35: 18527-18541, 2022.
Impact Statement
This paper presents work whose goal is to advance the cur- rent understandings for the generalization of neural network
Bouritsas, G., Frasca, F., Zafeiriou, S., and Bronstein, M. M. Improving graph neural network expressivity via sub-

---

Supercharging Graph Transformers with Advective Diffusion
graph isomorphism counting. IEEE Trans. Pattern Anal. Mach. Intell., 45(1):657-668, 2023.
Gilmer, J., Schoenholz, S. S., Riley, P. F., Vinyals, O., and Dahl, G. E. Neural message passing for quantum chem- istry. In International Conference on Machine Learning, pp. 1263-1272 2017.
Chamberlain, B., Rowbottom, J., Gorinova, M. I., Bron- stein, M. M., Webb, S., and Rossi, E. GRAND: graph neural diffusion. In International Conference on Machine Learning (ICML), pp. 1407-1418, 2021a.
Golub, G. H. and Van Loan, C. F. Matrix computations. John Hopkins University Press, 1989.
Chamberlain, B. P., Rowbottom, J., Eynard, D., Giovanni, F. D., Dong, X., and Bronstein, M. M. Beltrami flow and neural diffusion on graphs. In Advances in Neural Information Processing Systems (NeurIPS), 2021b.
Gravina, A., Bacciu, D., and Gallicchio, C. Anti-symmetric DGN: stable architecture for deep graph networks. In International Conference on Learning Representations, 2023.
Chandrasekhar, S. Stochastic problems in physics and as- tronomy. Reviews of modern physics, 15(1):1, 1943.
Hornik, K., Stinchcombe, M., and White, H. Multilayer feedforward networks are universal approximators. Neu- ral networks, 2(5):359-366, 1989.
Chasseigne, E., Chaves, M., and Rossi, J. D. Asymptotic be- havior for nonlocal diffusion equations. Journal de math- ématiques pures appliquées, 86(3):271-22) 2006.
Hu, W., Fey, M., Zitnik, M., Dong, Y., Ren, H., Liu, B., Catasta, M., and Leskovec, J. Open graph benchmark: Datasets for machine learning on graphs. In Advances in Neural Information Processing Systems, 2020.
Chen, R. T., Rubanova, Y., Bettencourt, J., and Duvenaud, D. K. Neural ordinary differential equations. In Advances in neural information processing systems, 2018.
Hu, W., Fey, M., Ren, H., Nakata, M., Dong, Y. and Leskovec, J. Ogb-Ise: A large-scale challenge for machine learning on graphs. arXiv preprint arXiv:2 2021.
Choi, J., Hong. S., Park, N., and Cho, S.-B. Gread: Graph neural reaction-diffusion equations. In International Con- ference on Machine Learning, 2023.
Kipf, T. N. and Welling, M. semi-supervised classifica- tion with graph convolutional networks. International Conference on Learning Representations (ICLR), 2017.
Choromanski, K. M., Likhosherstov, V., Dohan, D., Song, X., Gane, A., Sarlós, T., Hawkins, P., Davis, J. Q., Mo- hiuddin, A., Kaiser, L.., Belanger, D. B., Colwell, I.. J., and Weller, A. Rethinking attention with performers. In International Conference Learning Representations, 2021.
Klicpera, J., Weißenberger, S., and Günnemann, S. Dif- fusion improves graph learning. In Advances in neural information processing systems, 2019.
Di Giovanni, F., Rowbottom, J., Chamberlain, B. P., Markovich, T., and Bronstein, M. M. Understanding convolution on graphs via energies. Transactions on Ma- chine Learning Research, 2022.
Koh, P. W., Sagawa, S., Marklund, H., Xie, S. M., Zhang, M., Balsubramani, A., Hu, W., Yasunaga, M., Phillips, Guo, R. L.. Gao, I., Lee, T., David, E., Stavness, W., Earnshaw, B., Haque, I., Beery, S. M., Leskovec, J., Kundaje, A., Pierson, E., Levine, S., Finn, C., and Liang, P. WILDS: A benchmark of in-the-wild distri- bution shifts. In International Conference on Machine Learning (ICML), pp. 5637-5664, 2021.
Dwivedi, V. P. and Bresson, X. A generalization of trans- former networks to graphs. CoRR, abs/2012.09699, 2020.
Eells, J. and Sampson, J. H. Harmonic mappings of rieman- nian manifolds. American journal of mathematics, 86(1): 109-160, 1964.
Leveque, R. J. Numerical methods for conservation laws, volume 214. Springer, 1992.
Freidlin, M. and Wentzell, A. D. Diffusion processes on graphs and the averaging principle. The Annals of probability, pp. 2215-2248 1993.
Li, Z., Wellawatte, G. P., Chakraborty, M., Gandhi, H. A., Xu, C., and White, A. D. Graph neural network based coarse-grained mapping prediction. Chemical science, 11 (35):9524-9531, 2020.
Fu, D. and He, J. Dppin: A biological repository of dynamic protein- interaction network data. In 2022 IEEE International Conference on Big Data (Big Data), pp. 5269-5277. IEEE, 2022.
Lovász, L.. and Szegedy, B. Limits of dense graph sequences. Journal of Combinatorial Theory, Series B, 96(6):933 957, 2006.
Gallopoulos, E. and Saad, Y. Efficient solution of parabolic equations by krylov approximation methods. SIAM jour- nal on scientific and statistical computing, 13(5):1236 1264, 1992.
Ma, J., Deng, J., and Mei, Q. Subgroup generalization and fairness of graph neural networks. Advances in Neural Information Processing Systems, 2021.
10

---

Supercharging Graph Transformers with Advective Diffusion
Medvedev, G. S. The nonlinear heat equation on dense graphs and graph limits, SIAM Journal on Mathematical Analysis, 46(4):2743-2766, 2014.
Topping, J., Giovanni, F. D., Chamberlain, B. P., Dong, X., and Bronstein, M. M. Understanding over-squashing and bottlenecks on graphs via curvature. In International Conference on Learning Representations, 2022.
Morris, C., Ritzert, M., Fey, M., Hamilton, W. L., Lenssen, E.,, Rattan, G., and Grohe, M. Weisfeiler and leman go neural: Higher-order graph neural networks. AAAI Conference Artificial Intelligence, pp. 4602-4609, 2019.
Van Loan, C. The sensitivity of the matrix exponential. SIAM Journal Numerical Analysis, 14(6):971-981, 1977.
Vaswani, A., Shazeer, N.. Parmar, N., Uszkoreit. J., Jones, I.., Gomez, A. N., Kaiser, I.., and Polosukhin, I. Atten- tion is all you need. In Advances in Neural Information Processing Systems, pp. 5998-6008, 6008, 2017.
Papp, P. A., Martinkus, K., Faber, L., and Wattenhofer, R. Dropgnn: Random dropouts increase the expressiveness of graph neural networks. In Advances in Neural Infor- mation Processing Systems, pp. 21997-22009, 2021.
Velickovic, P., Cucurull, G., Casanova, A., Romero, A., Liò, P., and Bengio, Y. Graph attention networks. In International Conference on Learning Representations (ICLR), 2018.
Patané, G. Laplacian spectral distances and kernels on 3d shapes. Pattern Recognition Letters, 47:102 110, 10, 2014.
Rampásek, L., Galkin, M., Dwivedi, V. P., Luu, A. T. Wolf, G., and Beaini, D. Recipe for a general, powerful, scal- able graph transformer. In Advances in Neural Informa- tion Processing Systems, pp. 5998-6008, 2022.
Wu, F., Jr., A. H. S., Zhang, T., Fifty, C., Yu, T., and Wein- berger, K. Q. Simplifying graph convolutional networks. International Conference on Machine Learning, pp. 6861-6871, 2019.
Romeny, B. M. H. Geometry-driven diffusion in computer vision, volume 1. Springer Science & Business Media, 2013.
Wu, Q. Zhang, H., Yan, J., and Wipf, D. Handling distri- bution shifts on graphs: An invariance perspective. In International Conference on Learning Representations, 2022a.
Rossi, E., Kenlay, H., Gorinova, M. I., Chamberlain, B. P., Dong, X., and Bronstein, M. M. On the unreasonable effectiveness of feature propagation in learning on graphs with missing node features. In Learning Graphs Con- ference, pp. 11-1. PMLR, 2022.
Wu, Q., Zhao, W., Li, Z., Wipf, D., and Yan, Nodeformer: A scalable graph structure learning transformer for node classification. In Advances in Neural Information Pro- cessing Systems, 2022b.
Rozemberczki, B., Allen, C., and Sarkar, R. Multi-scale at- tributed node embedding. Journal of Complex Networks, 9(2). 2021.
Wu, Q., Yang. C.,, Zhao W., He, Y., Wipf, D., and Yan, J. Difformer: Scalable (graph) transformers induced by energy constrained diffusion. In International Conference on Learning Representations, 2023.
Rusch, T. K., Chamberlain, B. P., Mahoney, M. W., Bron- stein, M. M., and Mishra, S. Gradient gating for deep multi-rate learning on graphs. In International Confer- ence on Learning Representations, 2023.
Wu, Q. Nie, F., Yang, C., Bao, T., and Yan, J. Graph out- of-distribution generalization via causal intervention. In The Web Conference, pp. 350-860, 2024a.
Scarselli, F., Gori, M., Tsoi, A. C., Hagenbuchner, M., and Monfardini, G. The graph neural network model. IEEE transactions on neural networks, 20(1):61-80, 2008.
Wu, Q., Nie, F., Yang, C., and Yan, J. Learning divergence fields for shift-robust graph representations. In Interna- tional Conference on Machine Learning, 2024b.
Shalev-Shwartz, S. and Ben-David, Understanding ma- chine learning: From theory algorithms. Cambridge university press, 2014.
Wu, Q., Wipf, D., and Yan, J. Neural message passing induced by energy-constrained diffusion. arXiv preprint arXiv:2409.09) 2024c.
Snijders, T. A. and Nowicki, K. Estimation and prediction for stochastic blockmodels for graphs with latent block structure. Journal of classification, 14(1):75-100, 1997.
Wu, Z.. Jain, P., Wright, M. A., Mirhoseini, A., Gonzalez, J. F.., and Stoica, Representing long-range context for graph neural networks with global attention. In Advances in Neural Information Processing Systems, 2021.
Thorpe, M., Xia, H., Nguyen, T., Strohmer, T., Bertozzi, A. L.,, Osher, S. J., and Wang, B. GRAND++: graph neural diffusion with a source term. In International Conference on Learning Representations (ICLR), 2022.
Xu, K., Hu, W., Leskovec, J., and Jegelka, S. How powerful are graph neural networks? In International Conference on Learning Representations, 2019.

---

Supercharging Graph Transformers with Advective Diffusion
Xu, K., Li, J., Zhang, M., Du, S. S., Kawarabayashi, K., and Jegelka, S. How neural networks extrapolate: From feedforward to graph neural networks. In International Conference on Learning Representations (ICLR), 2021.
Yang, C., Wu, Q. Wang, J., and Yan, J. Graph neural networks inherently good generalizers: Insights by bridging gnns and mlps. In International Conference Learning Representations, 2023.
Yang, C., Wu, Q. Wipf, D., Sun, and Yan, J. How graph neural networks learn: Lessons from training dynamics. In International Conference Machine Learning, 2024.
Ying, C., Cai, T., Luo, S., Zheng, S., Ke, G., He, D., Shen, Y., and Liu, T. Do transformers really perform bad for graph representation? In Advances Neural Information Processing Systems, 2021.
Zhang, X., Wang, L., Helwig, J., Luo, Y., Fu, C.,, Xie, Y., Liu, M., Lin, Y., Xu, Z., Yan, K.. et al. Artificial intelli- gence for science in quantum, atomistic, and continuum systems. arXiv preprint arXiv:230..08423, 2023.
Zhao, K., Kang, Q. Song, Y., She, R., Wang, S., and Tay, W. P. Graph neural manvection-diffusioon with heterophily. In International Joint Conference on Artificial Intelli- gence, 2023.
12

---

Supercharging Graph Transformers with Advective Diffusion
A. Connection between Diffusion Equations and Message Passing
In this section, we provide a systematically introduction on the fundamental connections between graph diffusion equations and neural message passing, as supplementary technical background for our analysis and methodology presented in the main text. Consider graph diffusion equations of the generic form
( (C(Z(t): A) - ())(t), O with initial conditions Z(0) Penc (X).
(13)
As demonstrated by existing works, e.g., (Chamberlain et al., 2021a), using finitediifference numerical schemes for solving Eqn. 13 would induce the message passing neural networks of various forms. The latter is recognized as the common paradigm in modern graph neural networks and Transformers whose layer-wise updating aggregates the embeddings of other nodes compute the embeddings for the next layer.
A.1. Graph Neural Networks as Local Diffusion
Consider the explicit Euler's scheme as the commonly used finite- -nite-difference method for approximately solving the differential equations, and Eqn. 13 will induce the discrete iterations with step size T:
(++1) Z (k)
(C(Z R3 A) (1)
(14)
With some re-arranging we have
Z (--1) (1- (1-7) + +
(15)
with the initial states ( Deno (X.. The above updating equation gives one-layer update through residual connection and propagation with (*) A). There are some well-known graph neural network architectures that can be derived with different instantiations of the coupling matrix.
Simplifying Graph Convolution (SGC). If one considers ((Z(+);
À D -1/2AAD 1/2, then we will get the
one-layer updating rule:
z (++1) (1 ()) + DD 102 AD (1)
(16)
This can be seen as one-layer propagation of SGC (Wu al., 2019) with residual connection, and when T 1 becomes exactly the SGC layer. Since SGC model does not involve feature transformation layers and non-linearity throughout the message passing, one often uses a pre-computed propagation matrix for one-step convolution that is much faster than the multi-layer convolution:
(kk) -p= P
A/2 1/2
(17)
Graph Convolution Networks (GCN). The GCN network inserts feature transformation layers in-between the propagation layers. This can be achieved by considering K stacked piece-wise diffusion equations, where the k-th dynamics is given by the differential equation with time boundaries:
(K) az k) (C- E txt), with initial conditions Z(ta-1ik) Oins (Z(ta-lik= 1)), (18) Øt
where oie denotes the node-wise feature transformation of the k-th layer. Assume C = 1-2 A1112 Then consider
one-step seed-forward of the explicit Euler scheme for Eqn. 18, and one can obtain the updating rule at the k-th layer:
z (++1)
(k+1) -
(11-))Z()) + 1/2 AD 1/)2(())
(19)
This corresponds one GCN layer (Kipf & Welling, 2017) if one considers (k+1) as fully-connected neural layer with ReLU activation and simply sets 7 = 1.
High-Order Propagation. Besides the explicit numerical scheme, one can also utilize the implicit scheme and multi-step schemes (c.g., Hunge-KKutall for solving the diffusion equation, and the induced updating form will involve high-order information (Chamberlain et al., 2021a).
13

---

Supercharging Graph Transformers with Advective Diffusion
A.2. Graph Transformers Non-Local Diffusion
The original architectures of Transformers (Vaswani al., 2017) involve self-attention layers as the key module, where the attention measures the pairwise influence between arbitrary token pairs in the input. There recent works, e.g., (Dwivedi & Bresson, 2020; Ying et al., 2021; Wu et al., 2021; Rampásek al., 2022; Wu et al., 2022b) transferring the Transformer architectures originally designed for sequence inputs into graph-structured data, and the attention is computed for arbitrary node pairs in the graph, which can be seen as a counterpart of non-local diffusion (Wu et al., 2023; 2024c). In specific, the coupling matrix allows non-zero entries for arbitrary location pairs and can be instantiated as a global attention network. Then using the explicit Euler's scheme as Eqn. 15 we can obtain the self-attention propagation layer of common Transformers:
(k)
Z ( (1 () (
(k)
MZZ
(20)
Σ EEEV (/ (k)
(k)
For obtaining the fully-connected layers and non-linear activations adopted in Transformers, one can inherit the spirit of GCN and extend the diffusion model to K piece-wise equations as Eqn. 18.
B. Proofs for Technical Results
B.1. Proof for Theorem 3.1
According to the data generation hypothesis in Fig. 2, for given node latents U, 's, we can decompose the joint distribution into (we omit all conditions on U for brevity)
p(X, = PIIIEE)I-
(21)
Also, by definition in Sec. 3.1 we have
pX = Ev) = p Ete),
(22)
PY E E(r)
Ete).
(23)
Therefore we have - P We next consider the gap between R(Ts; o Evr) and R(T Ete):
IR(To: RT E.r)
EXXAAAYY", A Y)) MCA. Y)]
BX (XL' AA), Y')
YIJI
Ex B. A'),
EX EXX-XXAAAA--AAAAEEEAAAAAA-------
A), Y'Il
(24)
EX (B
( A), Y')]
+ of -x)
A), YIJ
< Ex wr -
+Bx=p(X) Anppp R= (o__ (
[I(To(X,A), -(T
Moreover, due the Lipschitz continuity of and Ôdee, we have
P | ZT; A/) - A))l2
(25)
[XT, A), Y') I(T A), Y) < L2 Y Y
(26)
where L1 and L2 denote the Lipschitz constants. Combing Eqn. 25 and Eqn. 26 with Eqn. 24, we have
RR( Ere) R(T';; Exr) 1 < L1 EA-
Z(T; ((TAA) L2 (AAAYAAAAAAAA
(27)
Y' -Y
The conclusion for the main theorem can obtained via combining Eqn. and Eqn. 9 using the triangle inequality.
14

---

Supercharging Graph Transformers with Advective Diffusion
B.2. Proof for Theorem 3.2
For the advective diffusion equation with the coupling matrix C pre-computed by attention network 77... (0), and fixed velocity V -D-1/2AD 1/2, we have its closed-fon solution
(1)) = - -1 1 2 0.
(28)
To prove the perturbation bound w.r.t. the change of graph structures, we introduce the following lemma. Lemma B.1. Ler X, E E Tgax,, and ler -
11 be submuliplicative matrix norm. Suppose there exist constants M 2 1, 2 O such that for all Y E Ruxx,, l|e x S Me x11, Then the following perturbation bound holds:
||e lex+E__ x
E M (((111)-
(29)
Proof. Define path X(s) == X + sE for E [0,
1]. Then:
x+F _ ex = (.¹
_ ds.
x(s)
x()
(30)
ds
Using the integral form of the derivative of the matrix exponential (Van Loan, 1977), have:
d x(x) ( (1-0)x(()
Ee ex(s) d0.
Be
(31)
ds
Therefore:
c X+E x= S/' (1o' (((-))x(0) De exis) db ds.
(32)
Taking norms and applying submultiplicatiiity:
(/ 1le |((1-1)x(()] | (x() d0ds.
||e X+x+ ex 11 |eex+B-e*11e
(33)
Using the growth bound assumption:
[e(l-))i(o S (1= |( S (6l
(34)
Multiplying these we have:
|e [100))((()|| (X])| s ~
(35)
Note that |((s)|
X
sEl
X
E SO:
M². f S/ cull Elddddd < E 2 ([I
x
E
(36)
The existence of M and w is guaranteed for every consistent matrix norm ( (Van Loan, 1977) such spectral norm || ||2 considered in our analysis.
L -I-C 3V and -I-C" BV. We then apply the above lemma that holds even if L and LT are not commutable:
Le- 112 S MºT IL' Lll2 Ti @TTIIL' Lill2
(37)
For spectral norm || ||2 in our case, the above result holds for M = 1 and w =
We next prove the bound for the last term of Eqn. 37 by construction. Notice that the initial states are given by the encoder MLP: (()) = Pene(X) According to our data generation hypothesis in Fig. 2, we know that node embeddings are generated from the latents of each node (we use u, to denote the realization of U.), i.c., X2 = g(u,; W) and the graph adjacency is generated through a pair-wil wise function auv h(u,, Up; W, F.). Since g is injective, we assume g as its inverse mapping.
15

---

Supercharging Graph Transformers with Advective Diffusion
We define by v dene the function composition of v and Pene that establishes a mapping from input node features X to the nntention-based coupling C. According the universal approximation results that hold for MLPs the compact set (Hornik al., 1989), can construct mapping induced by dene to obtain propagation matrix in the form of C = C - (,+ where C is independent from A and > 0 is an arbitrary small number. To be specific, the construction of the mapping can be achieved by 1) o Dene = m a o
g maps the input feature X to Uu:
h maps (Uw: Up) to Gaei
m maps Our to Curs where Cup denotes the (u, v)-th entry of C.
Then consider the difference of node representations under topological shifts and we have II(C' + + (C+ = ((1AAl)2). Since || is bounded, for any positive integer m, there exists > such that exp ||A (2) S |A mh Therefore, we have the conclusion:
eIIL' - L1l2 = = < O(1AA
(38)
The theorem can be concluded by combining the result of Eqn. 37.
B.3. Proof for Corollary 3.3
The conclusion follows by combing the results of Theorem 3.1 and Theorem 3.2.
B.4. Proof for Proposition 3.4
The diffusion equation with the constant coupling matrix has a closed-form solution Z(t) = (Z))) 0. To prove the proposition, we need to derive the bound of -eelee--"" (I-C'IT e- (I-CJT |l2 for any C' + C. According the result (3.5) of (Van Loan, 1977) we have
| 11--'-' C'JT - 112 CT Cllalle (I (L-CTT
12e I2ell(c
(39)
Given the fact C' C = A' À = AA, we have
|ee -1-c'' - - 0(14A12 [|
(40)
This gives rise to the conclusion that
A') -(T; O(IAA|12 exp(II 2))),
(41)
and we conclude the proof for the proposition.
B.5. Proof for Corollary 3.5
By combing the results Theorem 3.1 and Proposition 3.4, we have
Dood model (To, Etr, Ere) = O A(EA A A'PpA E.c) [(T;A) Z(T; A)ll)
(42)
O (BAA F.r) AA exp [A]))T))).
B.6. Extension with Feature Transformations
The conclusion of Proposition 3.4 and Corollary 3.5 can be extended to the cases incorporating feature transformations and non-linear activation in-between propagation layers used in common GNNs, like GCN (Kipf & Welling, 2017). In particular, the diffusion model becomes the piece-wise diffusion equations with K dynamics components defined by Eqn. 18:
(Z k)) = (C- I))(t;),, E (tt- tkl, with initial conditions Z(t+-1;k) = o (Z((( 1)), (43)
Øt
16

---

Supercharging Graph Transformers with Advective Diffusion
where Pint denotes the nose feature transformation of the k-th layer. Based on this, can re-use the reasoning line of proofs for Proposition 3.4 each component, and arrive at the exponential bound of node representation within the k-th dynamics:
IZ (12; A)) Z(4,; A, A)12 O ||] exp ([AA|la((=-(+<1)))) By stacking the results for each component, one can obtain the variation magnitude of the node representation yielded by the whole trajectory
-
(44)
(T; - Z(T; = 0 exp (127
(45)
C. Approximation Strategies for Diffusion PDE Solutions
The closed-form solutions of linear diffusion equations often involve the form of matrix exponential e e-Tx, which is intractable for computing its exact value. There are many established techniques based on numerical approximations, e.g., series expansion, in this fundamental challenge. In our presented model in Sec. 4, we propose two implementation versions based on two approximation ways for handling the closed-form solution of the advective diffusion equations on graphs.
Approximation with Linear Systems. One scalable scheme proposed by (Gallopoulos & Saad, 1992) is via the extension of the minimax Padé- Sade-Chebyshev theory to rational fractions (Golub & Van Loan, 1989). This approximation technique has been utilized by (Patané, 2014) as effective and efficient method for spectrum-free computation of the diffusion distances in 3D shape analysis. In specifie, the matrix exponential of the form Lt is approximated by the combination of multiple matrix inverses:
exp (-Lt) ~ Σ aL + +011)
(46)
where a and O2 can be pre-defined parameters (Gallopoulos & Saad, 1992). To unleash the capacity of neural networks, in Sec. 4, our model implementation (ADYDIFFORMER-) extends this scheme multi-head network where each head contributes to propagation with independently parameterized attention networks. The matrix inverse is computed with the linear system solver that is available in common deep learning tools (e.g., PyTorch) and supports automatic differentiation.
Approximation with Geometric Series. When the graph sizes become large, the matrix inverse can be computationally expensive. For better scalability, we can use the geometric series for approximation:
(L+ 1.0 E(-1) (1t2)) L
(
(47)
k=
In this way, the matrix exponential can be approximately computed via a combination of finite series:
cxp =
0
(k+1) L*
(48)
In our model, the closed-form solution for the PDE induces L (I C V)), and the summation in Eqn. 48 can be expressed as weighted sum of _ = (C+ + Vv)* for k = 0,... K. Our model implementation (ADYDIFFORMER-S)
proposed in Sec. 4 generalizes the weighted sum to a one-layer neural network.
D. Model Implementations and Algorithms
In this section, we provide detailed and self-contained descriptions about our model architectures in Appendix D. Then in Appendix D.2. we discuss how to apply our model to various graph-structured data with additional input information. To make the presentation clear and focused on the model implementation side, we will re-define some notations that are originally defined in Sec. 4, where we formulate the model with the terminology of the PDE domain.
D.1. Model Architectures
The model takes graph G- = (V,E, E X, A) input, and output prediction in the downstream tasks. We assume the number of nodes in the graph | = N, node feature matrix X E x and graph adjacency matrix A E {0, 1] 1 We use D to denote the diagonal degree matrix of A. The normalized adjacency is denoted by À = D- N d-dimensional column vector. In this subsection, we assume & has no edge weight or edge feature for presentation, and with loss of generality, we will discuss how to incorporate these additional attributes in Appendix D.2.
1/2 and 1 is an all-one
17

---

Supercharging Graph Transformers with Advective Diffusion
1.1.1. INSTANTIATIONS AND PARAMETERIZATTONS
Our model is comprised of three modules: the encoder denes the decoder Odecs and the propagation network in-between the first two.
Encoder:
The node features X = =[Xu]uEv E RNxD are first mapped to embeddings in the latent space Z = [zu EEV
RNxd via the encoder: ()) = Denc The encoder (enc(-) is instantiated as shallow MLP with non-linear activation ReLU).
Propagation:
The propagation network converts the initial node embeddings ( to the node representations Z Iuulev E Nxx (where (() and Z are the re-defined counterparts of (() and Z respectively, presented in Sec. 4).
The propagation network is implemented via multi-head network with H/ heads involving the attention network (.)))
and feature transformation network @ (-). The latter is instantiated as fully-connected layer Wo. and the attention
network is instantiated as normalized dot- not-product positive similarity function:
W Wo. (0)
W K. (0)
17 Iz Z
(b) (0) (0))
= 1
W Q.1.i. 1
[ K. 112
(0)
(49)
(2 (1) 211
C₁
(c(h))
(h) Cov
Ewev (1) Zu (0) (0))
where Wo. E Rdddd and WK, k. E Rddd are trainable weights for query and key, respectively, of the l-thh head. Then the node representations will be computed in different ways by two models.
For ADVDIFFORMER-1. the node representations are calculated via
L (1+ 0)1 C₁ - -BA, h h= 1.... H
Z (nnolver(Ln,... ) h 1. H
(50)
H
= Zn Wo.
h=1
where Wo, Raxd is trainable weight matrix. Alg. summarizes the feed-forward computation of ADYDIFFORMER-L.
For ADYDIFFORMER-S, the node representations computed by
P4 C₁ + AA 1, 1,...,
1 (k)
P
(k-1) Pr.z(-11) 1, K. h h=1,...,
(51)
H
z
(1) ; z(K /
h
k=1
where Wo, E R +1)ddx is a trainable weight matrix. To accelerate the computation Eqn. 51, we can inherit the strategy used in (Wu al., 2023) and alter the order of matrix products, which reduces the time and space complexity O(N) (see Appendix D.1.2 for detailed illustration). Alg. 2 presents the feed-forward computation of ADYDIFFORMER that only requires (NN) algorithmic complexity.
The decoder (.)e(.) transforms the node representations into prediction. Depending on the specific downstream tasks, the decoder can be implemented in different ways:
Decoder:
(node-level prediction): d.. = MLP(z.,)
(graph-level prediction): ý = M.L.PSSmPoolligg (zulwev)
(52)
(edge-level prediction): tuv MLP([Z.. ze]).
In particular, the softmax activation is used for output in classification tasks. For training, we adopt standard loss functions, i.e. cross-entropy for classification and mean square loss for regression.
18

---

Supercharging Graph Transformers with Advective Diffusion
Algorithm 1 Feed-Forward of the Model ADYDIFFORMER-I.
INPUT: Node feature matrix X and normalized adjacency matrix Å. ( dene (X)
for h = 1,
H do
Qq A =
Wo.
ZKA
W.. (9)
II 1 WEV
Wx,z 112 NEV
U 11 T Z...
C₁ diag (U,1) Un
LA (1 + 0 S BA
Z linsolver(L,, Z)
Z = Σ l==l Z Wo,r
OUTPUT: Node representations 2 and predicted labels with ddee (Z).
Algorithm 2 Feed-Forward of the Model ADYDIFFORMER-S (with O(N) complexity).
INPUT: Node feature matrix X and normalized adjacency matrix A.
(0) dene X)
for h H do
W Wo.
W K. 1 Wx,N
ZQ
ZK.
[Wa.aa uEV
MEV
N diag
(N + ZQ ((ZK,)
/ (0)
for k 1,... K do
z (k) -Na (11T2((-1) Zer (ZKK))/((111)
Z₁ [/h Z (1)
0)
(x)]
= >4 ZA W 0,
OUTPUT: Node representations Z and predicted labels with Ødec Pdec(Z). (Z).
D. ACCELERATION ADYDIFFORMER-S WITH LINEAR COMPLEXITY
We illustrate how achieve the propagation of ADYDIFFORMER-S in Eqn. 51 with O(N) complexity. With the query and key matrices defined by Zq,4
(o, (0)
2 (0) the attention matrix C2. in Eqn. 49 is
and ZK,
W Q.
[Wa.
dEV
computed by (in the matrix form used for implementation)
(N ZQ (ZK,)) 1) (11T + Zon (Zk,)")
C₁ diag
(53)
Computing the above result requires O(N2) time and space complexity. Still, if consider the feature propagation with C₁, we have
C diag -1 (N + Zq (ZK,) 1) [11' + ZZ. (ZK,b) 2//11)
(k)
(54)
diag (N Zq ((ZK,)) (T))
(1(Tz. x+) + ZQ.h (ZK,,)) (.4)
where the equality is achieved by altering the order of matrix products. The above computation only requires O(N) time and space complexity. The forward computation of ADYDIFFORMER-S with O(N) acceleration summarized in Alg. 2.
D.2. Applicability of Our Model
In the main paper, we assume unweighted graphs without edge attribute features for model formulation. Without loss of generality, we next discuss how to extend our model to handle the edge weights and edge features.
Edge Weights. For weighted graphs, the adjacency matrix A would become real matrix where the entry due denotes
19

---

Supercharging Graph Transformers with Advective Diffusion
the weight the edge (u, E &. In this situation, still have the corresponding normalized adjacency Á D A Å A--1-12------ 1/2, where D diag((du)uev) and d,, ".,(,,,)e Our- Our model implementations can be trivially generalized to this case by using A as the propagation matrix for local message passing.
Edge Features. If the graph contains edge features, denoted by E [e,o](u,o)e E IREE we introduce an encoding layer WE DR. for mapping the edge features into embeddings in the latent space and then incorporate them with node embeddings. In specific, first compute the edge-to-node signals:
M
mmllvv
m. Σ
Au, W egew-
(55)
es(MN)EE
For ADVDIFFORMER-1, we can modify Eqn. 50 as
L, (1 + 0JI - C₁ -3A,
/ linsolver
[L,, (Z (0) M))
(56)
H
Z ZE Wo.
For ADYDIFFORMER-S, we can modify Eqn. 51 be
Ps = Cn + 3A,
(k)) (Z (k-1) + M), h: 1,... K,
(57)
z Σ (Z) z(1), .... Z ( O.h. h=
E. Experiment Details
We supplement details for our experiments, regarding datasets, competitors, and implementations, for facilitating the reproducibility.
E.1. Datasets
The datasets we use for the experiments in Sec. 5 span diverse domains and learning tasks. We summarize the statistics and brief descriptions for each dataset in Table 3, with the detailed information presented in the following subsections.
Table 3: Statistics and descriptions for experimental datasets.
Dataset
#Nodes
#Edges
#Graphs
Train/Val/Test Split
Task
Metric
Synthetic-h
1,000
14,064 32,066
SRM (Homophily)
Node Regression
RMSF
Synthetic-d
1,000
7,785 13,912
SBM (Density)
Node Regression
RMSE
Synthetic-b
1,000
14,073 59.936
SBM (Block Number)
Node Regression
RMSE
Twitch
1,912 9,498
31,299 153,138
7
Geographic Domain Publication Time
Node Classification
ROC-AUC
Arxiv
169,343
1,166,243
Node Classification
Accuracy
OGB-BACE
10-97
10- 101
1.513
Molecular Scaffold
Graph Classification
ROC-AUC
DGB-SIDER
1-4
0-505
1,427
Molecular Scaffold
Graph Classification
ROC-AUC
DPPIN-nr
140-3,003
25,95,
Protein Identification Method
Node Regression
RMSE
DPPIN-cr
143 1,0-3
22-25.924
Protein Identification Method
Edge Regression
RMSE
DPPIN-IP
143 5,003
25.92.
12
Protein Identification Method
Link Prediction
ROC-AUC
HAM
25
7-
1.987
Relative Molecular Mass
Edge Classification
Accuracy
L. SYNTHETIC DATASETS
The synthetic datasets used in Sec. 5.1 simulate the graph data generation in Sec. 3.1, where the topological distribution shifts are caused by the difference of environments across training and testing data. In specific, we generate graphs of
20

---

Supercharging Graph Transformers with Advective Diffusion
|
1000 nodes, with the node features X, graph adjacency matrix A and labels Y generated by the following process.
Each node u E V is assigned with a scalar Uy randomly sampled from the uniform distribution U(00, 1].
For the generation of node features X = [xu]uEv, we instantiate the node-wise function g as a 2-layer MLP with ReLU activation and 4-dimensional output. Then the node feature X₂ is generated through X,, = ML.P(u)).
For the generation of graph adjacency A [ouu]u,ee, we instantiate the pairwise function h as the stochastic block model (Snijders & Nowicki, 1997) which generates edges according to the intra-block edge probability (pi) and the inter-ble block edge probability (p2). We map the nodes into b blocks by the following rule: for node u E V, we assign k) (where 1 < k < b). Then the edge aue is randomly generated from a bernoulli
it to the k-th block if Un. E [ =5) distribution with p1 if a and v are in the same block, and p2 otherwise.
For the generation of labels Y, we consider the regression tasks and each node has a label yo generated through an ensemble model of 2-layer GCN and l-layer DIFFormer (without using the graph-based propagation) with random initializations: Y gen U, A) difformer( U, A), where U = [uu]uEv.
Using the above data generation, we create graphs with the indices #1~ #12, and use the graph #1 for training, the graph #2 for validation, and the graphs #3~ #12 for testing. The topologiaal distribution shifts are introduced in three different ways as described in Sec. 5.1, where in each case, the detailed configurations for pt, P2 and are illustrated below.
Homophily Shift: p1 0.1, b 5 and p2 0.01 + 0.05 12 (i - 1) for the graph #i.
Density Shift: b = 5, 0.1 + 0.1 * 12 (i 1) and p2 = 0.01 + 0.1 * 12 (i 1) for the graph #i.
Block Shift: p1 = 0.1, P2 = 0.01 and b (i 1) for the graph #i.
1.2. INFORMATION NETWORKS
The citation network Arxiv provided by (Hu et al., 2020) consists of single graph with 0.16M nodes, where each node represents paper with the publication year (ranging from 1960 to 2020) and a subarea id (from 40 different subareas in total). The node attribute features are 128- 228-dimensional obtained by averaging the word embeddings of the paper's title and abstract. The edges are given by the citation relationship between papers. The predictive task is to estimate the paper's subarea. We use the publication years to split the data: papers published before 2014 for training, within the range from 2014 to 2017 for validation, and on 2018/2019/2022 for testing. Since there is a single graph, to increase the difficulty of generalization, we consider the inductive setting: the testing nodes are not contained in the training graph. Table demonstrates the dissimilar statistics for aaaninn/validationnessinn graphs, manifesting the existence of topological shifts. Following the common practice, we use Accuracy as the evaluation metric.
Table 4: Statistics for
graphs on Arxi ivv. v. There is single citation network that augments with time evolving, and with the data splits in the inductive setting, the previous graph is contained by the subsequent one.
Train (1960-2014)
Valid (2015-2017)
Test 1 (2018)
Test 2 (2019)
Test 3 (2020)
Target Nodes
41,125
49,816
24.799
39,711
8,892
# All Nodes
41.125
90,941
120,740
160.451
169,343
# All Edges
102,316
374,839
622,466
1,061,197
1,166,243
Max Degrees
275
3,036
6,251
12,006
13.161
Avg Degrees
4.98
8.24
10.31
13.23
13.77
Twitch (Rozemberezki et al., 2021) is comprised of seven dis-connected graphs, where each node represents a Twitch user and edges indicate the friendship. Each graph is collected from the social newtork in a particular region, including DE, ENGB, ES, FR, PTBR, RU and TW. The node features are multi-hot with 2,545 dimensions indicating the user's profile. The predictive task is to classify the gender of the user. The seven networks with sizes ranging from 2K to 9K have distinct structural characteristics (such as densities and maximum degrees) as observed by (Wu et al., 2022a). We therefore split the data according to the geographic information: use the network DE for training, ENGB for validation, and the remaining networks for testing. The evaluation metric is ROC-AUC for binary classification.
21

---

Supercharging Graph Transformers with Advective Diffusion
E.1. BIOLOGICAL PROTEIN INTERACTIONS
DPPIN (Fu & He, 2022) contains 12 individual dynamic network datasets at different scales, and each dataset a dynamic mooeiin-protei interaction network that describes the protein- interactions of yeast cells. Each graph dataset is obtained by one protein identification method and consists of 36 graph snapshots, wherein each node denotes protein that has sequence of 1-dimensional continuous features with 36 time stamps. This records the evolution of gene expression values within metabolic cycles of yeast cells. The edges in the graph are determined by co-expressed protein pairs at one time, and each edge is associated with a co-expression correlation coefficient.
We consider the predictive tasks within each graph snapshot and ignore the temporal evolution between different snapshots. In specific, we use the graph topology of each snapshot as the observed graph adjacency A and use the gene expression values at the previous 10 time steps as node features X. On top of this, we consider three different predictive tasks: 1) node regression for gene expression value at the current time (measured by RMSE); 2) edge regression for predicting the co-expression correlation coefficient (measured by RMSE); 3) link prediction for identifying co-expressed protein pairs (measured by ROC-AUC). Given the fact that each graph dataset has distinct sizes (ranging from 143 to 5,003 nodes) and distributions of 3-cliques and 4-cliques (ranging from 0 to hundreds) (Fu & He, 2022), we consider the dataset-level data which introduces topological distribution shifts.
splitting and use 6/1/5 graph datasets for
E. MOLECULAR MAPPING OPERATOR GENERATION
The Human Annotated Mappings (HAM) dataset (Li et al., 2020) consists of 1,206 molecules with expert annotated mapping operators, i.e., a representation of how atoms are grouped in molecule. The latter segments the atoms of a molecule into groups of varying sizes. As an important step in molecular dynamics simulation, generating coarse-graing mapping operators aims to reproduce the mapping operators produced by experts. This task can be modeled as a graph segmentation problem (Li et al., 2020) which takes a molccule graph as imput and outputs the labels for each edge that indicates there is cut needed to partition the source and end atoms into different groups.
For data splits, we calculate the relative molecular mass of each molecule using the RDKit package", and rank the molecules with increasing mass. Then we use the first 70% molecules for training, the following 15% for validation, and the remaining for testing. This splitting protocol partitions molecules with different weights, and requires generalization from small molecules in the training set to larger molecules in the testing set.
Table 5: The range of relative molecular mass for training/validaaionnessttng molecules in HAM.
Train
Valid
Test
Relative Molecular Mass
108. 18 ~ 273.34
273.34 ~ 311.14
311.14- ~ 762.94
E.2. Competitors
In our experiments, we compare with peer encoder backbones for graph learning tasks. The competitors span three aspects: 1) classical GNNs, 2) Liffusion-based GNNs, and 3) graph Transformers. We briefly introduce the competitors and illuminate their connections with our model.
GCN (Kipf & Welling, 2017) is a popular model that propagates node embeddings over observed graphs for com- puting node representations, which can be seen as the discretized version of graph diffusion equations with feature transformations.
GAT (Velickovic et al., 2018) introduces attention networks for computing pairwise weights for neighboring nodes in the graph and propagates node signals with adaptive strengths given by the attention weights. GAT can be seen as the discretized version of the graph diffusion equation with time-dependent coupling matrices.
SGC (Wu et al., 2019) proposes to simplify the GCN architecture by removing the feature transformations in-between propagation layers, reducing multi-layer propagation to one-layer. This brings up significant acceleration for training and inference. SGC can be seen as the discretization of the linear diffusion equation on graphs
/https://gt thub. com rrkit//i///.... ttkii..
22

---

Supercharging Graph Transformers with Advective Diffusion
GDC (Kliepera et al., 2019) extends the graph convolution operator to graph diffusion convolution derived from the linear diffusion equation on graphs. We use its implementation version based the heat kernel for diffusion coefficients.
GRAND (Chamberlain et al., 2021a) proposes graph neural diffusion, a continuous PDE model, that generalizes manifold diffusion to graphs and then uses numerical schemes to solve the PDE. We compare with its linear version that implements the linear graph diffusion equation.
A-DGNs (Gravina et al., 2023) is a stable graph neural architecture inspired by ODE on graphs that has provable capability to preserve Long-range information between nodes and avoid gradient vanishing or explosion in training.
CDE (Zhao et al., 2023) is recently proposed continuous model derived from convection diffusion equations that is designed for addressing heterophilic graphs
GraphTram eWu al., 2021) is a recently proposed Transformer for graph-structured data that satisfies the permutation- invariant property. The model architecture sequentially combines GNNs and Transformers in order, where the GNN can learn local, short-range structures and the Transformer can capture global, long-range relationships.
GraphGPS (Rampásek et al., 2022) introduces a scalable and powerful Transformer model class for graph data and achieves state-of-the-an results on molecular property prediction benchmarks. We use its scalable implementation version with the Performer attentions (Choromanski et al., 2021).
DIFFormer (Wu et al., 2023) is a scalable Transformer inspired by diffusion on graphs. The model comprised of principled attention layers, which implements the diffusion iterations minimizing a global energy. The architecture integrates graph-based feature propagation and global attention each layer. We use its version with simple diffusivity that only requires linear complexity and yields state-of-the-art results some large-graph benchmarks.
E.3. Implementation Details
Computation Systems. All the experiments are run NVIDIA 3090 with 24GB memory. The environment is based on Ubuntu 18.04.6. Cuda 11.6, Pytorch 1.13.0 and Pytorch Geometric 2.1.0.
Evaluation Protocol. For all the experiments, we run the training and evaluation of each model with five independent trials, and report the mean and standard deviation results our tables and figures. In each run, we train the model with fixed budget of epochs and record the testing performance produced by the epoch where the model yields the best performance on validation data.
We use the grid search for hyper-parameter tuning on the validation dataset with the searching space described below.
For information networks, hidden size d {32, 64, 128), learning rate E {0.0001, 0.001), head number H (1,2,4), the weight for local message passing 3 E {0.2, 0.5, 0.8, 1.0], and the order of propagation (only used for ADYDIFFORMER-S. K 41,
For molecular datasets, hidden size d 256, learning rate E (0.01, 0.005, 0.001, 0.0005, 0.0001, 0.00005), dropout E {0.0, 0.1, 0.3, 0.5). head number H 21, 4), the weight for local message passing / {0.5, 0.75, 1.0). the coefficient for identity matrix (only used for ADVDIFFORMER-I) 0 {0.5, 1.0), and the order of propagation (only used for ADYDIFFORMER-S. K E 2,42,
For protein interaction networks, hidden size E {32, 64), learning rate E {0.01, 0.001, 0.0001), head number H e {1, 2, 4], the weight for local message passing 3 E {0.3, 0.5, 0.8, 1.0), the coefficient for identity matrix (only used for ADDDIFFORMER-I) 0 E {0.5, 1.0], and the order of propagation (only used for ADYDIFFORMER-S) K E (1,2,3,4,
F. Additional Experimental Results
In this section, we supplement more experimental results including additional results for main experiments, ablation studies and hyper- parameerr analysis.
23

---

Supercharging Graph Transformers with Advective Diffusion
Arxiv
DDPIN-nr
DDPIN-er
0.165
1.65
0.54
1.60
Accuracy
1.55
0.155
1.50
0.150
0.48
1.45
0.145
0.46
0.140
0.44
Test 2 (2019)
1.35
Test 3 (2020)
0.42
0.130
2
Figure 6: Model performance on Arxiv and DPPIN with different settings of K. The latter involves node regression (nr) and edge regression (er) tasks.
F.1 Supplementary Results for Main Experiments
In Table 6, we present the ROC-AUC for each graph of Twitch. Fig. 7 and 8, we show the generated results for more testing cases of molecular mapping operators in HAM.
Table 6: Result of ROC-AUC for each graph Twitch where nodes in different networks to split the training, validation and testing data.
Train (DE)
Valid (ENGB)
Test I (ES)
2 (FR)
Test (PTBR)
Test (RU)
T (TW)
MLP
75.26 + 1.49
63.48 + 0.15
65.19 0.37
62.25 0.28
65.01 + 0.19
$4,92 + 0.33
58.23 4:
GCN
69.55 +: 0.34
60.76 + 0.21
63.75 0.44
61.56 = 0.56
63.26 0.42
54.51 #021
5
GAT
0:00:0000
+:
SGC
71.68 +:
+
65.12 0.15
0.12
64.14 0.19
56.83
GDC
80.73 1.69
62.14 0.30
66.33
0.51
64.21
$6.60 0.24
58.97
GRAND
79.17 0.74
62.48
0.22
61.62 0.62
64.44
5.33
59.27
A-DGNs
78.91 0.52
61.52 : 0.34
65.82
60.59 0.56
63.49
5..74
5.53
CDE
80.21 0.35
62.51 0.21
65.62
60.93 0.55
63.92
55.79 0.31
58.42 0.42
GraphTrans
79.17 + 0.74
62.48 + 0.39
66.52 0.23
61.62 0.62
64.44 0.73
56.42 0.38
59,27 0.57
GraphGPS
74.49 + 1.35
63.40 + 0.31
66.85
0.32
63.74 0.37
65.03 + 0.58
56.39 0.39
58.63 0.83
DIFFormer
73.12 + 0.52
63.06 1 0.09
66.68 1: 0.15
64.44 0.13
65.23 0.20
55.75 0.12
58.91 0.37
ADYDIFFORMER-S
75.46 +: 0.28
63.53 + 0.14
66.78 +: 0.14
63,35 0.10
65.68 +: 0.06
56.27 0.06
60.48 + 0.21
E.2. Ablation Studies and Hyper-Parameter Analaysis
We next conduct more analysis on proposed model by ablation studies some key components and investigating the impact of hyper-parameters.
Diffusion and Advection. We conduct ablation studies on the advection term (i.e., the local message passing) and the diffusion term (i.e., the global attention). In Table we report the results for ADYDIFFORMER-S on Arxiv, which shows that the two modules are indeed effective for producing superior generalization on node classification tasks.
Table 7: Ablation studies for ADYDIFFORMERR-S Arxiv.
Train (1960-2014)
Valid (2015-2017)
Test (2018)
Test (2019)
Test (2020)
ADYDIFFORMER AIDFIFFRRMER diffusion
63.79 + 0.66
55.25 + 0.14
53.41 0.48
51.53 + 0.60
49 64 +: 0.54
64.65 1.10
55.00 + 0.12
52.45 0.27
50.18
48.01 +
DIFFOORRR
54.31
0.59
49.65
Impact of K. The pypeer-arrameter K (used for ADVDIFFORMER-S. controls the number of propagation orders in the model. In fact, the value of K would impact how the structural information is utilized by the model. If K is small, the model only utilizes the low-order structure, and large K enables the usage of high-order structural information. Fig. 6 presents the model performance on Arxiv and DPP IN with K ranging from to 6. We observe that the optimal settings for K are different across cases, and using larger can not necessarily yield better performance. This is because these cases, the low-order structural information is informative enough for desired generalization.
24

---

Supercharging Graph Transformers with Advective Diffusion
Impact of 0. Finally, study the impact of 0 used for computing L, in ADVDIFFORMER-I. Table 8 shows the performance of ADYDIFFORMR DPPIN with different 0's. We found that using 0 close to bring up stably good performance, which is consistently manifested by experiments on other cases. Still, if 0 is too small, e.g., close to 0, it would sometimes lead to numerical instability. This is due to that, in such case, the matrix L₁ could become singular matrix.
Table 8: Testing accuracy of ADVDIFFORMER-I with different 0's in the edge regression task DPPIN.
0
0.5
1.0
2.0
Accuracy 0.241
0.154
0.147
0.149
G. Current Limitations and Future Works
The generalization analysis in the current work focuses the data- generating mechanism described in Fig. 2 which is inspired and generalized by the random graph model. While this mechanism can in principle reflect real-world data generation process in various graph-structured data, in the open-world regime, there could exist situations involving topological distribution shifts by diverse factors or their combination. Future works can extend our framework for such cases where inter-dependent data is generated with different causal mechanisms. Another future research direction lies in the instantiation of the diffusion and advection operators in our model. Besides our choice of MPNN architecture to implement the advection process, other possibilities include structural and positional embeddings. We leave this line of exploration for the future, along with the analysis for the generalization capabilities of more general (e.g., non-linear) versions of the advective diffusion equation and other architectural choices.
25

---

Supercharging Graph Transformers with Advective Diffusion
Ground Truth
AdvilIFFommer (1.00)
GCN (1.00)
GAT (0.69)
GraphGPS (0.92)
Difformer (0.92)
Ground Truth
AdyDIFFormer (1.00)
GCN (0.22)
GAT (0.22)
GraphGPS (0.22)
Diftormer (0.22)
(b)
Ground Truth
AdvDlFForrme (1.00)
GCN (1.00)
GAT (0.89)
GraphGPS (0.86)
Difformer (0.86)
(c)
Groun Truth
AdvDIF-Former (0.90)
(0.54)
(0.87)
GraphGPS (0.74)
Difformer (0.74)
(d)
Ground Truth
AdvDIFFormer (0.83)
GCN (0.70)
GAT (0.70)
GraphGPS (0.63)
Difformer (0.60)
Figure 7: Additional testing cases for molecular mapping operators generated by different models and the expert annotations (ground-truth) For each case, we report the score (the higher is better) that measures the closeness between the generated results and the expert annotations.
26

---

Supercharging Graph Transformers with Advective Diffusion
Ground Truth
AdvIIFFoomer (1.00)
GCN (0.87)
GAT (0.77)
GraphGPS (0.85)
Difformer (0.87)
Ground Truth
AdyDIFFormer (1.00)
GCN (0.84)
GAT (0.76)
GraphGPS (0.70)
Diftormer (0.64)
b
Ground Truth
AdvDIFFormer (1.00)
GCN (0.72)
GAT (0.74)
GraphGPS (0.56)
Difformer (0.57)
Ground Truth
AdvDIFFormer (1.00)
GCN (0.77)
GAT (0.77)
GraphGPS (0.77)
Difformer (0.89)
Ground Truth
AdvDIFFormer (1.00)
GCN (0.64)
GAT (0.51)
GraphGPS (0.51)
Difformer (0.51)
Figure 8: Additional testing cases for molecular mapping operators generated by different models and the expert annotations (ground-truth). For each case, we report the score (the higher is better) that measures the closeness between the generated results and the expert annotations.
27


Reviews:
[{"review_id": "Qz0m3skHFV", "reviewer": "Reviewer_ZJpC", "summary": "This paper proposes advective diffusion transformer, a model that provably controls generalization error with topological shifts. The model has been evaluated on synthetic and several real-world datasets that verify its superiority compared with existing local and non-local graph neural networks/transformers.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your time in reviewing our paper and for the constructive feedback.\n\n**Q1: \"How is the performance of the proposed approach on the OGB benchmarks?\"**\n\nWe appreciate the suggestion to include additional datasets. We added results on OGBL-COLLAB, a link prediction task suited for evaluating models under topological shifts. We follow the temporal splits from [2], which naturally introduce distribution shifts. We found that our model (ADiT-Series) achieves substantial improvements over GNN/Transformer baselines:\n\n| | GCN  | GAT | GraphGPS | DIFFormer | ADiT-Series\n| - | - | - | - | - | - |\n| **Test Hit@50 ($\\uparrow$)** | 50.42±1.13 | 51.50±0.96 | 53.98±0.98 | 53.24±0.42 | 56.24±0.75 |\n\n\n**Q2: \"Are there any results on the benchmarks from [1]?\"**\n\nWe also added results on Wiki-CS from [1]. As the original split uses random sampling (i.e., no distribution shift), we introduce a harder setting: nodes are sorted by degree and split into 25%/25%/50% for train/val/test. This allows us to evaluate performance under topological shifts. ADiT-Series matches or slightly outperforms baselines under the original split, and shows notable gains in the harder setting with topological shifts:\n\n| | GCN | GAT  | GraphGPS | DIFFormer | ADiT-Series\n| - | - | - |  - |  - |  - |\n| **Test Acc ($\\uparrow$) with original split** | 77.46±0.85 | 76.90±0.82 | 78.34±0.88 | 79.39±0.62 | 79.53±0.42\n| **Test Acc ($\\uparrow$) with harder split** | 61.98±0.63 | 63.52±0.40 | 64.21±0.98 | 65.34±1.03 | 68.21±0.79\n\n**Q3: \"Not a fully novel approach since there are a series of existing works on graph neural diffusion\"**\n\nWe agree that graph neural diffusion has been studied extensively. However, our work introduces novelty in two key aspects: (1) we address topological shifts, a challenging but underexplored setting where training and test graphs differ; and (2) our approach is derived from a continuous PDE formulation that unifies local and non-local propagation, in contrast to prior works that focus primarily on local diffusion and static structures.\n\n[1] Dwivedi et al. Benchmarking graph neural networks.\n\n[2] Hu et al. Open graph benchmark: Datasets for machine learning on graphs, NeurIPS 2020.\n\nPlease let us know if further clarification is needed. Thank you again!", "claims_and_evidence": "The claims are clear with convincing evidence.", "methods_and_evaluation_criteria": "The method and its evaluation is technically sound.", "theoretical_claims": "The claims seem to be sound though I did not carefully check the proof.", "experimental_designs_or_analyses": "The experimental designs are valid and rigor.", "relation_to_broader_scientific_literature": "A general graph diffusion framework with competitive empirical performance, though not a fully novel approach since there are a series of existing works on graph neural diffusion.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "Nice unified view of global and local message passing.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "1. How is the performance of the proposed approach on the OGB benchmarks and OGB-LSC PCQM4Mv2 dataset?\n\n2. Are there any results on the benchmarks from [1]?\n\n[1] Dwivedi et al. Benchmarking graph neural networks.", "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": "ggc7IwSGhQ", "reviewer": "Reviewer_tsoh", "summary": "This paper proposes a framework called Advective Diffusion Transformer (ADIT) to address topological distribution shifts in graph-structured data. By formulating a continuous diffusion equation augmented with an advection term, the authors effectively combine non-local propagation with local message passing. They justify ADIT via a theoretical analysis showing that the model can control and limit the impact of topological changes on performance. Empirically, ADIT outperforms standard graph neural networks and transformer-based baselines across synthetic, informational (citation/social), and molecular/protein datasets.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful review and constructive feedback.\n\n**Q1: Theoretical assumptions and justification on more relaxed real-world conditions**\n\nOur PDE model indeed serves as a simplified abstraction, but it does not introduce new assumptions beyond the standard setting of graph representation learning [1–4]. Despite the simplification, PDE-based analysis allows for understanding GNN behavior, even on irregular and large-scale graphs, as shown in prior work [1-4] and our own.\n\nThe injectivity assumption in Thm 3.2 is indeed nontrivial but reasonable: since $g$ maps from a low-dimensional latent space to a higher-dimensional observation space, collisions (i.e., degenerate mappings) are rare in practice. Removing this assumption would be interesting, but it introduces substantial technical challenges that we'd like to leave to future work.\n\nCrucially, our model performs well on real datasets that do not satisfy the theoretical assumptions, confirming that the PDE framework yields effective architectures even under relaxed real-world conditions. While we agree that shifts in the wild can be more complex, our goal is to establish a principled framework and demonstrate its robustness to a wide range of topological shifts—including more diverse settings added below.\n\n**Q2: Adding experiments with more complex distribution shifts (combination of multiple shift factors, new nodes/edges and partial label availability)**\n\nBefore presenting new results, we kindly note that Table 1 already included dynamic shifts where new nodes and edges appear over time. We adopt an inductive setting (see details in Appendix E.1.2) where test nodes and their edges are excluded during training. Apart from this, we add experiments on a new dataset OGBL-COLLAB that is also a dynamic graph where new nodes/edges are added to the test graph. See our rebuttal to **Q1** of Reviewer ZJpC for results and details.\n\nTo further strengthen evaluation, we add a more challenging variant of the Arxiv dataset combining multiple shift types (node/edge changes and partial label availability). On top of the data splits that already involve dynamic shifts, we mask 10 of 40 classes for training data, and evaluate on test data with full labels. As shown below, while all models degrade under this setting, ADiT-Series maintains clear superiority:\n\n||GCN|GAT|GraphGPS|DIFFormer|ADiT-Series\n|-|-|-|-|-|-|\n|**Test Acc ($\\uparrow$) with single shift**|46.46±0.85|46.50±0.21|46.46±0.95|44.30±2.02|49.64±0.54|\n|**Test Acc ($\\uparrow$) with multiple shifts**|39.29±0.91|38.82±0.52|39.46±1.33|37.30±1.14|44.64±0.82|\n\n**Q3: Scalability to large graphs and comparison of computational costs**\n\nADiT-Series scales linearly with graph size, enabling it to scale up to large graphs. In practice, we use small values of \n$K$ (e.g., 1–4), yielding strong performance and efficiency. We add comparison of computational costs on the Arxiv graph (0.2M nodes). ADiT-Series consumes moderate GPU memory and much fewer time costs than Transformer competitors (DIFFormer and GraphGPS):\n\n||GCN|GAT|GraphGPS|DIFFormer|ADiT-Series\n|-|-|-|-|-|-|\n|**Train Time per Epoch (second)**|7.66|18.12|198.32|73.99|25.21|\n|**Train GPU Memory (GB)**|2.4|10.8|13.7|5.3|4.2|\n\n**Q4: Systematic procedure for hyperparameter selection**\n\nWe provided systematic studies for all hyper-parameters introduced by our model (such as $\\beta$, $\\theta$ and $K$) in Sec. 5.2 and Appendix F.2, with searching procedures and spaces detailed in Appendix E.3. \n\n**Q5: Comparison with non-PDE-based robust or domain-adversarial GNNs**\n\nWe appreciate the suggestion that can increase our impact to broader area of robust graph learning. We add a comparison on the DPPIN edge regression task against SR-GNN [5] and EERM [6], two robust GNNs—EERM uses domain-adversarial training. While all models show similar average RMSE, our models substantially improve worst-case performance, highlighting better stability under distribution shifts:\n\n||SR-GNN|EERM|ADiT-Inverse|ADiT-Series\n|-|-|-|-|-|\n|**Test Average RMSE ($\\downarrow$)**|0.170±0.003|0.172±0.006|0.166±0.006|0.167±0.004|\n|**Test Worst RMSE ($\\downarrow$)**|0.201±0.013|0.207±0.018|0.184±0.011|0.188±0.010|\n\nThis suggests that PDE-inspired architectures can offer robustness comparable to domain-invariant GNNs, without requiring adversarial training schemes.\n\n[1] GRAND: graph neural diffusion, ICML 2021\n\n[2] GRAND++: graph neural diffusion with a source term, ICLR 2022\n\n[3] Gradient gating for deep multi-rate learning on graphs, ICLR 2023\n\n[4] Understanding convolution on graphs via energies, TMLR 2024\n\n[5] Shift-robust gnns: Overcoming the limitations of localized graph training data, NeurIPS 2021\n\n[6] Handling Distribution Shifts on Graphs: An Invariance Perspective, ICLR 2022\n\nPlease let us know if further clarification is needed. Thank you again!", "claims_and_evidence": "The authors claim that ADIT maintains stable performance under distribution shifts in graph structure and the paper presents bounds indicating that ADIT's embeddings respond in a polynomially bounded way to changes in adjacency. However, the authors primarily focus on carefully controlled synthetic scenarios and specific real-world splits. In practice, shifts may be more diverse (e.g., entirely new nodes, partial re-labeling, dynamic edges). It remains unclear whether the same polynomial-bounded embedding behavior would hold under more chaotic or evolving graph processes.\nAlthough the theoretical discussion is convincing in isolated PDE contexts, real data can deviate from the carefully posited assumptions (e.g., injectivity of the latent generator). Thus, additional evidence or bounding under more relaxed conditions could strengthen the claim.", "methods_and_evaluation_criteria": "ADIT solves a PDE that unifies local and global propagation. Two variants approximate the continuous-time solution. However, there are some weaknesses:\nFirst, the computational overhead of matrix-inversion-based or large K-step expansions could be significant on huge graphs. While the authors propose optimizations (e.g., multi-head parallelism, Chebyshev expansions), the paper does not thoroughly compare memory/time profiles across different scales or under real-time constraints.\nSecond, the weighting hyperparameter is introduced to balance advection vs. diffusion, but the method to select it is somewhat heuristic.\nFor the evaluation criteria, it remains uncertain how these chosen splits map onto the PDE-based assumptions.", "theoretical_claims": "The authors present bounding arguments for out-of-distribution generalization error, contrasting polynomial vs. exponential sensitivity to adjacency changes. However, the authors assume “injectivity” in the latent generator g; real data might see degenerate or repeated node embeddings, which can weaken the claims. Moreover, the PDE analysis hinges on classical results that require smoothness assumptions and well-posedness of the underlying continuous system. Graph data can be highly irregular, with discrete or non-smooth connectivity. Some justification of how classical PDE frameworks carry over to more irregular or large-scale networks would be valuable.", "experimental_designs_or_analyses": "The topological splits for real datasets (by year or region) do not fully guarantee that all topological changes (e.g., changes in the node sets, or partial re-labeling) are tested. Additional experiments varying the node set or combining multiple shift factors might reveal more nuanced results.", "relation_to_broader_scientific_literature": "Recent research (e.g., robust GNNs under adversarial attacks or continual learning over evolving graphs) also addresses shifting topologies from different perspectives. A more explicit comparison or synergy with these areas (e.g., adversarially perturbed edges) would highlight whether ADIT can handle worst-case shifts or only mild, distributional ones.\nThe authors stress the PDE analogy but do not benchmark or contrast directly with non-PDE-based robust or domain-adversarial GNN approaches. Further references and direct comparisons might clarify whether PDE-based methods inherently outperform, or if domain-invariant approaches could match ADIT’s stability.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "Other strengths:\nADIT offers a clear PDE-based argument for combining local and global message passing; this is theoretically principled and practically relevant.\nOther weaknesses:\nFor extremely large graphs, it is unclear if the matrix exponential expansions remain viable or how approximations degrade performance.\nThe model introduces several PDE-inspired hyperparameters. A systematic procedure for selecting them under real constraints or partial label availability is not thoroughly explored.", "other_comments_or_suggestions": "No other comments.", "questions_for_authors": "No other questions.", "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": "npcTx81cwa", "reviewer": "Reviewer_sDvY", "summary": "The paper analysis the generalization capabilities of a specific class of diffusion-based graph models, namely those following the advective diffusion transformer, under topology distribution shifts. Specifically, it highlights that while the generalization gap for local diffusion—relying solely on the adjacency matrix—grows exponentially with changes in adjacency, incorporating a non-local propagation mechanism, in a transformer-like style, alongside the adjacency matrix reduces this gap to a polynomial scale. These findings are validated through experiments on a synthetic benchmark as well as several real-world datasets.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your time and thoughtful feedback.\n\n**Q1: Concerns about theoretical assumptions and (3.5) of (Van Loan, 1977)**\n\nWe appreciate the opportunity to clarify this point. The reviewer is correct that the result (3.5) of (Van Loan, 1977) [1] requires the matrix to commute with its transpose. This condition **does hold for Prop 3.4**, as the normalized adjacency matrix $\\tilde {\\mathbf A} = \\mathbf D^{-1/2}\\mathbf A \\mathbf D^{-1/2}$ is symmetric under the standard assumption that the observed graph is undirected—a common setting adopted by GNNs [e.g., 2–4] and justified via the graphon model described in Sec. 3.1. We’ll clarify this assumption in the revised paper.\n\nFor Thm 3.2, we agree that the symmetry of $\\mathbf I - \\mathbf C - \\beta \\mathbf V$ may not generally hold, especially since $\\mathbf C$ (attention matrix) may be asymmetric. While enforcing symmetry on $\\mathbf C$ would technically resolve this, it would overly restrict the model’s design space. Instead, we provide an alternative proof that **bypasses the need for (3.5) of (Van Loan, 1977) entirely**:\n\n> **Lemma.** Let $X, E \\in \\mathbb{R}^{n \\times n}$, and let $||\\cdot||$ be a submultiplicative matrix norm. Suppose there exist constants $M \\geq 1$, $\\omega \\geq 0$ such that for all $Y \\in \\mathbb{R}^{n \\times n}$, $||e^Y|| \\leq M e^{\\omega ||Y||}$. Then the following perturbation bound holds:\n$$\n||e^{X+E} - e^X|| \\leq ||E|| \\cdot M^2 \\cdot e^{\\omega (||X|| + ||E||)}.\n$$\n**Proof:** Define path $X(s) := X + sE$ for $s \\in [0, 1]$. Then:\n$$\ne^{X+E} - e^X = \\int_0^1 \\frac{d}{ds} e^{X(s)} \\ ds.\n$$\nUsing the integral form of the derivative of the matrix exponential [5], we have:\n$$\n\\frac{d}{ds} e^{X(s)} = \\int_0^1 e^{(1 - \\theta)X(s)} E e^{\\theta X(s)} \\ d\\theta.\n$$\nTherefore:\n$$\ne^{X+E} - e^X = \\int_0^1 \\left( \\int_0^1 e^{(1 - \\theta)X(s)} E e^{\\theta X(s)} \\ d\\theta \\right) ds.\n$$\nTaking norms and applying submultiplicativity:\n$$\n||e^{X+E} - e^X|| \\leq \\int_0^1 \\int_0^1 ||e^{(1 - \\theta)X(s)}|| \\cdot ||E|| \\cdot ||e^{\\theta X(s)}|| \\ d\\theta ds.\n$$\nUsing the growth bound assumption:\n$$\n||e^{(1 - \\theta)X(s)}|| \\leq M e^{\\omega (1 - \\theta)||X(s)||}, \\quad\n||e^{\\theta X(s)}|| \\leq M e^{\\omega \\theta ||X(s)||}.\n$$\nMultiplying these we have:\n$$\n||e^{(1 - \\theta)X(s)}|| \\cdot ||e^{\\theta X(s)}|| \\leq M^2 e^{\\omega ||X(s)||}.\n$$\nNote that $||X(s) || = ||X + sE|| \\leq ||X|| + ||E||$, so:\n$$\n||e^{X+E} - e^X|| \\leq ||E|| \\cdot M^2 \\cdot \\int_0^1 \\int_0^1 e^{\\omega ||X + sE||} d\\theta ds \\leq ||E|| \\cdot M^2 \\cdot e^{\\omega(||X|| + ||E||)}.\n$$\n\nThe existence of $M$ and $\\omega$ is guaranteed for every consistent matrix norm [5] such as spectral norm $||\\cdot||_2$ considered in our analysis.\n\nNow consider the proof for our Thm 3.2, let $\\mathbf L = \\mathbf I - \\mathbf C - \\beta \\mathbf V$ and $\\mathbf L' = \\mathbf I - \\mathbf C' - \\beta \\mathbf V'$. We can apply the above lemma even if $\\mathbf L$ and $\\mathbf L^\\top$ are not commutable:\n$$\n||e^{-\\mathbf L' T} - e^{-\\mathbf LT}||_2 \\leq M^2 T \\cdot ||\\mathbf L' - \\mathbf L||_2 \\cdot e^{\\omega T ||\\mathbf L||_2 } \\cdot e^{\\omega T ||\\mathbf L' - \\mathbf L||_2  }.\n$$\nFor spectral norm $||\\cdot||_2$ considered in our analysis, the above result holds for $M=1$ and $\\omega=1$. Also, the last term can be bounded using the same argument of Line 727-747 in our proof which leads to\n$$\ne^{||\\mathbf L' - \\mathbf L||_2} = e^{||(\\mathbf C' + \\beta \\mathbf V') - (\\mathbf C + \\beta \\mathbf V)||_2} \\leq O(||\\Delta \\tilde{\\mathbf A}||_2^m).\n$$\nTherefore, $||e^{-\\mathbf L' T} - e^{-\\mathbf LT}||_2$ is bounded with polynomial order w.r.t. $||\\Delta \\tilde{\\mathbf A}||_2$ and Thm 4.2 can be concluded. \n\nWe will revise the paper to replace the use of (3.5) of (Van Loan, 1977) in Thm. 3.2 with this updated, assumption-free proof.\n\n**Q2: GraphTrans vs. ADiT-Series memory cost**\n\nGraphTrans uses the original global attention mechanism, resulting in quadratic time and memory complexity w.r.t. node number. This makes it infeasible for large graphs like Arxiv (0.2M nodes), leading to the out-of-memory issue in Table 1.\n\nIn contrast, ADiT-Series employs a scalable global attention that preserves all-pair interactions while reducing time and memory complexity to linearity w.r.t. node number. This enables our model to handle large graphs efficiently. Details of this acceleration are presented in Appendix D.1.2, and we’ll highlight this distinction more clearly in the revision.\n\n**Q3: Missing references in Appendix**\n\nThank you for catching these. We will fix these typos in the revised paper. \n\n[1] Van Loan, C. The sensitivity of the matrix exponential.\n\n[2] GRAND: graph neural diffusion, ICML 2021\n\n[3] Diffusion improves graph learning, NeurIPS 2019\n\n[4]  Semi-supervised classification with graph convolutional networks, ICLR 2017\n\n[5] Higham et al. Functions of Matrices: Theory and Computation. SIAM, 2008.\n\nPlease let us know if further clarification is needed. Thank you again!", "claims_and_evidence": "Most of the proofs presented in the paper are based on the result (3.5) of (Van Loan, 1977). That result requires the matrix to commute with its traspose. However, that is not necessary the case for the $I-C-\\beta V$ matrix. In my opinion, additional assumptions needs to be done regarding the original graph topology and the coupling matrix C in order for the results to remain true. Since all the theoretical results presented in the paper (Theorem 3.2, Proposition 3.4 and everything derived from them) relies on this result, I would appreciate a clarification from the authors regarding the precise conditions under which their results hold. Specifically, could the authors clarify whether their framework implicitly assumes certain structural properties of C and the underlying graph, or whether alternative justifications exist for applying Van Loan’s result in this context?", "methods_and_evaluation_criteria": "The setup presented in the paper, with a test split consisting of graphs with shifted topological distribution represent an important and under-explored  area in geometric deep learning. The experimental results are interesting and in line with the theoretical observations.", "theoretical_claims": "Please see \"Claims And Evidence\" section.", "experimental_designs_or_analyses": "In Table 1 the paper reports out-of-memory results for the GraphTrans architecture. However, the proposed model presented in Section 4 and Algorithm 2 also relies on a fully connected graph, with the same type of attention mechanism as in the GraphTrans model. I am wondering why the GraphTrans memory requirement is higher than the one for AdiT-Series model presented in this paper.", "relation_to_broader_scientific_literature": "The performance of graph neural networks under distribution shift remains an under-explored area within the graph community. In this regard, I believe the paper introduces theoretical novelty. As for the proposed model, its architecture can be seen as an interpolation between a local graph diffusion and a fully connected transformer, which somewhat diminishes its individual contribution. However, this generalization also offers advantages, as the theoretical results can be specialized to these two widely used standard architectures.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "As mentioned above, my main concern regards the fault in the proofs. In my opinion, additional assumptions are made without being stated which can reduce the contribution of the work. If there is a way of bypassing them I am happy to raise my score to acceptance.", "other_comments_or_suggestions": "Section C in the appendix contains a missing reference (row 797, row 823, row 808, row 830)", "questions_for_authors": "Please see the boxes above.", "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:

Meta-review: The paper studies generalization under topological shifts and proposes a model motivated by the advective diffusion equation. They provide both theoretical justification and empirical evaluation of the resulting model. Reviewers raised concerns about some of the theoretical analysis and requested more extensive evaluation. The authors clarified both points in the rebuttal, leading to one score being raised, and thus all reviewers rating the paper "weak accept".
Decision: Accept (poster)