Linear Transformer Topological Masking with Graph Random Features

We thank the reviewer for their comments. We are pleased that they think our research is important, and that they appreciate the theoretical contributions and ablation analysis. We respond to all their questions and concerns below.

1. `Background and related works section is limited’. We respectfully point the reviewer to Appendix B, which gives a very detailed account of the relationship to existing work. We wonder whether they may have missed it, and will be sure to signpost it better in the main body. With regards to the specific papers they suggest:

• Nodeformer [1]: This paper focuses on learning latent graph structures, rather than incorporating information about an existing graph information as a structural inductive bias. It is not directly relevant.
• KDLGT [2]: This paper does indeed incorporate a structural inductive bias into Transformers, but it does so with a simple additive RPE mechanism. This is not able to modulate the attention by a multiplicative graph function, which is our core goal (see Eq. 2). Nonetheless, we agree that it might be of interest to readers so have added a brief description in the Appendix. Thanks for the pointer.

2. Experiments and evaluations on traditional graph datasets. Respectfully, we already include experiments on three different data modalities: images, videos and point clouds for robotics. We also include extensive ablation studies. Experiments like node classification on datasets such as Citeseer are not the focus of this paper, but might make for interesting future work in a separate publication.
3. Confusing notation in line 148. $\mathcal{N} = \\{v_1, v_2, …, v_N \\}$ denotes the set of graph nodes. $N$ is the size of this set, i.e. the total number of graph nodes. In the Transformer context, this is equal to the number of tokens. This notation is very standard.

Once again, we thank the reviewer. We have pointed them to our existing related work section (now supplemented with their additional suggestion), and clarified some points of minor misunderstanding. We very much hope that, on reflection, they will raise their score.

[1] Nodeformer: A Scalable Graph Structure Learning Transformer for Node Classification, Wu et al., NeurIPS 2022, https://arxiv.org/abs/2306.08385
[2] KDLGT: A Linear Graph Transformer Framework via Kernel Decomposition Approach, Wu et al., IJCAI 2023, https://www.ijcai.org/proceedings/2023/0263.pdf

We thank the reviewer for their comments. We are pleased that they find the paper well organised and written, and that they appreciate the importance of our methods when $N$ becomes large.

1. Time complexity is $\mathcal{O}(Nmd)$. The reviewer is correct that, as with any low-rank attention method with $N$ tokens of dimensionality $d$ and features of dimensionality $m$, the time complexity of Eq. 3 is $\mathcal{O}(Nmd)$. We state this in line 121. Since we are primarily interested in scaling with respect to the number of tokens, we follow convention in the literature by shortening this to $\mathcal{O}(N)$ throughout the text [1,2,3]. For example, for the Interlacer experiment with massive point clouds, $d=64$, $m=16$ but $N=32768$, so the scaling with respect to $N$ is most interesting and important. We believe our notation choice to be standard and unambiguous, and it reduces clutter. However, we agree that it is important to be as clear as possible so will add this comment to the manuscript. Thanks.
2. ’Line 161 is about graph kernels, not graph node kernels’. Thanks for the comment. Graph kernels and graph node kernels are closely related: graph kernels define a product graph using the two constituent graphs, compute the corresponding graph node kernel, then sum up its (weighted) entries. See e.g. page 4 of Bogwardt [4]. For this reason, we still believe the reference to be relevant for our work.
3. Number of walks and computation time. Thanks for the suggestion. We already include extensive ablations: number of walkers $n$ vs. Transformer performance, the presence or absence of importance sampling, and use of fully learnable features instead of GRFs. Nonetheless, we agree that including number of walkers vs. number of FLOPs would be a nice addition, so have now added a further plot to Appendix C.2. The cost of sampling random walks is linear in the number of walkers. The cost of computing the corresponding masked attention in Eq. 7 initially grows at most linearly (by Lemma 3.2), but will eventually saturate once the features become dense. Our plot reflects this. Thanks again for prompting a nice addition.
4. Section 4 headers. Thanks for the stylistic comments about using section headers instead of paragraphs. We agree that this may be clearer, so have updated the manuscript.
5. Relationship to GNNs. Thanks for the comment. We actually already discuss the relationship to GNNs in the paragraph beginning on line 367. GATs [5] are a particular type of Transformer/GNN with a very strong structural inductive bias, where nodes only attend to their neighbors. Our scheme could be considered a stochastic relaxation of this, still injecting information about the graph but now including longer-range attention by importance sampling random walks. Also, our method can be considered more expressive than GNNs because it is able to distinguish graphs identical under the 1-dimensional Weisfeiler Lehman graph isomorphism heuristic.

We again thank the reviewer for their feedback. We warmly invite them to respond with any further questions, and respectfully ask that they consider raising their score.

[1] Rethinking Attention with Performers, Choromanski et al., ICLR 2021, https://arxiv.org/abs/2009.14794
[2] From Block-Toeplitz Matrices to Differential Equations on Graphs: Towards a General Theory for Scalable Masked Transformers, Choromanski et al., ICML 2022, https://doi.org/10.48550/arXiv.2107.07999
[3] Reformer: The Efficient Transformer, Kitaev et al., ICLR 2020, https://arxiv.org/pdf/2001.04451
[4] Protein Function Prediction via Graph Kernels, Bogwardt et al., Bioinformatics 2005, https://doi.org/10.1093/bioinformatics/bti1007
[5] Graph Attention Networks, Veličković et al., ICLR 2018, https://doi.org/10.48550/arXiv.1710.10903

We thank the reviewer for their comments. We are pleased that they agree the paper is well-motivated and that they find the explanations clear. We address all their concerns and questions below.

1. ‘Can the author also present other results, like the total training time or total flops, to validate the proposed method’s efficiency?’ Thanks for the suggestion. We actually already include an experiment showing how the total number of FLOPs scales with the input size: see Fig. 3. As expected, our method is linear. There is a constant multiplicative cost incurred by topological masking with GRFs, and we are much more efficient than the softmax baseline. To supplement with some wall clock times, for the robotics PCT experiment the MP Interlacer baseline [1] (previous SOTA) trains at 0.925 steps/second. Our method, the GRFs Interlacer, trains at 0.982 steps/second. Hence, our method is not only more accurate, but also faster by 6%. (The unmasked baseline, which struggles to capture local structure so gives poor accuracy, trains at 1.49 steps/second). We have added these results to the manuscript; thanks for the suggestion. Finally, we emphasise that, for point clouds with $>30k$ nodes, computing full-rank softmax attention is not even possible on any reasonable hardware – the time and memory requirements are too great. The fact that we can train our model and make predictions on such a massive graph provides further experimental evidence that it is extremely efficient. As such, we respectfully suggest that the paper includes plenty of experimental results showing our algorithm’s efficiency, in addition to our detailed theoretical arguments (Sec. 3.4).
2. Difference between GRF Interlacer and MP Interlacer. The GRF Interlacer (our method) uses GRF-masked linear attention, whereas the MP Interlacer [1] uses GNN-style message passing layers. Our method achieves higher image SSIM (prediction accuracy) because it is more expressive and better able to model complex dynamics in the point cloud. The reviewer is correct to note that the difference between the GRF and MP Interlacers seems to become smaller after many rollout timesteps. This may be because applying several GNN layers in sequence gradually increases the receptive field, becoming closer to topological masking (where instead nodes can attend if their ensembles of random walks overlap). However, as is very often the case with deep learning methods, this interpretation is speculative.
3. Besides GRFs, are there any other methods to do implicit graph masking? A key aspect of our work is that it is the first linear algorithm for $\mathcal{O}(N)$ topological masking on general graphs – so no, there is no equally efficient alternative to directly compare against. The closest previous algorithm is block Toeplitz masking [2], which is $\mathcal{O}(N\log N)$ rather than $\mathcal{O}(N)$ and can only be applied to grid graphs. We include this benchmark for the image experiments in Table 1. We find that, despite our method being cheaper, it still tends to match or beat block Toeplitz.

Once again, we thank the reviewer. We believe that we have resolved all their questions and concerns, and have added the extra wall clock times for the Interlacer experiment to the manuscript. We warmly invite them to respond with any further questions, and respectfully request that they consider raising their score.

[1] Modelling the Real World with High Density Visual Particle Dynamics, Whitney et al., CoRL 2024, https://arxiv.org/abs/2406.19800
[2] From Block-Toeplitz Matrices to Differential Equations on Graphs: Towards a General Theory for Scalable Masked Transformers, Choromanski et al., ICML 2022, https://doi.org/10.48550/arXiv.2107.07999

(Part 2/2)

Focus on topological masking c.f. improving linear attention. The reviewer is correct that our paper focuses on developing new topological masking (graph RPE) techniques that are compatible with existing linear attention algorithms. Our methods are agnostic to the particular feature map $\phi$ used to replace softmax in linear attention: optimising $\phi$ is not the goal of the work. However, we agree that it may be interesting to see the gains that topological masking provides to other linear attention variants, so we are running extra experiments with positive random features [FAVOR+, 4] instead of ReLU. Preliminary results are already complete and, as expected, we again see strong gains from incorporating GRFs. For ViT-B/16 trained from scratch on ImageNet with FAVOR+ attention ($m=256$ random features), we see a relative improvement of +2.3% from our algorithm compared to the unmasked baseline ($p=0.1$, $n=100$ walks). Once the rest of the additional experiments are complete, we will add them to the manuscript. Thanks for suggesting this.

Minor points:

1. $f: \mathbb{N} \to \mathbb{R}$ can be interpreted as a function mapping from the natural numbers to real numbers. $(f_k)^\infty_{k=0}$ is a sequence of reals, intended to denote the evaluations of some particular $f$ for the natural numbers rather than a sequence of functions. We will make this clear. Thanks for the suggestion.
2. Normalisation of $\mathbf{W}$. Please see earlier comments. This is standard practice to ensure that the kernel converges.

We again thank the reviewer for their time. Having clarified misunderstandings about our theoretical contributions and added extra experiments with the FAVOR+ linear attention mechanism, we hope that they will raise their score. We warmly invite them to respond with any further questions.

[1] Kernels and regularization on graphs, Smola and Kondor, COLT 2003, https://people.cs.uchicago.edu/~risi/papers/SmolaKondor.pdf
[2] Spectral graph theory, Chung, 2007, https://mathweb.ucsd.edu/~fan/research/revised.html
[3] General graph random features, Reid et al., ICLR 2024, https://arxiv.org/pdf/2310.04859
[4] Rethinking attention with Performers, Choromanski et al., ICLR 2021, https://arxiv.org/pdf/2009.14794

(Part 1/2)

We thank the reviewer for their comments. We are pleased that they note our work’s strong theoretical foundations, concrete performance improvements, diverse applications and ablation studies. We are grateful for their time.

However, we strongly disagree that any of the theoretical analysis in the paper is incorrect. We will clarify points of misunderstanding below.

To recapitulate the results in the paper, our key theoretical contributions are as follows.

1. Theorem 3.1 gives the first known concentration bound for GRFs. The bound depends on the number of walkers $n$, and a scalar denoted $c$. As the reviewer notes, $c$ depends on the underlying graph via its maximum node degree and edge weight: specifically, via $\max_{v_i, v_j \in \mathcal{N}}(|\mathbf{W}_{ij}| d_i)$ (line 266).
2. Meanwhile, Lemma 3.2 bounds the GRF sparsity for any graph, given some number of walkers $n$ and termination probability $p_\textrm{halt}$. The simple proof is based on bounding the number of hops a walker can take before it terminates, supposing its length is geometrically distributed. For given $n$ and $p_\textrm{halt}$, this bound is independent of $\mathcal{G}$ (though of course how tight it is depends on the structure of $\mathcal{G}$; it is generally tighter for denser, bigger graphs).
3. Corollary 3.3 combines the two observations above to show that, provided $c$ remains constant as the graph grows – meaning we do not have to update our bound so we can use the same number of walkers $n$ and termination probability $p_\textrm{halt}$ without compromising estimator accuracy – then the number of nonzero GRF entries is bounded by a constant with high probability when $N$ gets large. As the reviewer correctly notes, $\mathcal{O}(N)$ time complexity follows.

To be clear, Theorem 3.1 and Lemma 3.2 hold for any graph. Corollary 3.3 does indeed require that $c$ remains constant as the graph grows, in order that the same bound holds so we can safely use the same $p_\textrm{halt}$ and $n$. We are totally upfront about this requirement. In Sec. 3.4, we emphasise it three times, including in the core statement of the corollary: see lines 278, 295 and 313. We initially omitted this technical detail from the introduction, but on reflection agree that the interested reader might benefit from encountering it earlier in the text. Therefore, we have now updated the paper introduction to flag it even more explicitly – thanks for the suggestion.

$c$ remaining constant as $N$ grows is not very restrictive. Though sparsity can be used to ensure that $c$ remains constant as $N$ grows, this is not a necessary condition. For instance, in the reviewer’s example of a complete graph with edge weights $1/(N-1)$ , we have that $\max_{v_i, v_j \in \mathcal{N}}(|\mathbf{W}_{ij}| d_i) = 1$ which is independent of $N$. This is an example of a dense graph for which our assumption about $c$ holds, so it is not the case that the analysis only holds for sparse graphs. Stepping back to take a broader perspective, the reason we need to control $c$ is to prevent the spectral radius of the adjacency matrix $\mathbf{W}$ diverging as the graph becomes large. If we do not do this, the underlying exact kernel (defined in Eq. 4) will also diverge, in which case one clearly cannot approximate it with GRFs or otherwise. 'Regularising’ or 'normalising’ $\mathbf{W}$ to control its spectral radius is very standard in the graph node kernel literature [1,2,3]. It is not a weakness of our specific approach to topological masking.

Regularising W. To build on the above, one typically 'regularises' by taking $W_{ij} \to W_{ij}/ \sqrt{d_i d_j}$ with $d_i = \sum_j W_{ij}$. This bounds its spectral radius to $1$ [2]. Since we start with unweighted adjacency matrices this gives edge weights $1/\sqrt{d_i d_j}$, but our method does not formally require this.

The reviewer’s proposed counterexample does not consider asymptotic $N$. `Big O’ notation describes the limiting behaviour of the time complexity when $N$ tends to infinity. In contrast, the reviewer’s proposed counterexample looks at the small $N$ regime, where any asymptotic analysis inevitably breaks down. To make this concrete, consider a complete graph with $n=1$ walks and termination probability $p_\textrm{halt}$. When the number of nodes $N$ becomes very large, the probability of a walker backtracking becomes small, so at every timestep it hops to a new, unvisited node. The walker length will be $\frac{1}{p_\textrm{halt}}$ on average, so at asymptotic $N$ only $\sim \frac{1}{p_\textrm{halt}}$ coordinates of the GRF will be nonzero. This is manifestly independent of $N$. In contrast, the reviewer’s example considers the small $N$ regime, where time complexity scaling results are not expected to hold. In experiments, this behaviour is shown by the small nonlinear regime at the far left of Fig. 3. There is no problem with our theoretical claims.

