Supercharging Graph Transformers with Advective Diffusion

Thank you for your time and thoughtful feedback.

Q1: Concerns about theoretical assumptions and (3.5) of (Van Loan, 1977)

We 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.

For 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:

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:

$$||e^{X+E} - e^X|| \leq ||E|| \cdot M^2 \cdot e^{\omega (||X|| + ||E||)}.$$

Proof: Define path $X(s) := X + sE$ for $s \in [0, 1]$. Then:

$$e^{X+E} - e^X = \int_0^1 \frac{d}{ds} e^{X(s)} \ ds.$$

Using the integral form of the derivative of the matrix exponential [5], we have:

$$\frac{d}{ds} e^{X(s)} = \int_0^1 e^{(1 - \theta)X(s)} E e^{\theta X(s)} \ d\theta.$$

Therefore:

$$e^{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.$$

Taking norms and applying submultiplicativity:

$$||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.$$

Using the growth bound assumption:

$$||e^{(1 - \theta)X(s)}|| \leq M e^{\omega (1 - \theta)||X(s)||}, \quad ||e^{\theta X(s)}|| \leq M e^{\omega \theta ||X(s)||}.$$

Multiplying these we have:

$$||e^{(1 - \theta)X(s)}|| \cdot ||e^{\theta X(s)}|| \leq M^2 e^{\omega ||X(s)||}.$$

Note that $||X(s) || = ||X + sE|| \leq ||X|| + ||E||$, so:

$$||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||)}.$$

The existence of $M$ and $\omega$ is guaranteed for every consistent matrix norm [5] such as spectral norm $||\cdot||_2$ considered in our analysis.

Now 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:

$$||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 }.$$

For 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

$$e^{||\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).$$

Therefore, $||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.

We 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.

Q2: GraphTrans vs. ADiT-Series memory cost

GraphTrans 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.

In 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.

Q3: Missing references in Appendix

Thank you for catching these. We will fix these typos in the revised paper.

[1] Van Loan, C. The sensitivity of the matrix exponential.

[2] GRAND: graph neural diffusion, ICML 2021

[3] Diffusion improves graph learning, NeurIPS 2019

[4] Semi-supervised classification with graph convolutional networks, ICLR 2017

[5] Higham et al. Functions of Matrices: Theory and Computation. SIAM, 2008.

Please let us know if further clarification is needed. Thank you again!

Thank you for your time in reviewing our paper and for the constructive feedback.

Q1: "How is the performance of the proposed approach on the OGB benchmarks?"

We 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:

Q2: "Are there any results on the benchmarks from [1]?"

We 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:

Q3: "Not a fully novel approach since there are a series of existing works on graph neural diffusion"

We 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.

[1] Dwivedi et al. Benchmarking graph neural networks.

[2] Hu et al. Open graph benchmark: Datasets for machine learning on graphs, NeurIPS 2020.

Please let us know if further clarification is needed. Thank you again!

Thank you for your thoughtful review and constructive feedback.

Q1: Theoretical assumptions and justification on more relaxed real-world conditions

Our 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.

The 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.

Crucially, 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.

Q2: Adding experiments with more complex distribution shifts (combination of multiple shift factors, new nodes/edges and partial label availability)

Before 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.

To 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:

Q3: Scalability to large graphs and comparison of computational costs

ADiT-Series scales linearly with graph size, enabling it to scale up to large graphs. In practice, we use small values of $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):

Q4: Systematic procedure for hyperparameter selection

We 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.

Q5: Comparison with non-PDE-based robust or domain-adversarial GNNs

We 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:

This suggests that PDE-inspired architectures can offer robustness comparable to domain-invariant GNNs, without requiring adversarial training schemes.

[1] GRAND: graph neural diffusion, ICML 2021

[2] GRAND++: graph neural diffusion with a source term, ICLR 2022

[3] Gradient gating for deep multi-rate learning on graphs, ICLR 2023

[4] Understanding convolution on graphs via energies, TMLR 2024

[5] Shift-robust gnns: Overcoming the limitations of localized graph training data, NeurIPS 2021

[6] Handling Distribution Shifts on Graphs: An Invariance Perspective, ICLR 2022

Please let us know if further clarification is needed. Thank you again!