CONTINUED BELOW.

We thank the reviewer for their quick reply. We are grateful for their perspective

The reviewer is correct to note that, for a given $p_\textrm{halt}$ and number of walkers $n$ (selected to guarantee sharp kernel estimates by Thm. 3.1), the actual sparsity of the corresponding GRFs can depend on the particular graph being considered. For sparse graphs GRFs will tend to be more sparse (since there is a higher probability of backtracking with fewer neighbours), and for dense graphs GRFs will tend to be less sparse (because walkers are more likely to visit a new, previously unvisited node). However, the bound in Lemma 3.2 already considers the worst possible case, when walkers never backtrack. You can see this in lines 288 and 289, when we say ‘$n$ walkers are all $b$ or shorter with probability $(1-(1-p_\textrm{halt})^b)^n$... at most $bn$ entries of $\widehat{\phi}(v_i)$ can then be nonzero’. Clearly, this corresponds exactly to the scenario when the number of visited nodes is equal to the total number of hops. Therefore, our bound will actually be tighter for denser graphs. Put simply, our algorithm will be even more efficient for sparse graphs, whereas for dense graphs the time complexity will be closer to the theoretical worst case. Of course, it is $\mathcal{O}(N)$ in both scenarios.

We agree that the question of how GRF sparsity and mask (kernel) approximation quality are related is an interesting one. We direct the reviewer to ‘General Graph Random Features’ [1], the paper which introduced the algorithm, which provides a detailed investigation of this question in App A.6. Respectfully, this kind of ablation for GRFs is not the main focus of this paper, which instead provides new concentration bounds and applies them to a range of tasks in Transformers. Nonetheless, we do agree that the interested reader may benefit from some brief empirical investigation to avoid needing to separately refer to this previous work.

For this reason, following the reviewer’s suggestion, we have now added a new experiment to the Appendix – please see App. E. Taking different Erdős–Rényi graphs with edge probabilities between $0.1$ and $0.9$, we show how 1) the mask approximation quality and 2) the GRF sparsity varies. As anticipated, both become slightly worse then plateau as the graphs become denser. However, they vary within a narrow range of $\sim$ 10%, so this need not be a big practical concern. We also provide similar plots for a varying number of graph nodes $N$ with fixed graph sparsity (edge probability 0.5), where we again find that GRF performance tends to be a little better for smaller graphs but quickly plateaus for big graphs. Meanwhile, the GRF sparsity drops as $N$ grows: at most $\mathcal{O}(1)$ entries are nonzero, so the proportion of nonzero entries goes down as $\mathcal{O}(1/N)$. We emphasise once again that these experiments serve to show how tight the bounds are, but the bounds themselves still hold in every case – put simply, one may get an even faster algorithm for very sparse, small examples, but our theoretical results prove that the algorithm is efficient in all scenarios. Note especially that in every case the mask approximation quality remains excellent, characterised by a tiny relative Frobenius norm error, even for dense, big topologies.

We thank the reviewer for prompting us to add these experiments regarding scaling behaviour with multiple walkers on dense graphs to the paper. We trust that they have allayed any remaining practical concerns. We also politely remind them that our Transformer experiments already include three different data modalities across multiple datasets and tasks: images, videos and point cloud data for robotics. Previously proposed algorithms for topological masking for special graphs (sequences, grids or trees) have typically focussed on just one [2,3]. As such, we respectfully suggest that we have already shown that our algorithm is very practical in a broad range of settings. Nonetheless, we agree that these latest additions have improved the paper – thanks for the suggestions.

With these additions and clarifications (of course, as well as the additional FAVOR+ experiments we are running for the reviewer), would the reviewer please consider raising their score and recommending acceptance?

[1] General Graph Random Features, Reid et al., ICLR 2024, https://arxiv.org/abs/2310.04859
[2] From Block-Toeplitz Matrices to Differential Equations on Graphs: Towards a General Theory for Scalable Masked Transformers, Choromanski et al., ICML 2022, https://doi.org/10.48550/arXiv.2107.07999.
[3] Stable, Fast and Accurate: Kernelized Attention with Relative Positional Encoding, Luo et al., NeurIPS 2021. URL https://doi.org/10.48550/arXiv.2106.12566

We thank the reviewer for their comments. We are pleased that they find our theoretical contributions valuable, and that they recognise the improvements to predictive performance. We answer all their questions and address points of misunderstanding below.

1. Mathematical presentation. We are sorry that the reviewer found some of the technical exposition difficult to follow. We took great care to include high-level passages (line 224, ‘at a high-level…’) and a visual schematic (Fig. 2), designed to help build intuition. We agree that the notation around Thm 3.1 and the use of McDiarmid’s inequality becomes dense. Are there any specific sentences or mathematical contributions that the reviewer did not understand, which we may clarify for them?. We would be very happy to try to rephrase any parts the reviewer can flag as confusing.
2. Limitations. Respectfully, we do not agree with the reviewer’s list of suggested limitations.

• Dependence on graph structure. One of our key theoretical contributions is that GRF estimators are sharp for any graph with bounded $c$ (Thm. 3.1), so it is not true to say that the method depends heavily on graph specifics. This is in sharp contrast to previously published algorithms which are restricted to e.g. grids or trees [1,2]. Intuitively, since we do importance sampling of random walks, we still capture long range dependencies. We can sample these long walks by chance, and then upweight them depending on their probability. Please see App. C.2 for detailed ablations. We emphasise that our method works very well for image classification, which may in general require these long range dependencies. This gives evidence that our method captures them.
• Dynamic and evolving graphs. Our point cloud experiment on page 9 exactly considers this case of a dynamic graph where nearest neighbours must be re-computed at each timestep. Our method is found to be extremely efficient, and is the best-performing variant compared to the baselines. Indeed, it is not only more accurate, but its training wall clock times are actually 6% faster compared to the previous SOTA (MP Interlacer baseline). We have added this wall clock time result to the paper. Whilst we agree that computing (approximate) $k$-nearest neighbours on massive point clouds can in general be expensive, this is not a specific of our algorithm. It is a property of this type of dynamic data, and is also required for the MP Interlacer baseline or indeed any GNN-type model in this context.
• Choice of hyperparameters. The reviewer is correct that our method is stochastic, similar to the popular and highly-cited Performers paper [3]. However, one of our key theoretical contributions is how to choose the number of walkers $n$ and termination probability $p_\textrm{halt}$ to guarantee sharp kernel estimates with high probability (Thm 3.1). We provide a very specific recipe for choosing them based on our novel results for the estimator concentration; the hyperparameters do not need to be manually tuned by the practitioner. Therefore, it is not correct to say that the randomised nature of the algorithm makes it difficult to generalise across graph structures – in fact, the converse is true.

CONTINUED BELOW.

We thank the reviewer. We are pleased that, despite having a different research background, they appreciate that our algorithms are novel and scale well to massive datasets. We answer their questions in detail below.

1. Applicability of GRFs. Graph random features were only recently introduced [1,2], so they have not yet been used for the full range of graph-based learning tasks. For instance, our paper is the first to apply them in Transformers. However, GRFs have previously been used for efficient kernelised node clustering for Citeseer [3, see Table 2].
2. Computation times c.f. baselines. We direct the reviewer for Fig. 3, which shows the total number of FLOPs vs. the number of graph nodes for our method (linear + GRFs) compared to the baselines (softmax and unmasked linear). We report FLOPs because it is agnostic to the hardware being used. This experiment confirms that our method is $\mathcal{O}(N)$ time complexity and is faster than full-rank softmax. To supplement with some wall clock times, for the robotics PCT experiment the MP Interlacer baseline [4] (previous SOTA) trains at 0.925 steps/second. Our method, the GRFs Interlacer, trains at 0.982 steps/second. Hence, our method is not only more accurate, but also faster by 6%. (The unmasked baseline, which struggles to capture local structure so gives poor accuracy, trains at 1.49 steps/second). We have added these results to the manuscript; thanks for the suggestion.

We again thank the reviewer. Having answered their questions and added some extra wall clock times, we respectfully ask that they consider raising their score.

[1] Taming Graph Kernels with Random Features, Choromanski, ICML 2023, https://arxiv.org/abs/2305.00156
[2] General Graph Random Features, Reid et al., ICLR 2024, https://arxiv.org/abs/2310.04859
[3] Quasi Monte Carlo Graph Random Features, Reid et al., NeurIPS 2023, https://arxiv.org/abs/2305.12470
[4 Modelling the Real World with High Density Visual Particle Dynamics, Whitney et al., CoRL 2024, https://arxiv.org/abs/2406.19800

We would like to sincerely thank the Reviewer for the question. Efficient linear-attention Transformers with geometrically modulated attention via GRFs can indeed be leveraged in settings, where GNNs are used. This is the case since they can be naturally applied to graph data. In this context, the GRF mechanism serves as a relative positional encoding method, capable of discounting direct interactions between tokens faraway in the metric induced by the particular graph kernel (potentially learnable). For the comparison, in the GNN setting this discounting is usually implemented by modeling direct interactions only between pairs of adjacent nodes. In fact our application of particle-based dynamics is an example of the usability in the setting where GNNs are used on a regular basis (particle-based dynamics via GNNs is a subject of the voluminous literature). Therefore other potential applications of our methods include in particular bio-informatics (e.g. drug design and molecular biology, where GNNs are machine learning methods of choice).