Linear Transformer Topological Masking with Graph Random Features

## Abstract

Abstract When training transformers on graph-structured data, incorporating information about the underlying topology is crucial for good performance. Topological masking , a type of relative position encoding, achieves this by upweighting or downweighting attention depending on the relationship between the query and keys in a graph.
In this paper, we propose to parameterise topological masks as a learnable function of a weighted adjacency matrix – a novel, flexible approach which incorporates a strong structural inductive bias.
By approximating this mask with graph random features (for which we prove the first known concentration bounds), we show how this can be made fully compatible with linear attention, preserving 𝒪 ⁢ ( N ) 𝒪 𝑁 \mathcal{O}(N) caligraphic_O ( italic_N ) time and space complexity with respect to the number of input tokens.
The fastest previous alternative was 𝒪 ⁢ ( N ⁢ log ⁡ N ) 𝒪 𝑁 𝑁 \mathcal{O}(N\log N) caligraphic_O ( italic_N roman_log italic_N ) and only suitable for specific graphs.
Our efficient masking algorithms provide strong performance gains for tasks on image and point cloud data, including with > 30 absent 30 >30 > 30 k nodes.

## 1 Introduction

Across data modalities, transformers have emerged as a leading machine learning architecture (Vaswani, ; Dosovitskiy et al., ; Arnab et al., ).
They derive their success from modelling complex dependencies between the tokens using an *attention mechanism*.
However, in its vanilla instantiation the transformer is a set function, meaning it is invariant to permutation of its inputs.
This is often undesireable; e.g. words tend to be arranged in a particular sequence, and graph nodes are connected to a particular set of neighbours.
Therefore, it has become standard to modify the attention matrix depending on any underlying structure, building in inductive bias.
For instance, *causal attention* applies a triangular binary mask to ensure that tokens only attend to preceding elements of the sequence (Yang et al., ).
Meanwhile, *relative position encoding* (RPE) captures the spatial relationship between the tokens (Shaw et al., ).
In some cases, the queries and keys can be arranged in a graph $\mathcal{G}$ and a function $\mathbf{M}(\mathcal{G})$ can be used to modulate the transformer attention.
This is referred to as *topological masking* (Choromanski et al., ).
It provides a powerful avenue for incorporating structural inductive bias from $\mathcal{G}$ into transformers, in some cases closing the performance gap with the best-customised graph neural networks (Ying et al., ).

Topological masking of low-rank attention.
Topological masking is simple when the attention is full-rank.
In this case, the entire attention matrix $\mathbf{A}\in\mathbb{R}^{N\times N}$ (with $N$ the number of input tokens) is materialised in memory, so its individual entries $\mathbf{A}_{ij}$ can be pointwise multiplied by a (fixed or learnable) mask $\mathbf{M}(\mathcal{G})_{ij}$.
But full-rank attention famously incurs a quadratic $\mathcal{O}(N^{2})$ space and time complexity cost, which may become prohibitive on large inputs.
This has motivated a number of *low-rank decompositions* of $\mathbf{A}$.
These improve scalability by avoiding computing $\mathbf{A}$ explicitly, instead doing so implicitly in feature space.
The question of how to perform topological masking in the low-rank setting is challenging and not fully solved.
Previously proposed algorithms are limited to inexpressive functions $\mathbf{M}$ and restricted classes of graphs $\mathcal{G}$ (e.g. trees and grids), and are often $\mathcal{O}(N\log N)$ time complexity rather than $\mathcal{O}(N)$ (App. B; Choromanski et al., ; Liutkus et al., ).
Addressing these shortcomings is the chief focus of this manuscript.
Our central question is:
*how can we implement efficient yet flexible topological (graph-based) masking of transformer attention, in the low-rank setting where $\mathbf{A}$ is never explicitly instantiated?*

Graph random features.
To answer the question posed above, we propose to use *graph random features* (GRFs) (Choromanski, ; Reid et al., ).
GRFs are sparse, randomised vectors whose dot product gives an unbiased estimate of arbitrary functions of a weighted adjacency matrix, a class that captures the structure of $\mathcal{G}$ and is popular in classical graph-based machine learning (Smola and Kondor, ; Kondor and Lafferty, ). GRFs have been used for approximate kernelised node clustering and in scalable graph-based Gaussian processes (Reid et al., ; ), but never before in transformers.
Advancing understanding of the convergence properties of GRFs is another key goal of this work.

Our contributions. We present the first known $\mathcal{O}(N)$ algorithm for topological masking of linear attention transformers with general $\mathcal{G}$, using graph random features (Fig. 1).

- 1.
We propose to modulate the transformer attention mechanism with a *learnable function of a weighted adjacency matrix* – a simple, general topological masking trick which brings a strong inductive bias and substantial improvements to predictive performance.
- 2.
For scalability, we replace the exact, deterministic mask with a sparse, stochastic estimate constructed using *graph random features* (GRFs). We use a novel exponential concentration bound to prove that our method is linear time and space complexity with respect to input size, the first known algorithm with this property.
- 3.
We demonstrate strong accuracy gains on image data, as well as for modelling the dynamics of massive point clouds ($>30$k particles) in robotics applications where efficiency is essential.

Figure: Figure 1: Schematic overview. Regular attention is $\mathcal{O}(N^{2})$, with $N$ the number of input tokens. Topological masking modulates $\mathbf{A}$ by a graph function $\mathbf{M}(\mathcal{G})$, improving predictive performance. Linear attention reduces the time complexity to $\mathcal{O}(N)$ by leveraging a low-rank decomposition. Our contribution (blue) is *the first algorithm to achieve both* – $\mathcal{O}(N)$ topological masking of low-rank attention – by approximating $\mathbf{M}(\mathcal{G})$ with graph random features (GRFs). GRFs are sparse vectors (denoted ‘sp.’) computed by sampling random walks, constructed so that $\mathbb{E}(\widehat{\mathbf{\Phi}}_{\mathbf{Q},\mathcal{G}}\widehat{\mathbf{ \Phi}}_{\mathbf{K},\mathcal{G}}^{\top})=\mathbf{A}\odot\mathbf{M}(\mathcal{G})$ with strong concentration properties (Thm. 3.1).

## 2 Preliminaries

Masked attention.
Denote by $N$ the number of input tokens and $d$ the dimensionality of their latent representations.
Consider matrices $\mathbf{Q},\mathbf{K},\mathbf{V}\in\mathbb{R}^{N\times d}$ whose rows are given by the query ($\{\bm{q}_{i}\}_{i=1}^{N}$), key ($\{\bm{k}_{i}\}_{i=1}^{N}$) and value vectors ($\{\bm{v}_{i}\}_{i=1}^{N}$) respectively.
The *attention mechanism*, a basic processing unit of the transformer, can be written:

$$ $\textrm{Att}(\mathbf{Q},\mathbf{K},\mathbf{V})\coloneqq\mathbf{D}^{-1}\mathbf{ A}\mathbf{V},\quad\mathbf{A}\coloneqq\mathcal{K}(\mathbf{Q},\mathbf{K}),\quad \mathbf{D}\coloneqq\mathbf{A}\mathbf{1}_{N}.$ (1) $$

Here, the function $\mathcal{K}$ assigns a similarity score to each query-key pair and $\mathbf{1}_{N}\coloneqq[1]_{i=1}^{N}\in\mathbb{R}^{N}$.
For regular softmax attention, one simply takes $\mathcal{K}(\mathbf{Q},\mathbf{K})_{ij}=\exp(\bm{q}_{i}^{\top}\bm{k}_{j})$, optionally also normalising by $d$.
To incorporate *masking*, we make the modification

$$ $\mathbf{A}_{\mathbf{M}}\coloneqq\mathbf{M}\odot\mathcal{K}(\mathbf{Q},\mathbf{ K}),$ (2) $$

where $\mathbf{M}\in\mathbb{R}^{N\times N}$ is a (hard or soft) *masking matrix* and $\odot$ denotes the Hadamard (element-wise) product.
For softmax, this is equivalent to adding a bias term $b_{ij}\coloneqq\log\mathbf{M}_{ij}$ to $\bm{q}_{i}^{\top}\bm{k}_{j}$ before exponentiating.
As remarked in Sec. 1, this can be used to enforce causality or as a relative position encoding (RPE) (Shaw et al., ). Taking $\mathbf{M}(\mathcal{G})$ so that attention is modulated by a function of the underlying graph structure is called *topological masking*.
Topological masking is *not* just a mechanism to sparsify attention for efficiency gains; its goal is to incorporate topological signal from the underlying graph, improving predictive performance via a useful inductive bias.

Improving efficiency with low-rank attention.
The time-complexity of Eq. 1 is $\mathcal{O}(N^{2})$, which famously limits the ability of regular transformers to scale to long inputs.
A leading approach to mitigate this and recover $\mathcal{O}(N)$ complexity is *linear* or *low-rank attention* (Katharopoulos et al., ; Choromanski et al., ).
This takes a feature mapping $\phi:\mathbb{R}^{d}\to\mathbb{R}^{m}$ with $m\ll N$ and constructs the matrices $\mathbf{\Phi}_{\mathbf{Q}}\coloneqq[\phi(\mathbf{q}_{i})]_{i=1}^{N}$ and $\mathbf{\Phi}_{\mathbf{K}}\coloneqq[\phi(\mathbf{k}_{i})]_{i=1}^{N}$, where $\mathbf{\Phi}_{\mathbf{Q},\mathbf{K}}\in\mathbb{R}^{N\times m}$.
One then defines $\mathcal{K}_{\textrm{LR}}(\mathbf{Q},\mathbf{K})\coloneqq\mathbf{\Phi}_{
\mathbf{Q}}\mathbf{\Phi}_{\mathbf{K}}^{\top}$ and computes

$$ $\textrm{Att}_{\textrm{LR}}(\mathbf{Q},\mathbf{K},\mathbf{V})\coloneqq\mathbf{D }^{-1}\left(\mathbf{\Phi}_{\mathbf{Q}}\left(\mathbf{\Phi}_{\mathbf{K}}^{\top} \mathbf{V}\right)\right),\quad\mathbf{D}\coloneqq\mathbf{\Phi}_{\mathbf{Q}} \left(\mathbf{\Phi}_{\mathbf{K}}^{\top}\mathbf{1}_{N}\right),$ (3) $$

where the parentheses show the order of computation.
Exploiting associativity, the time complexity is reduced to $\mathcal{O}(Nmd)$, albeit usually with some sacrifice in performance compared to full-rank softmax attention.
Different choices of feature maps $\phi$ exist; for example, one can take a *random feature map* that provides a Monte Carlo estimate of the softmax kernel (Performer; Choromanski et al., ; ; Likhosherstov et al., ), or a deterministic nonlinearity like $\textrm{elu}(x)+x$ (linear transformer; Katharopoulos et al., ).

Efficient masking of low-rank attention.
As Eq. 3 suggests, low-rank attention derives its speed from avoiding computing $\mathbf{A}=\mathcal{K}_{\textrm{LR}}(\mathbf{Q},\mathbf{K})$ explicitly, preferring to evaluate faster products like $\mathbf{\Phi}_{\mathbf{K}}^{\top}\mathbf{V}$ and $\mathbf{\Phi}_{\mathbf{K}}^{\top}\mathbf{1}_{N}$ first and use them to weight the query features $\mathbf{\Phi}_{\mathbf{Q}}$.
This makes the direct Hadamard product implementation of masking in Eq. 2 impossible: if the matrix $\mathbf{A}$ is never explicitly materialised in memory, one cannot pointwise multiply each of its elements.
Algorithms have been proposed to mask *implicitly* (without instantiating $\mathbf{A}$ or $\mathbf{M}$) in very simple cases, e.g. causal masking (Choromanski et al., ) or one-dimensional RPE (Liutkus et al., ; Luo et al., ).
Subquadratic implicit masking is also possible when $\mathbf{M}$ is very structured: e.g. functions of shortest path distance on grid graphs and trees (Choromanski et al., ), or of the degrees of the pair of input nodes (Chen et al., ).
See App. B for a review.
However, an $\mathcal{O}(N)$ algorithm capable of masking with flexible functions on general graphs has, until now, remained out of reach.

Remainder of manuscript. In Sec. 3.1 we parameterise topological masks $\mathbf{M}(\mathcal{G})$ as learnable functions of weighted adjacency matrices, which incorporates strong structural inductive bias into attention and boosts performance.
Sec. 3.2 shows how this can be implemented *implicitly*, without instantiating $\mathbf{A}$ or $\mathbf{M}$ in memory.
Sec. 3.3 offers a faster stochastic alternative using GRFs, which we prove is $\mathcal{O}(N)$ time and space complexity.
Sec. 4 gives experimental evaluations across data modalities, including images and point clouds.
The Appendices present proofs and extra details too long for the main text, and extended discussion about related work.

## 3 Towards linear topological masking

This section details our main results.
Consider an undirected graph $\mathcal{G}(\mathcal{N},\mathcal{E})$ where $\mathcal{N}\coloneqq\{v_{1},...,v_{N}\}$ is the set of nodes and $\mathcal{E}$ is the set of edges, with $(v_{i},v_{j})\in\mathcal{E}$ if and only if there exists an edge between $v_{i}$ and $v_{j}$ in $\mathcal{G}$.
The size of the graph $N$ corresponds to the number of tokens.
Denote by $\mathbf{M}(\mathcal{G})\in\mathbb{R}^{N\times N}$ the topological masking matrix.

### 3.1 Graph node kernels as topological masks

As remarked above, a common choice for $\mathbf{M}(\mathcal{G})_{ij}$ is a function of the shortest path distance between nodes $v_{i}$ and $v_{j}$ (Ying et al., ).
Though effective at improving performance, this requires computation for every pair of nodes so it cannot straightforwardly be applied to low-rank attention in a scalable way – except for special structured $\mathcal{G}$ where the mask becomes simple (Choromanski et al., ). Shortest path distances are also very sensitive to changes like the addition or removal of edges, and need access to the whole graph upfront.

A new topological mask: graph node kernels.
Given that we intend $\mathbf{M}(\mathcal{G})_{ij}$ to quantify some sense of topological similarity between $v_{i}$ and $v_{j}$, a more principled choice may be *graph node kernels*: positive definite, symmetric functions $k:\mathcal{N}\times\mathcal{N}\to\mathbb{R}$ mapping from pairs of graph nodes to real numbers.
These are widely used in bioinformatics (Borgwardt et al., ), community detection (Kloster and Gleich, ) and recommender systems (Yajima, ), with more recent applications in manifold learning for deep generative modelling (Zhou et al., ) and solving single- and multiple-source shortest path problems (Crane et al., ).
They are also used to model the covariance function in geometric Gaussian processes (Zhi et al., ; Borovitskiy et al., ).

Graph node kernels are often parameterised as follows.
Let $\mathbf{W}\in\mathbb{R}^{N\times N}$ be a *weighted adjacency matrix*, so that $\mathbf{W}\coloneqq[w_{ij}]_{i,j\in\mathcal{N}}$ with $w_{ij}$ nonzero if and only if $(v_{i},v_{j})\in\mathcal{E}$. A common choice is $w_{ij}=1/\sqrt{d_{i}d_{j}}$, with $d_{i}$ the degree of $v_{i}$.
Then consider the power series

$$ $\mathbf{M}_{\alpha}(\mathcal{G})\coloneqq\sum_{k=0}^{\infty}\alpha_{k}\mathbf{ W}^{k}$ (4) $$

where $(\alpha_{k})_{k=0}^{\infty}\subset\mathbb{R}$ is a set of real-valued Taylor coefficients.
This is positive definite if and only if $\sum_{k=0}^{\infty}\alpha_{k}\lambda_{i}^{k}>0\hskip 1.99997pt\forall\hskip 1.
99997pt\lambda_{i}\in\Lambda(\mathbf{W})$, where $\Lambda(\mathbf{W})$ denotes the set of eigenvalues of $\mathbf{W}$.
With different choices for $(\alpha_{k})_{k=0}^{\infty}$, Eq. 4 includes many of the most popular graph node kernels in the literature, e.g. heat, diffusion, cosine and $p$-step random walk.
Kernel learning can be implemented by optimising $\alpha_{k}$ (Reid et al., ; Zhi et al., ).

Eq. 4 provides a very effective parameterisation for topological attention masks, striking the right balance between flexibility and hard-coded structural inductive bias.
In Sec. 4 we will show that it boosts performance compared to unmasked attention.
However, explicitly computing $\mathbf{M}_{\alpha}(\mathcal{G})\in\mathbb{R}^{N\times N}$ and storing it in memory is incompatible with $\mathcal{O}(N)$ low-rank attention.
The next section will show how this can be avoided by masking *implicitly*.

### 3.2 Implicit masking with graph features

Linear attention is fast because it replaces softmax attention $\exp(\mathbf{q}_{i}^{\top}\mathbf{k}_{j})$ with a ‘featurised’ alternative $\phi(\mathbf{q}_{i})^{\top}\phi(\mathbf{k}_{j})$. Can we do the same for $\mathbf{M}_{\alpha}(\mathcal{G})$?
The following is true.

###### Remark 3.1 (Explicit $N$ -dimensional features for graph node kernels (Reid et al., 2024b ) ) .

Let $(f_{k})_{k=0}^{\infty}\subset\mathbb{R}$ denote the deconvolution of $(\alpha_{k})_{k=0}^{\infty}$, i.e. $\sum_{p=0}^{k}f_{p}f_{k-p}=\alpha_{k}\hskip 5.0pt\forall\hskip 5.0ptk$.
There exists a feature matrix $\mathbf{\Phi}_{\mathcal{G}}\in\mathbb{R}^{N\times N}$ satisfying
$\mathbf{M}_{\alpha}(\mathcal{G})=\mathbf{\Phi}_{\mathcal{G}}\mathbf{\Phi}_{
\mathcal{G}}^{\top}$,
given by $\mathbf{\Phi}_{\mathcal{G}}\coloneqq\sum_{k=0}^{\infty}f_{k}\mathbf{W}^{k}.$

This is can be shown in a few lines of algebra; see App. A.1.
We refer to the rows of the feature matrix $\mathbf{\Phi}_{\mathcal{G}}$ as *graph features*, letting $\mathbf{\Phi}_{\mathcal{G}}\eqqcolon[\phi_{\mathcal{G}}(v_{i})]_{i=1}^{N}$ with $\phi_{\mathcal{G}}(v_{i})\in\mathbb{R}^{N}$ the graph feature of node $v_{i}$.
For every kernel there trivially exists a feature map (Gretton, ), but unusually in this case it is finite-dimensional and admits a simple closed form.

Let $\otimes:\mathbb{R}^{m}\times\mathbb{R}^{N}\to\mathbb{R}^{m\times N}$ denote the *outer product* which maps a pair of vectors to a matrix, and $\textrm{vec}:\mathbb{R}^{m\times N}\to\mathbb{R}^{mN}$ denote the *vectorising* operation that flattens a matrix into a vector.
Then $\phi(\mathbf{q}_{i})\otimes\phi_{\mathcal{G}}(v_{i})$ is an $m\times N$ matrix whose $jk$th element is $\phi(\mathbf{q}_{i})_{j}\phi_{\mathcal{G}}(v_{i})_{k}$, and $\mathrm{vec}\left(\phi(\mathbf{q}_{i})\otimes\phi_{\mathcal{G}}(v_{i})\right)$ is an $mN$-dimensional vector whose $(N(j-1)+k)$th entry is $\phi(\mathbf{q}_{i})_{j}\phi_{\mathcal{G}}(v_{i})_{k}$.
Observe the following:(^1^11Recall the old adage: the dot product of outer products is the product of dot products.)

$$ $\begin{multlined}\mathrm{vec}\left(\phi(\mathbf{q}_{i})\otimes\phi_{\mathcal{G }}(v_{i})\right)^{\top}\mathrm{vec}\left(\phi(\mathbf{k}_{j})\otimes\phi_{ \mathcal{G}}(v_{j})\right)=\sum_{s=1}^{mN}\mathrm{vec}\left(\phi(\mathbf{q}_{i })\otimes\phi_{\mathcal{G}}(v_{i})\right)_{s}\mathrm{vec}\left(\phi(\mathbf{k} _{j})\otimes\phi_{\mathcal{G}}(v_{j})\right)_{s}\vspace{-3mm}\\ =\sum_{k=1}^{m}\sum_{l=1}^{N}\phi(\mathbf{q}_{i})_{k}\phi_{\mathcal{G}}(v_{i}) _{l}\phi(\mathbf{k}_{j})_{k}\phi_{\mathcal{G}}(v_{j})_{l}=\left(\phi(\mathbf{q }_{i})^{\top}\phi(\mathbf{k}_{j})\right)\left(\phi_{\mathcal{G}}(v_{i})^{\top} \phi_{\mathcal{G}}(v_{j})\right)={\mathcal{K}_{\textrm{LR}}}_{ij}{\mathbf{M}_{ \alpha}}_{ij}.\end{multlined}\mathrm{vec}\left(\phi(\mathbf{q}_{i})\otimes\phi _{\mathcal{G}}(v_{i})\right)^{\top}\mathrm{vec}\left(\phi(\mathbf{k}_{j}) \otimes\phi_{\mathcal{G}}(v_{j})\right)=\sum_{s=1}^{mN}\mathrm{vec}\left(\phi( \mathbf{q}_{i})\otimes\phi_{\mathcal{G}}(v_{i})\right)_{s}\mathrm{vec}\left( \phi(\mathbf{k}_{j})\otimes\phi_{\mathcal{G}}(v_{j})\right)_{s}\vspace{-3mm}\\ =\sum_{k=1}^{m}\sum_{l=1}^{N}\phi(\mathbf{q}_{i})_{k}\phi_{\mathcal{G}}(v_{i}) _{l}\phi(\mathbf{k}_{j})_{k}\phi_{\mathcal{G}}(v_{j})_{l}=\left(\phi(\mathbf{q }_{i})^{\top}\phi(\mathbf{k}_{j})\right)\left(\phi_{\mathcal{G}}(v_{i})^{\top} \phi_{\mathcal{G}}(v_{j})\right)={\mathcal{K}_{\textrm{LR}}}_{ij}{\mathbf{M}_{ \alpha}}_{ij}.$ (5) $$

Hence, taking $\{\mathrm{vec}\left(\phi(\left\{\mathbf{q}_{i},\mathbf{k}_{i}\right\})\otimes
\phi_{\mathcal{G}}(v_{i})\right)\}_{v_{i}\in\mathcal{N}}\subset\mathbb{R}^{mN}$ as features, we can compute masked attention implicitly, without materialising $\mathbf{A}$ and $\mathbf{M}_{\alpha}$.
Concretely, we have that:

$$ $\mathbf{\Phi}_{\mathbf{Q},\mathcal{G}}\mathbf{\Phi}_{\mathbf{K},\mathcal{G}}^{ \top}=\mathcal{K}_{\textrm{LR}}\odot\mathbf{M}_{\alpha}\quad\textrm{if}\quad \mathbf{\Phi}_{\{\mathbf{Q},\mathbf{K}\},\mathcal{G}}\coloneqq\left[\mathrm{ vec}\left(\phi(\left\{\mathbf{q}_{i},\mathbf{k}_{i}\right\})\otimes\phi_{ \mathcal{G}}(v_{i})\right)\right]_{i=1}^{N}\in\mathbb{R}^{N\times Nm}.$ (6) $$

Masked low-rank attention can then be computed using $\mathbf{\Phi}_{\{\mathbf{Q},\mathbf{K}\},\mathcal{G}}$ as the feature matrices:

$$ $\textrm{Att}_{\textrm{LR},\mathbf{M}}(\mathbf{Q},\mathbf{K},\mathbf{V}, \mathcal{G})\coloneqq\mathbf{D}^{-1}\left(\mathbf{\Phi}_{\mathbf{Q},\mathcal{G }}\left(\mathbf{\Phi}_{\mathbf{K},\mathcal{G}}^{\top}\mathbf{V}\right)\right), \quad\mathbf{D}\coloneqq\mathbf{\Phi}_{\mathbf{Q},\mathcal{G}}\left(\mathbf{ \Phi}_{\mathbf{K},\mathcal{G}}^{\top}\mathbf{1}_{N}\right).$ (7) $$

Time complexity.
Eqs 6 and 7 give a recipe for computing
masked attention implicitly in feature space, without needing to materialise $\mathbf{A}$ or $\mathbf{M}$ in memory.
However, this approach relies on computing the feature matrix $\mathbf{\Phi}_{\{\mathbf{Q},\mathbf{K}\},\mathcal{G}}\in\mathbb{R}^{N\times Nm}$ and evaluating products like
$\mathbf{\Phi}_{\mathbf{K},\mathcal{G}}^{\top}\mathbf{1}_{N}$, which still incurs $\mathcal{O}(N^{2})$ space and time complexity.
For a linear masking algorithm, we must reduce this to $\mathcal{O}(N)$.
This can be achieved by approximating the features.

###### Remark 3.1 (Explicit N 𝑁 N italic_N -dimensional features for graph node kernels (Reid et al., 2024b ) ) .

### 3.3 Faster implicit masking with graph random features

Our goal is now to replace $\mathbf{\Phi}_{\{\mathbf{Q},\mathbf{K}\},\mathcal{G}}$ by Monte Carlo estimates $\widehat{\mathbf{\Phi}}_{\{\mathbf{Q},\mathbf{K}\},\mathcal{G}}$ that satisfy

$$ $\mathbb{E}\left(\widehat{\mathbf{\Phi}}_{\mathbf{Q},\mathcal{G}}\widehat{ \mathbf{\Phi}}_{\mathbf{K},\mathcal{G}}^{\top}\right)=\mathbf{\Phi}_{\mathbf{Q },\mathcal{G}}\mathbf{\Phi}_{\mathbf{K},\mathcal{G}}^{\top}=\mathcal{K}_{ \textrm{LR}}\odot\mathbf{M}_{\alpha},\quad\widehat{\mathbf{\Phi}}_{\mathbf{K}, \mathcal{G}}^{\top}\bm{v}\sim\mathcal{O}(N)\hskip 2.84526pt\forall\hskip 2.845 26pt\bm{v}\in\mathbb{R}^{N},$ (8) $$

meaning we apply the desired topological mask in expectation but in linear time.
This can be achieved using *graph random features* (GRFs) (Choromanski, ; Reid et al., ).

At a high-level, GRFs can be understood as follows.
From Remark 3.1, the graph feature for node $v_{i}$ is given by $\phi_{\mathcal{G}}(v_{i})\coloneqq[\sum_{k=0}^{\infty}f_{k}\mathbf{W}^{k}_{ij}
]_{j=1}^{N}$.
Since $\mathbf{W}$ is an adjacency matrix, $\mathbf{W}^{k}_{ij}$ counts the number of graph walks of length $k$ between nodes $i$ and $j$, weighted by the product of traversed edge weights.
So $\phi_{\mathcal{G}}(v_{i})_{j}$ can be interpreted as a weighted sum over *all* walks of any length between nodes $v_{i}$ and $v_{j}$, with
longer contributions suppressed since edge weights are at most $1$ (by normalisation of $\mathbf{W}$).
GRFs apply importance sampling, approximating this infinite sum by a Monte Carlo estimate $\widehat{\phi}_{\mathcal{G}}(v_{i})$ using random walks.

Using GRFs. Concretely, to build the GRF $\widehat{\phi}_{\mathcal{G}}(v_{i})$ one samples $n$ random walks $\{\omega^{(i)}_{k}\}_{k=1}^{n}\subset\Omega$ out of node $v_{i}$, where $\Omega\coloneqq\left\{(v_{i})_{i=1}^{l}\,|\,v_{i}\in\mathcal{N},(v_{i},v_{i+1}
)\in\mathcal{E},l\in\mathbb{N}\right\}$ is the set of all walks.
Denote by $p(\omega)$ the probability of walk $\omega$ which is known from the sampling mechanism.
$p(\omega)$ is typically chosen so that walks randomly halt, prioritising sampling lower powers of $\mathbf{W}$ that contribute more to the kernel.
Let $\widetilde{\omega}:\Omega\to\mathbb{R}$ be a function computing the product of traversed edge weights,

$$ $\vspace{-1.5mm}\widetilde{\omega}(\omega)\coloneqq\left\{\prod_{i=1}^{\textrm{ len}(\omega)}\mathbf{W}_{\omega[i]\omega[i+1]}\textrm{ if }\textrm{len}(\omega )>1,\quad 1\textrm{ otherwise}\right\},$ (9) $$

where $\textrm{len}(\omega)$ is the number of hops and $\omega[i]$ is the node visited at the $i$th timestep. Then:

$$ $\widehat{\phi}_{\mathcal{G}}(v_{i})_{q}\coloneqq\frac{1}{n}\sum_{k=1}^{n}\sum_ {\omega_{iq}\in\Omega_{iq}}\frac{\widetilde{\omega}(\omega_{iq})f_{\textrm{len }(\omega_{iq})}}{p(\omega_{iq})}\mathbb{I}(\omega_{iq}\textrm{ prefix subwalk of }\omega_{k}^{(i)}),\quad q=1,...,N,$ (10) $$

where $\Omega_{iq}$ denotes the set of all walks between nodes $v_{i}$ and $v_{q}$, of which $\omega_{iq}$ is a member.
‘Prefix subwalk’ means that the walk $\omega_{k}^{(i)}$ beginning at node $v_{i}$ initially follows the sequence of nodes $\omega_{iq}$, then optionally continues.
This complicated notation belies a very simple algorithm: one simulates $n$ random walks out of node $v_{i}$.
At every visited node, one adds a contribution to the corresponding GRF coordinate that depends on 1) the product of traversed edge weights, 2) the probability of the subwalk and 3) the function $f$.
The result is normalised by dividing by the number of walks $n$.

GRFs give unbiased approximation of graph node kernels: $\mathbb{E}[\widehat{\phi}_{\mathcal{G}}(v_{i})^{\top}\widehat{\phi}_{\mathcal{
G}}(v_{j})]={\mathbf{M}_{\alpha}}_{ij}.$
This follows because $\mathbb{E}[\mathbb{I}(\omega_{iq}\textrm{ prefix subwalk of }\omega_{k}^{(i)})
]=p(\omega_{iq})$.
Unbiasedness is one of our required properties; we also need $\widehat{\mathbf{\Phi}}_{\{\mathbf{Q},\mathbf{K}\},\mathcal{G}}$ to support $\mathcal{O}(N)$ matrix-vector multiplication.
We will show this by providing the first known proof that GRFs are *sparse*: $\widehat{\phi}(v_{i})$ has $\mathcal{O}(1)$ nonzero entries $\hskip 1.99997pt\forall\hskip 1.99997ptv_{i}\in\mathcal{N}$.

### 3.4 Novel theoretical results for GRFs: sharp estimators with sparse features

It is intuitive that GRFs are sparse.
The walk sampler $p(\omega)$ is chosen so that walks are short, e.g. taking simple random walks with geometrically distributed lengths.
From Eq. 10, this means that $\widehat{\phi}_{\mathcal{G}}(v_{i})$ is only nonzero at nodes ‘close’ to $v_{i}$. Most of its entries are 0 0.
However, to prove that $\widehat{\mathbf{\Phi}}_{\{\mathbf{Q},\mathbf{K}\}\mathcal{G}}$ supports $\mathcal{O}(N)$ matrix-vector multiplication we need to rigorously relate this to the quality of approximation of $\mathbf{M}_{\alpha}(\mathcal{G})$.
This requires concentration inequalities for GRFs, but no such bounds were previously known. The following result is novel.

###### Theorem 3.1 (GRF exponential concentration bounds) .

Consider a graph $\mathcal{G}$ with adjacency matrix $\mathbf{W}$ and node degrees $\{d_{i}\}_{v_{i}\in\mathcal{N}}$.
Suppose we construct GRFs $\{\widehat{\phi}_{\mathcal{G}}(v_{i})\}_{v_{i}\in\mathcal{G}}$ by sampling $n$ random walks $\{\omega_{k}^{(i)}\}_{v_{i}\in\mathcal{N},\hskip 1.39998ptk\in[\![1,n]\!]}$ with geometrically distributed lengths, terminating with probability $p_{\textrm{halt}}$ at each timestep.
Let
$c\coloneqq\sum_{k=0}^{\infty}|f_{k}|(\max_{v_{i},v_{j}\in\mathcal{N}}\frac{|
\mathbf{W}_{ij}|d_{i}}{1-p_{\textrm{halt}}})^{k}.$
Suppose $c$ is finite, which is guaranteed e.g. for bounded $f$ and suitably normalised $\mathbf{W}$, or if there exists some $i_{\textrm{max}}\in\mathbb{N}$ such that $f_{i}$ vanishes for all $i>i_{\textrm{max}}$ (see App. A.2).
Then:

$$ $\textrm{\emph{Pr}}\left(|\widehat{\phi}_{\mathcal{G}}(v_{i})^{\top}\widehat{ \phi}_{\mathcal{G}}(v_{j})-{\mathbf{M}_{\alpha}}_{ij}|>t\right)\leq 2\exp\left (-\frac{t^{2}n^{3}}{2(2n-1)^{2}c^{4}}\right).$ (11) $$

*Proof sketch.* Full details are in App. A.2; here, we provide an overview.
Treat each pair of random walks $X_{k}\coloneqq(\omega_{k}^{(i)},\omega_{k}^{(j)})\in\Omega\times\Omega$, simulated out of nodes $v_{i}$ and $v_{j}$ respectively, as a random variable.
Let $k\in[\![1,n]\!]$ so we sample $n$ such pairs of walks in total.
Consider modifying one of the pairs of walks.
If $c$ is finite, the maximum load it is capable of depositing is finite so we can bound the $L_{1}$-norm of the contribution made to $\widehat{\phi}_{\mathcal{G}}(v_{i})$ and $\widehat{\phi}_{\mathcal{G}}(v_{j})$.
This allows us to bound the change in the estimator $\widehat{\phi}_{\mathcal{G}}(v_{i})^{\top}\widehat{\phi}_{\mathcal{G}}(v_{j})$ if $X_{k}$ is modified.
Applying McDiarmid’s inequality, the result follows. ∎

Thm 3.1 provides an exponential bound on the probability of a topological mask estimate deviating from its true value.
Moreover, assuming that $c$ remains constant (i.e. we fix a maximum edge weight and node degree as the graph grows), this probability is *independent of the number of nodes $N$*.
Therefore, if we fix $t$ and the probability Pr, we can solve for the minimum number of walkers $n$ to guarantee a good mask estimate with high probability, independent of graph size.
Now observe the following.

###### Lemma 3.2 (GRF sparsity) .

Consider a graph random feature $\widehat{\phi}_{\mathcal{G}}(v_{i})$ constructed using $m$ random walkers each terminating with probability $p_{\textrm{halt}}$ at every timestep.
With probability at least $1-\delta$, $\widehat{\phi}_{\mathcal{G}}(v_{i})$ will have $n\log(1-(1-\delta)^{1/n})\log(1-p_{\textrm{halt}})^{-1}$ or fewer nonzero entries.

*Proof.*
Since the walk lengths are geometrically distributed, the probability that a single walk is of length $b$ or shorter is $1-(1-p_{\textrm{halt}})^{b}$.
$n$ independent walkers are all $b$ or shorter with probability $\left(1-(1-p_{\textrm{halt}})^{b}\right)^{n}$.
Let this be equal to $1-\delta$ and solve for $b$.
At most $bn$ entries of $\widehat{\phi}_{\mathcal{G}}(v_{i})$ can then be nonzero, so the sparsity guarantee follows. ∎

Lemma 3.2 means that, if we fix the number of walkers $n$ (e.g. using the bound in Eq. 11), we can also bound GRF sparsity with high probability.
Crucially: the number of nonzero GRF entries is independent of the graph size $N$.
We conclude the following.

###### Corollary 3.3 (GRFs implement $\mathcal{O}(N)$ topological masking) .

Suppose we construct a set of graph random features $\{\widehat{\phi}_{\mathcal{G}}(v_{i})\}_{v_{i}\in\mathcal{N}}$, choosing the number of walkers $n$ as described above with $c$, $p_{\textrm{halt}}$, $\delta$, $t$ and the deviation probability $\textrm{\emph{Pr}}\left(|\phi_{\textrm{\emph{GRF}}}(v_{i})^{\top}\phi_{\textrm
{\emph{GRF}}}(v_{j})-{\mathbf{M}_{\alpha}}_{ij}|>t\right)$ fixed as hyperparameters.
Then the number of nonzero entries in $\widehat{\mathbf{\Phi}}_{\{\mathbf{Q},\mathbf{K}\},\mathcal{G}}\in\mathbb{R}^{
N\times Nm}$ is *linear* in $N$, which means that topological masking is implemented with $\mathcal{O}(N)$ time complexity.

*Proof.* Lemma 3.2 shows that the number of nonzero entries in $\widehat{\phi}_{\mathcal{G}}(v_{i})$ is independent of the number of nodes $N$.
This must also be true of $\textrm{vec}(\phi(\mathbf{k}_{i})\otimes\widehat{\phi}_{\mathcal{G}}(v_{i}))$, so the number of nonzero entries in $\widehat{\mathbf{\Phi}}_{\mathbf{K},\mathcal{G}}\coloneqq[\mathrm{vec}(\phi(
\mathbf{k}_{i})\otimes\widehat{\phi}_{\mathcal{G}}(v_{i}))]_{i=1}^{N}$ is proportional to $N$.
The same holds for the equivalent matrix for the queries, $\smash{\widehat{\mathbf{\Phi}}_{\mathbf{Q},\mathcal{G}}}$.
The time complexity of the sparse matrix-matrix multiplication $\widehat{\mathbf{\Phi}}_{\mathbf{K},\mathcal{G}}^{\top}\mathbf{V}$ is $\mathcal{O}(Nmd)$ and $\widehat{\mathbf{\Phi}}_{\mathbf{K},\mathcal{G}}^{\top}\mathbf{1}_{N}$ is $\mathcal{O}(Nm)$.
Both these time complexities are linear in $N$.
This is also the case for the subsequent computations involving $\widehat{\mathbf{\Phi}}_{\mathbf{Q},\mathcal{G}}$.
∎

Further comments on theoretical results.
We have provided the first concentration bounds for GRFs and quantification of the tradeoff between sparsity and kernel approximation quality.
These results may be of interest independent from topological masking, e.g. for applications in scalable geometric Gaussian processes (Reid et al., ; Borovitskiy et al., ).
Thm. 3.1 is remarkably general, requiring only the modest assumption about $\mathcal{G}$ and $f$ that $c$ is finite.
To guarantee $\mathcal{O}(N)$ scaling we assume that $c$ is fixed as $\mathcal{G}$ changes, which is guaranteed by e.g. fixed maximum node degree and edge weight.

###### Theorem 3.1 (GRF exponential concentration bounds) .

###### Lemma 3.2 (GRF sparsity) .

###### Corollary 3.3 (GRFs implement 𝒪 ⁢ ( N ) 𝒪 𝑁 \mathcal{O}(N) caligraphic_O ( italic_N ) topological masking) .

### 3.5 Algorithm and intuitive picture

Alg. 1 presents our method.
The results of Secs 3.1-3.4 can be intuitively understood as follows (see Fig. 2 for an overview).
Graph node kernels, which compute an infinite weighted sum of walks between pairs of input nodes, provide an effective RPE mechanism for transformers.
They capture the topological relationship between queries and keys. Using graph features, we can implement this implicitly, without needing to actually materialise the mask matrix.
Replacing these dense features by graph *random* features, a sparse approximation, we can reduce the cost of implicit masking to $\mathcal{O}(N)$ with low performance loss.
Eq. 11 gives a remarkably tight bound on the quality of mask estimation.
Physically, masking with GRFs means that nodes $v_{i}$ and $v_{j}$ only attend to one another if members of their respective ensembles of random walks *hit*, such that they visit some of the same nodes.
This incorporates a strong structural inductive bias from $\mathcal{G}$.
It upweights crucial interactions between nearby nodes, whilst also sampling longer-range attention with lower probability.

Figure: Figure 2: Visual overview. A graph $\mathcal{G}$ (left) has a weighted adjacency matrix $\mathbf{W}$ (centre left). A learnable power series $\mathbf{M}_{\alpha}(\mathcal{G})\coloneqq\sum_{i=0}^{\infty}\alpha_{i}\mathbf{ W}^{i}$ is an effective topological mask or graph RPE (centre right). $\mathbf{M}_{\alpha}(\mathcal{G})$ can be efficiently approximated using graph random features (centre top), which perform importance sampling of halting random walks. The feature $\widehat{\phi}_{\mathcal{G}}(v_{i})$ is only nonzero at entries visited by the ensemble of walks beginning at $v_{i}$. In Thm. 3.1, we prove that the number of such entries is $\mathcal{O}(1)$ whilst still accurately estimating $\mathbf{M}_{\alpha}(\mathbf{W})_{ij}$ with high probability, so GRFs are sparse. This unlocks $\mathcal{O}(N)$ topological masking. Note that $\widehat{\phi}_{\mathcal{G}}(v_{i})^{\top}\widehat{\phi}_{\mathcal{G}}(v_{j})$ is only nonzero if the features are nonzero at some of the same coordinates, which happens if their respective ensembles of walks ‘hit’. This incorporates a strong structural inductive bias. The equations on the right formalise our method mathematically.

Distributed computation, expressivity and GATs.
Alg. 1 is simple to distribute for massive graphs that cannot be stored on a single machine; to simulate random walks one only needs a node’s immediate neighbours rather than all of $\mathcal{G}$ upfront.
This desirable property is not shared by masking algorithms that rely on the fast Fourier transform (Luo et al., ; Choromanski et al., ).
Moreover, our method can distinguish graphs identical under the $1$-dimensional Weisfeiler Lehman graph isomorphism heuristic (with colours replaced by node features) because their adjacency matrices and hence masks differ.
In this sense, GRF transformers are *more expressive than standard graph neural networks and unmasked transformers*, which notoriously fail this test (Morris et al., ; Xu et al., ).
Finally, we remark that our algorithm can be interpreted as a stochastic extension of graph attention networks (GATs; Veličković et al., ), with a greater focus on linear attention and scalability.
Tokens still attend to a subset of other tokens informed by the topology of $\mathcal{G}$, but importance sampling random walks permits interactions beyond nearest-neighbour.
Incorporating learnable parameters $(f_{i})_{i=0}^{i_{\textrm{max}}}$ lets us estimate a principled function of $\mathbf{W}$.

Figure: Algorithm 1 $\mathcal{O}(N)$ topologically-masked attention for general graphs

## 4 Experiments

In this section, we test our algorithms for topological masking with GRFs.
We consider data modalities with different graph topologies: images and point clouds.

Figure: Figure 3: Number of FLOPs vs. number of graph nodes for softmax attention, linear attention, and linear attention with GRF topological masking (ours).
Refer to caption: x1.png

Time complexity.
For a hardware-agnostic comparison, we first compute the total number of FLOPs for evaluating (i) unmasked softmax, (ii) unmasked linear and (iii) GRF-masked linear attention for graphs of different sizes $N$.
Fig. 3 shows the results.
For concreteness we use $1$-dimensional grid graphs, though the theoretical results in Sec. 3.4 are independent of $\mathcal{G}$ (provided $c$ is upper bounded).
We take $n=4$ random walkers per node with termination probability $p_{\textrm{halt}}=0.5$.
The latent dimension is $d=8$ and the feature size is $m=8$. In the masked case we show the *mean* number of FLOPs over $10$ seeds since GRF sparsity is itself a random variable.
As anticipated, the time complexity of our mechanism is $\mathcal{O}(N)$; there is a constant multiplicative cost for topological masking.
We clearly improve upon the time complexity of $\mathcal{O}(N^{2})$ softmax attention.

ViTs: grid graphs and image data.
Next, we use our topological masking algorithm to improve the vision transformer (ViT; Dosovitskiy et al., ).
The underlying graph $\mathcal{G}$ is a square grid of size num_x_patches $\times$ num_y_patches, with nodes connected by an edge if they are neighbours in $x$ or $y$.
Since the graph is fixed, random walks can be pre-computed and frozen for each attention head and layer before training and inference. Full architectural and training details are in App. C.
All our ViT models also endow tokens with learned additive absolute position embeddings.

Table 1 shows the final test accuracies for ImageNet (Deng et al., ), iNaturalist2021 (Horn et al., ) and Places365 (Zhou et al., ).
For the *slow* (superlinear) variants, we include (i) unmasked softmax and (ii) Toeplitz-masked linear (Choromanski et al., ) attention.
Toeplitz can only be used for this specific, structured graph topology.
We also include (iii) $\mathbf{M}_{\alpha}(\mathcal{G})$-masked linear, which takes an explicit power series in $\mathbf{W}$ (Eq. 4) to directly modulate the $N\times N$ attention matrix.
This is equal to our Monte Carlo GRFs algorithm in the $n\to\infty$ limit, so it indicates the asymptotic performance of GRFs with very many random walkers.
For the *fast* (linear) variants, we include (iv) unmasked linear and (v) GRF-masked linear attention.
Our algorithm gives strong improvements above the unmasked baseline, by $\mathbf{+3.7\%}$, $\mathbf{+2.2\%}$ and $\mathbf{+0.5\%}$ respectively.
Remarkably, it is competitive with much more expensive variants: GRFs often match or beat Toeplitz.
For ImageNet and iNaturalist2021, it even substantially narrows the gap with $\mathcal{O}(N^{2})$ softmax.

**Table 1: Final transformer test accuracies across attention masking variants and datasets. Our algorithms are shown in bold. Topological masking with GRFs is the best $\mathcal{O}(N)$ mechanism, matching or beating more expensive alternatives.**
|  | Variant | Time comp. | ImageNet | iNaturalist2021 | Places365 (Small) |
| --- | --- | --- | --- | --- | --- |
|  |  |  | 1M | 2.7M | 1.8M |
| SLOW | Unmasked softmax | $\mathcal{O}(N^{2})$ | 0.741 | 0.699 | 0.567 |
| Toeplitz-masked linear | $\mathcal{O}(N\log N)$ | 0.733 | 0.674 | 0.550 |  |
| $\mathbf{M}_{\alpha}(\mathcal{G})$-masked linear | $\mathcal{O}(N^{2})$ | 0.741 | 0.692 | 0.551 |  |
| FAST | Unmasked linear | $\mathcal{O}(N)$ | 0.693 | 0.667 | 0.543 |
| GRF-masked linear | $\mathcal{O}(N)$ | 0.730 | 0.689 | 0.548 |  |

Ablation studies.
In App. C.2, we perform ablation studies for the ViT experiments.
We find that predictive performance improves with the number of walkers $n$ because the accuracy of the Monte Carlo estimate $\widehat{\mathbf{M}}_{\alpha}(\mathcal{G})$ increases.
Moreover, importance sampling walks is crucial; if we fail to upweight improbable walks, the gains vanish.
Lastly, the inductive bias from parameterising $\mathbf{M}_{\alpha}(\mathcal{G})$ as a function of $\mathbf{W}$ is key.
Fully learnable graph features of the same dimensionality $N$, an impractical and expensive alternative to our method, give much smaller gains.

PCTs: predicting particle dynamics for robotics.
As a final task, we use our algorithm to model *high-density visual particle dynamics* (HD-VPD; Whitney et al., ; ), predicting the physical evolution of scenes involving a bi-manual robot using a point cloud transformer (PCT).
Images from a pair of cameras are unprojected to a set of 3D particles.
Their interactions, conditioned on specific robot actions, are modelled by a dynamics network.
We predict updates to the latent point cloud, render the result to an image with a conditional NeRF (Mildenhall et al., ), and train end-to-end with video prediction loss.
Such learned dynamics models are attracting growing interest in robotics as a more flexible alternative to physical simulators, but to scale to the massive point clouds needed for detailed predictions efficiency and effective structural inductive biases are key.

For the dynamics network, we use an *Interlacer*: a PCT which alternates unmasked linear attention and message passing-style updates between neighbouring particles (Whitney et al., ).
This architecture is twice as fast as GNNs with the same prediction quality, and can handle $4\times$ bigger point clouds.
To test the ability of GRFs to efficiently capture local topological structure, we replace message passing by GRF-masked linear attention, taking query/key features $\phi(\cdot)=\textrm{ReLU}(\cdot)$.
Implementation details for HD-VPD, including the robotic arm and camera configurations, are the same as reported by Whitney et al. (); see App. C.3 for discussion.
To disambiguate, we will refer to our model as the *GRF Interlacer* and Whitney et al.’s as the *message passing (MP) Interlacer*.

Figure: Figure 4: Rendered rollouts. NeRF renderings of the predictive dynamics of a bimanual Kuka robot, conditioned on the initial scene and sequence of robot actions. We show four tasks: using a dustpan and brush, lifting a can, moving a green block, and dropping a can. GRF Interlacers model point cloud dynamics more accurately, so the predicted frame rendings are closer to the ground truth.
Refer to caption: extracted/5928799/iclr2024/final_renderings/2/00001.jpg

Figure: Figure 5: Accuracy comparison. Structural similarity index measure (SSIM) between ground truth camera frames and predictive NeRF renderings after $100$k training steps, plotted against rollout timestep. Higher is better. GRFs improve the accuracy of dynamics prediction compared to the baselines.
Refer to caption: x2.png

In contrast to our previous experiments, for HD-VPD the underlying graph $\mathcal{G}$ is no longer fixed so walks cannot be pre-computed.
Instead, we construct them on the fly, using an approximate $k$-nearest neighbours solver to get adjacency lists of length $k=3$ for each node.
We use repelling random walks for variance reduction (Reid et al., ), sampling $3$ hops.
For easy integration with existing Interlacer code, we use *asymmetric* GRFs (App. D; Reid et al., ).
We consider point clouds with $N=32768$ particles, which is too many to explicitly instantiate $\mathbf{A}$ or $\mathbf{M}$ in memory.

Fig. 4 shows example NeRF renderings of test set predictive dynamics for four rollouts.
Since the model is deterministic, predictions blur with long time horizons.
Incorporating a diffusion head (Song et al., ) is left to future work.
Fig. 5 shows predicted video frame accuracy (image SSIM) vs. rollout timestep for each model.
GRFs outperform the vanilla transformer and MP Interlacer baselines, modelling the particle dynamics more faithfully.
We include examples of predictive videos in the supplementary material.

## 5 Conclusion

We have presented the first linear algorithm for transformer topological masking on general graphs, based on sampling random walks.
It uses graph random features, for which we have proved the first known concentration inequalities and sparsity guarantees.
Our algorithm brings strong performance gains across data modalities, including on massive point clouds with $>30$k particles for which computational efficiency is essential.
Future work might include improving hardware and libraries for high-dimensional sparse linear algebra to fully benefit from the guaranteed complexity improvements, and investigating whether dimensionality reduction techniques like the Johnson-Lindenstrauss transform (Dasgupta et al., ) can be applied to GRFs for further speedups.

## 6 Ethics and reproducibility

Ethics statement: Our work is foundational with no immediate ethical concerns apparent to us.
However, increases in scalability provided by improvements to efficient transformers could exacerbate existing and incipient risks of machine learning, from bad actors or as unintended consequences.

Reproducibility statement:
We have made every effort to ensure the work’s reproducibility.
The core algorithm is presented clearly in Alg. 1.
Theoretical results are proved with accompanying assumptions in Sec. 3.4 and App. A.2.
Proof sketches are also included for clarity.
We will make source code available once accepted.
Apart from the robotics dataset, all datasets are standard and freely available online.
Exhaustive experimental details about the training and architectures are reported in App. C; see especially Table 2.

## 7 Contributions and acknowledgements

Relative contributions.
IR and KC conceptualised the project and co-designed the algorithms.
IR proved the theoretical results, implemented the GRFs code, ran the Interlacer experiments and wrote the text.
KAD was crucially involved throughout, running the ViT and ViViT experiments and guiding the project’s direction.
DJ and WW helped set up the Interlacer experiments.
AA, JA, AB, MJ, AM, DR, CS, RT, RW and AW provided thoughtful feedback on the project and manuscript, helped set up experimental infrastructure, and made Cambridge and Google great places to work.

Acknowledgements and funding.
IR acknowledges support from a Trinity College External Studentship.
This work was completed whilst employed as a student researcher at Google.
IR sincerely thanks Silvio Lattanzi for making this possible.
RET is supported by Google, Amazon, ARM, Improbable, EPSRC grant EP/T005386/1, and the EPSRC Probabilistic AI Hub (ProbAI, EP/Y028783/1).
AW acknowledges support from a Turing AI fellowship under grant EP/V025279/1 and the Leverhulme Trust via CFI.

## Appendix A Proofs

### A.1 Proof of Remark 3.1

This result was presented by Reid et al. []; a short demonstration is included below for the reader’s convenience.
Observe that

$$ $\begin{multlined}\mathbf{M}_{\alpha}(\mathcal{G})_{ij}\coloneqq\sum_{k=0}^{ \infty}\alpha_{k}\mathbf{W}^{k}_{ij}\overset{(1)}{=}\sum_{v=1}^{N}\sum_{k=0}^{ \infty}\alpha_{k}\mathbf{W}^{k-p}_{iv}\mathbf{W}^{p}_{vj}\overset{(2)}{=}\sum_ {v=1}^{N}\sum_{k=0}^{\infty}\sum_{p=0}^{k}f_{k-p}f_{p}\mathbf{W}^{k-p}_{iv} \mathbf{W}^{p}_{vj}\\ \overset{(3)}{=}\sum_{v=1}^{N}\left(\sum_{k_{1}=0}^{\infty}f_{k_{1}}\mathbf{W} ^{k_{1}}_{iv}\right)\left(\sum_{k_{2}=0}^{\infty}f_{k_{2}}\mathbf{W}^{k_{2}}_{ jv}\right)^{\top}\overset{(4)}{=}\mathbf{\Phi}_{\mathcal{G}}\mathbf{\Phi}_{ \mathcal{G}}^{\top}\end{multlined}\mathbf{M}_{\alpha}(\mathcal{G})_{ij} \coloneqq\sum_{k=0}^{\infty}\alpha_{k}\mathbf{W}^{k}_{ij}\overset{(1)}{=}\sum_ {v=1}^{N}\sum_{k=0}^{\infty}\alpha_{k}\mathbf{W}^{k-p}_{iv}\mathbf{W}^{p}_{vj} \overset{(2)}{=}\sum_{v=1}^{N}\sum_{k=0}^{\infty}\sum_{p=0}^{k}f_{k-p}f_{p} \mathbf{W}^{k-p}_{iv}\mathbf{W}^{p}_{vj}\\ \overset{(3)}{=}\sum_{v=1}^{N}\left(\sum_{k_{1}=0}^{\infty}f_{k_{1}}\mathbf{W} ^{k_{1}}_{iv}\right)\left(\sum_{k_{2}=0}^{\infty}f_{k_{2}}\mathbf{W}^{k_{2}}_{ jv}\right)^{\top}\overset{(4)}{=}\mathbf{\Phi}_{\mathcal{G}}\mathbf{\Phi}_{ \mathcal{G}}^{\top}$ (13) $$

where $\mathbf{\Phi}_{\mathcal{G}}\coloneqq\sum_{k=0}^{\infty}f_{k}\mathbf{W}^{k}$.
Here, step (1) splits the $\mathbf{W}^{k}$ into a product of $\mathbf{W}^{k-p}$ and $\mathbf{W}^{p}$ with $p\in\mathbb{N}$; step (2) replaces $\alpha_{k}$ by the convolution $\alpha_{k}\eqqcolon\sum_{p=0}^{k}f_{p}f_{k-p}$; step (3) changes the summation indices; step (4) arrives at the desired form.

### A.2 Proof of Thm. 3.1 (novel)

The following concentration bounds are new.
Referring back to Eq. 10, recall that graph random features are computed by

$$ $\widehat{\phi}_{\mathcal{G}}(v_{i})_{q}\coloneqq\frac{1}{n}\sum_{k=1}^{n}\psi( \omega_{k}^{(i)})_{q},\quad q=1,...,N,$ (14) $$

where we define the *projection vector* $\psi(\cdot):\Omega\to\mathbb{R}^{N}$ such that

$$ $\psi(\omega_{k}^{(i)})_{q}\coloneqq\sum_{\omega_{iq}\in\Omega_{iq}}\frac{ \widetilde{\omega}(\omega_{iq})f_{\textrm{len}(\omega_{iq})}}{p(\omega_{iq})} \mathbb{I}(\omega_{iq}\textrm{ prefix subwalk of }\omega_{k}^{(i)}).$ (15) $$

$\Omega_{iq}$ is the set of walks between nodes $v_{i}$ and $v_{q}$; $p(\omega_{iq})$ is the probability of a particular walk $\omega_{iq}$, known from the sampling mechanism; $\widetilde{\omega}(\omega_{iq})$ is a function that returns the product of weights of edges traversed by $\omega_{iq}$ (Eq. 9); $\textrm{len}(\omega_{iq})$ is the length of $\omega_{iq}$; and $f:\mathbb{N}\to\mathbb{R}$ is a sequence of a real numbers that determines the graph node kernel to be estimated.
Following convention, we will assume that we sample random walks with *geometrically distributed lengths*, halting with probability $p_{\textrm{halt}}$ at each timestep, though in principle different importance sampling schemes are possible.

Suppose we are to estimate a single entry of the graph node kernel, $\mathbf{M}_{\alpha}(\mathcal{G})_{ij}$.
We do so by taking

$$ $\widehat{\mathbf{M}}_{\alpha ij}\coloneqq\widehat{\phi}_{\mathcal{G}}(v_{i})^{ \top}\widehat{\phi}_{\mathcal{G}}(v_{j})=\frac{1}{n^{2}}\sum_{k_{1}=1}^{n}\sum _{k_{2}=1}^{n}\psi(\omega_{k_{1}}^{(i)})^{\top}\psi(\omega_{k_{2}}^{(j)})$ (16) $$

where $\mathbb{E}(\widehat{\mathbf{M}}_{\alpha ij})={\mathbf{M}}_{\alpha ij}$ since the estimate is unbiased.
We sample $n$ independent random walks starting from every node of the graph.

Let us now consider a *single* pair of random walks from $v_{i}$ and $v_{j}$ respectively.
Denote this by the random variable

$$ $X_{k}\coloneqq(\omega_{k}^{(i)},\omega_{k}^{(j)})\in\Omega\times\Omega,$ (17) $$

where $n$ such random variables $\{X_{k}\}_{k=1}^{n}$ are used in total.
Suppose we modify $X_{k}$.
Under mild assumptions, we can prove that the change in the estimator $|\Delta\widehat{\mathbf{M}}_{\alpha ij}|$ is *bounded*, which will enable us to use derive concentration inequalities.

This can be understood as follows.
To begin, we will bound the $L_{1}$-norm of $\psi(\omega_{k}^{(i)})$.
Observe that

$$ $\|\psi(\omega_{k}^{(i)})\|_{1}\leq\max_{\omega_{k}^{(i)}\in\Omega}\sum_{\omega \textrm{ p.s. }\omega_{k}^{(i)}}\left|\frac{\widetilde{\omega}(\omega)f_{ \textrm{len}(\omega)}}{p(\omega)}\right|\leq\sum_{k=0}^{\infty}|f_{k}|\left( \max_{v_{i},v_{j}\in\mathcal{N}}\frac{|\mathbf{W}_{ij}|d_{i}}{1-p_{\textrm{ halt}}}\right)^{k}\eqqcolon c,$ (18) $$

which corresponds to the case of an infinite length walk that happens almost never.
Here, $\omega\textrm{ p.s. }\omega_{k}^{(i)}$ means that $\omega$ is a ‘prefix subwalk’ of $\omega_{k}^{(i)}$, so the first $\textrm{len}(\omega)$ nodes of $\omega_{k}^{(i)}$ are exactly the nodes of $\omega$.

Suppose that $c$ is finite.
This can be guaranteed for bounded $f$ by normalisation of $\mathbf{W}$, multiplying by a scalar so that $|\max_{v_{i},v_{j}\in\mathcal{N}}\frac{\mathbf{W}_{ij}d_{i}}{1-p}|<1$.
More pragmatically, it can be achieved by stipulating that there exists some integer $i_{\textrm{max}}\in\mathbb{N}$ such that $f_{i}=0\hskip 1.99997pt\forall\hskip 1.99997pti>i_{\textrm{max}}$, so we only consider walks up to a certain length.
This is a natural choice in experiments since we want to keep the set of learnable parameters finite.
It also coincides with the intuition that very long walks should not reflect the relationship between the nodes of a graph of finite size.

We have seen that $0<\|\psi(\omega_{k}^{(i))}\|_{1}\leq c$.
It follows that $-c^{2}\leq\psi(\omega_{k_{1}}^{(i)})^{\top}\psi(\omega_{k_{2}}^{(j)})\leq c^{2
}\hskip 1.99997pt\forall\hskip 1.99997pt(k_{1},k_{2},i,j)$, whereupon the change in $\widehat{\mathbf{M}}_{\alpha ij}$ if we modify $X_{k}$ obeys

$$ $|\Delta\widehat{\mathbf{M}}_{\alpha ij}|\leq 2\frac{2n-1}{n^{2}}c^{2}.$ (19) $$

With straightforward application of McDiarmid’s inequality [Doob, ], it follows that

$$ $\textrm{{Pr}}\left(|\widehat{\phi}_{\mathcal{G}}(v_{i})^{\top}\widehat{\phi}_{ \mathcal{G}}(v_{j})-{\mathbf{M}_{\alpha}}_{ij}|>t\right)\leq 2\exp\left(-\frac {t^{2}n^{3}}{2(2n-1)^{2}c^{4}}\right),$ (20) $$

completing the proof. ∎

As noted in the main text, this is a remarkably general and powerful result.
It is the first known concentration inequality for GRFs.

## Appendix B Extended discussion on related work

In this appendix, we give a detailed explanation of the relationship to previous work.

- 1.
Stochastic position encoding (SPE) [Liutkus et al., ].
To our knowledge, this is the earliest RPE strategy compatible with $\mathcal{O}(N)$ transformers, incorporating information about the ‘lag’ between tokens in the linear attention setting by drawing samples from random processes with a prescribed covariance structure.
The time complexity is $\mathcal{O}(N)$ but the method is only applicable to sequence data, i.e. the very special case that $\mathcal{G}$ is a $1$-dimensional grid.
- 2.
Toeplitz masks and the fast Fourier transform (FFT) [Luo et al., ].
Again considering sequence data, suppose that RPE between a query $\bm{q}_{i}$ and a key $\bm{k}_{j}$ is impemented by adding a bias term $b_{j-i}$ to the dot product before exponentiating.
Then the corresponding mask $\mathbf{M}_{ij}=e^{b_{j-i}}$ is *Toeplitz* (constant on all diagonals).
From numerical linear algebra, multiplying a Toeplitz matrix and a vector can be achieved with $\mathcal{O}(N\log N)$ time complexity rather than $\mathcal{O}(N^{2})$, using the FFT.
This is subquadratic, though not linear, and again only specific $\mathcal{G}$ are considered.
- 3.
Graphormer [Ying et al., ].
In this paper, positional encodings based on the shortest path distances between the nodes of the graph (spatial encoding) and their respective degrees (centrality encoding) are seen to substantially improve transformer performance, closing the gap with graph neural networks.
There is no attempt to incorporate these techniques with linear attention, but the paper provides strong evidence of the effectiveness of graph-based attention modulation (topological masking).
- 4.
Block Toeplitz matrices and differential equations on graphs [Choromanski et al., ].
This paper proves that subquadratic matrix-vector multiplication is sufficient for masking to be applied efficiently, and suggests a number of algorithms to achieve this in special cases.
•
*Block Toeplitz matrices.* This method extends the work of Luo et al. [] to show that the analogous mask on a more general $d$-dimensional grid is $d$-level *block* Toeplitz.
Hence, this mask parameterisation also supports $\mathcal{O}(N\log N)$ masking, but now for a slightly more general class of graphs.
•
*Affine mappings on forests.* If $\mathcal{G}$ is a forest and $\mathbf{M}_{ij}=\exp(\tau(\textrm{dist}(v_{i},v_{j})))$, where $\textrm{dist}(v_{i},v_{j})$ is the shortest path distance between nodes $v_{i}$ and $v_{j}$ and $\tau$ is an affine mapping, then topological masking can be implemented in $\mathcal{O}(N)$.
The algorithm uses dynamic programming techniques for rooted trees.
Notwithstanding its speed, the flexibility of this parameterisation is limited and only particular topologies can be treated.
•
*Heat kernels on hypercubes.* Kondor and Lafferty [] derived a closed form for the heat kernel on a hypercube, which also turns out give a Toeplitz mask that can be applied in $\mathcal{O}(N\log N)$.
This is a special instantiation of our power series masks $\mathbf{M}_{\alpha}(\mathbf{W})$ for a specific choice of $\mathcal{G}$.
The authors did not empirically test heat kernels; we found them to perform very poorly compared to our more general, faster method.
•
*Random walk graph node kernels.*
Here, the authors define an ad-hoc graph kernel by taking the dot product between two ‘frequency vectors’, recording the nodes visited by the random walks beginning at the two nodes of interest.
This method is the closest to our approach, but with some crucial differences: (i) we use *importance* sampling of random walks, which we find to be crucial for good performance (see App. C); (ii) we use *learnable* weights that depend of the length of random walk, permitting flexible upweighting or downweighting of longer- and shorter-range interactions; (iii) our method provides a Monte Carlo estimate of a learnable function of a weighted adjacency matrix, which enjoys certain regularisation properties [Reid et al., ].
More importantly, we are able to derive strong concentration bounds that *prove* that our method is $\mathcal{O}(N)$, whereas Choromanski et al. [] do not explicitly guarantee any time complexities.
- 5.
Centrality degree encoding [Chen et al., ].
This paper takes the mask $\mathbf{M}_{ij}=\sin(\frac{\pi}{4}[z(d_{i})+z(d_{j})])$, where $d_{\{i,j\}}$ is the degree of node $v_{\{i,j\}}$ and $z:\mathbb{N}\to(0,1)$ is a learnable function thereof.
This can be implemented with linear transformers on general graphs in $\mathcal{O}(N)$ time complexity, but since it is completely insensitive to the relative positions of the nodes in the graph it cannot be considered a general topological masking mechanism.
- 6.
Graph attention networks [Veličković et al., ].
This seminal work incorporates masked self-attention into graph neural networks, with each node attending only to its neighbours.
The focus is on developing a convolution-style neural network for graph-structured data rather than on efficient masking of linear attention transformers.
Our mask parameterisation is more general than this hard nearest-neighbours example, but shares the finding that modulating attention by local topological structure boosts performance.

## Appendix C Further experiments and details

Here, we include extra experiments and details too long for the main text.

### C.1 Architectures, hyperparameters and training details

Table 2 gives details for the ViT experiments.

**Table 2: Architecture, hyperparameters and training details for ViT experiments.**
|  | ImageNet | iNaturalist2021 | Places365 (Small) |
| --- | --- | --- | --- |
|  | 1M | 2.7M | 1.8M |
| Num. layers | 12 | 12 | 12 |
| Num. heads | 12 | 12 | 12 |
| Num. patches | $16\times 16$ | $16\times 16$ | $16\times 16$ |
| Hidden size | 768 | 768 | 768 |
| MLP dim. | 3072 | 3072 | 3072 |
| Optimiser | Adam | Adam | Adam |
| Epochs | 90 | 5 | 5 |
| Base learning rate | $3\times 10^{-3}$ | $3\times 10^{-3}$ | $3\times 10^{-3}$ |
| Final learning rate | $1\times 10^{-5}$ | $1\times 10^{-5}$ | $1\times 10^{-5}$ |
| Learning rate schedule | Linear warmup ($10^{4}$ steps), constant, cosine decay |  |  |
| Batch size | 4096 | 64 | 64 |
| $\phi(\cdot)$ | $\textrm{ReLU}(\cdot)$ | $\textrm{ReLU}(\cdot)$ | $\textrm{ReLU}(\cdot)$ |
| Pretrained? | ✗ | ✓ | ✓ |
| Num. walks | 100 | 20 | 20 |
| $p_{\textrm{halt}}$ | 0.1 | 0.1 | 0.1 |
| Max. walk length ($i_{\textrm{max}}$) | 100 | 10 | 10 |

### C.2 ViT ablations

Here, we run extra experiments to further evidence the findings reported in the main text. Unless explicitly stated, architectures and training details are as reported in Table 2.

Number of walkers $n$.
We construct GRFs using $\{1,10,100,1000\}$ walkers with a termination probability $p_{\textrm{halt}}=0.5$, retraining from scratch using the respective topological masks.
See Fig. 6.
Larger $n$ reduces the mask variance and improves predictive performance, but also makes the features $\{\widehat{\phi}_{\mathcal{G}}(v_{i})\}_{v_{i}\in\mathcal{N}}$ less sparse so increases computational cost.
Since the estimator is unbiased, the $n\to\infty$ limit gives dense features $\{{\phi}_{\mathcal{G}}(v_{i})\}_{v_{i}\in\mathcal{N}}$ whose dot product is exactly a Taylor mask $\mathbf{M}_{\alpha}(\mathcal{G})$.

Importance sampling.
The crux of our approach is modulating attention by an unbiased Monte Carlo approximation of a function of a weighted adjacency matrix $\mathbf{W}$.
Using GRFs, this involves sampling random walks which we reweight by (i) their probability under the sampling mechanism, (ii) the product of traversed edge weights $\widetilde{\omega}(\omega)$, and (iii) a learnable function $f$ (see Eq. 10).
To investigate the importance of this principled sampling mechanism, we ablate (i) and (ii) and compare with an ad-hoc empirical random walk kernel defined by the features

$$ $\phi_{\textrm{ad-hoc}}(v_{i})_{q}\coloneqq\frac{1}{n}\sum_{k=1}^{n}\sum_{ \omega_{iq}\in\Omega_{iq}}f_{\textrm{len}(\omega_{iq})}\mathbb{I}(\omega_{iq} \textrm{ prefix subwalk of }\omega_{k}^{(i)}),\quad q=1,...,N.$ (21) $$

Eq. 21 still simulates random walks and records their destinations, giving an empirical measure of the overlap between respective nodes’ neighborhoods.
However, in absence of dependence on $\widetilde{\omega}(\omega)$ and $p(\omega)$ this cannot be interpreted as a MC estimate of a function of $\mathbf{W}$.
This is similar to the approach employed by Choromanski et al. [] but with extra learnability via $f$.

Comparing ViT performance with $n=100$ random walks and $p_{\textrm{halt}}=0.5$, the results are striking; see Table 3.
In contrast to our method, the mask defined by Eq. 21 is *worse than the unmasked baseline*.
This shows the crucial importance of the theory in Sec. 3.4.
Intuitively, without upweighting long, improbable walks by $p(\omega)^{-1}$ it is difficult to capture long-range dependencies in the topological mask.

Fully learnable features.
As a final check, we endow each graph with a fully learnable, dense $N$-dimensional feature $\phi_{\textrm{dense, learned}}(v_{i})\in\mathbb{R}^{N}$, relaxing the parameterisation as a function of a weighted adjacency matrix $\mathbf{M}_{\alpha}(\mathcal{G})$.
This uses many more parameters and is $\mathcal{O}(N^{2})$ rather than $\mathcal{O}(N)$, but it technically includes GRFs as a special case.
In making this choice we have removed all structural inductive bias from $\mathcal{G}$, allowing the model to implicitly learn a masking kernel without any constraints (except feature dimension).
The model performs badly, providing only a very slight improvement over the unmasked linear baseline.
See Table 3.

Figure: Figure 6: Final test accuracy vs. number of walkers $n$. Monte Carlo mask variance drops with the number of walks so predictions improve.
Refer to caption: x3.png

**Table 3: Ablation studies. Importance sampling of random walks and mask parameterisation as a function of $\mathbf{W}$ are crucial for good performance.**
| Variant | Test acc. |
| --- | --- |
| GRF-masked | 0.730 |
| Unmasked | 0.693 |
| Ad-hoc masked (Eq. 21) | 0.689 |
| Fully learnable features | 0.696 |

### C.3 HD-VPD experimental details

Here, we describe the high density visual particle dynamics (HD-VPD) experiment from Sec. 4 in more detail.
The setup closely follows that of [Whitney et al., ], to which we direct the interested reader for an exhaustive information about hardware and modelling techniques.
We will provide a brief overview for convenience.

Dynamics modelling as video prediction.
The HD-VPD stack involves three steps: (1) encode the input images as 3D point clouds; (2) predict updates to the point clouds, conditioned on robot actions; (3) render the predicted new state into an image.
We train end-to-end using video prediction loss.
In more detail:

- 1.
*Encode the input images.* The model receives two RGB-D (colour channels plus depth) images at times $t-1$, $t-2$.
Each RGB image is encoded into per-pixel features using a U-Net [Ronneberger et al., ], and unprojected using known camera intrinsics and extrinsics to a point cloud $(\bm{x}_{i},\bm{f}_{i})_{i=1}^{N_{\textrm{pixels}}}$ with $\bm{x}\in\mathbb{R}^{3}$ and $\bm{f}\in\mathbb{R}^{16}$.
For every timestep, the particles across different cameras are merged and subsampled uniformly at random to get $N=32768$.
The robot actions at timesteps $\{t-2,t-1,t\}$, representing the robot joints and where they plan to move, are also encoded as *kinematic particles*.
- 2.
*Predict the point cloud dynamics.*
With an *Interlacer* (see below), use the point clouds and kinematic particles to predict $(\Delta\bm{x}_{i},\Delta\bm{f}_{i})_{i=1}^{N}$, and hence the updated state $(\bm{x}_{i}^{t-1}+\Delta\bm{x}_{i},\bm{f}_{i}^{t-1}+\Delta\bm{f}_{i})_{i=1}^{N}$.
- 3.
*Render the new state to an image.*
Use a ray-based renderer similar to Point-NeRF [Xu et al., ] to render an image for the updated point cloud.

The dynamics model can be recursively applied to make predictions multiple rollout timesteps into the future.
To train, we sample a subset of rays to render at each timestep and compute the $L_{2}$ loss between the predicted and observed RGB values.

Interlacers as PCTs.
The Interlacer is a type of point cloud transformer (PCT) designed to scale to tens of thousands of particles.
It alternates linear-attention transformer layers that capture global dependencies between tokens with local *neighbour-attender* layers that model fine-grained geometric structure (see Fig. 3 by Whitney et al. []).
This is reported to outcompete both graph neural networks and vanilla linear-attention transformers.
For the neighbour-attender, the authors use a message passing-type algorithm that updates the features of a subset of anchor particles depending on the $k$-nearest neighbours’ position and feature vectors, and then updates the remaining particles depending on their closest anchor.
In Sec. 4, we show that this can be replaced by GRF-masked linear attention. This also captures local structure and obviates expensive message passing.

Training and architecture details.
This exactly follows App. E by Whitney et al. [], apart from:

- 1.
We average over $5$ seeds, discarding one where training does not converge.
- 2.
For our algorithm, we use *asymmetric* GRFs. (Sec. D).
We use $n=3^{3}=27$ repelling random walks [Reid et al., ].
Instead of terminating, we sample $3$ hops deterministically since the computational bottleneck is finding the $k$-nearest neighbours list rather than sampling lengths once it has been computed.
We take $\phi(\cdot)=\textrm{ReLU}(\cdot)$ and compute GRF-masked attention with a single head, projecting the queries and keys with a dense layer.
- 3.
To construct $\mathcal{G}$, we find $k=3$ neighbours for every node. This is done on the fly since the neighbours can change at every timestep.
- 4.
For Fig. 5, we train for 100k steps, using a learning rate of $3\times 10^{-4}$ until step 1000 then $1\times 10^{-4}$ until step 100k.

To enumerate some other important details: all models are trained with a batch size of 16; we use the AdamW optimiser [Loshchilov, ] with weight decay $10^{-3}$, clipping the gradient norm to 0.01; models are trained with 6 step rollouts, with losses computed on 128 sampled rays; $\Delta\bm{x}_{i}$ are constrained to $[-0.25,0.25]$ by a scaled $\tanh$ nonlinearity so particles do not move more than $25$cm per timestep.
For the rendering, like Whitney et al. [] we use four concentric annular kernels of radii $[0,0.01,0.02,0.05]$ with bandwidths $[0.01,0.01,0.01,0.05]$, approximated using $16$ nearest neighbours.
The near plane is 0.1m and the far plane is 2m, with the background set to solid white.

### C.4 ViViT: topological masking for video data

**Table 4: ViViT test accuracies on the Kinetics benchmark. Topological masking with GRFs boosts performance.**
| Variant | Test acc. |
| --- | --- |
| Softmax | $0.754$ |
| Softmax + GRFs | $0.758$ |

Here, we present preliminary results for incorporating topological masking with graph random features into video vision transformers [ViViT; Arnab et al., ].
Prompted by the original paper, we use a factorised encoder, where the spatial encoder follows the ViT architecture from earlier in the paper (see Sec. 4 and App. C.1) with patch size $16\times 16\times 2$.
The final index mixes pairs of successive frames.
The temporal encoder, which subsequently models interactions between the frame-level representations, also follows this architecture but with $4$ layers instead of $12$.
We train for $30$ epochs using the Adam optimiser, with base learning rate $10^{-1}$ and final learning rate $10^{-4}$. We take the same schedule as in Table 2.
We train and evaluate on the Kinetics 400 benchmark [Kay et al., ].
Table 4 shows the results for (i) unmasked softmax attention and (ii) softmax attention with GRFs.
Topological masking boosts predictive performance by $\mathbf{+0.4\%}$.
To our knowledge, this is the first application of topological masking to the video modality.

## Appendix D Asymmetric graph random features

In this Appendix, we present a variant of our core topological masking algorithm that relaxes the assumption that the function $f$ for the query and key ensembles of walkers are identical, instead using *asymmetric* graph random features.
These are believed (but not proved) to generally give higher variance mask estimates [Reid et al., ], but
making a judicious choice can also bring computational and implementation benefits.
Using asymmetric GRFs for topological masking is also a novel contribution of this paper.
In Sec. 4, it is found to be a convenient choice for the Interlacer [Whitney et al., ] because it can be straightforwardly swapped into existing code.

Asymmetric GRFs.
It is straightforward to see that Eq. 13 of App. A.1 generalises to

$$ $\mathbf{M}_{\alpha}(\mathcal{G})_{ij}=\sum_{v=1}^{N}\left(\sum_{k_{1}=0}^{ \infty}f_{k_{1}}^{(1)}\mathbf{W}^{k_{1}}_{iv}\right)\left(\sum_{k_{2}=0}^{ \infty}f_{k_{2}}^{(2)}\mathbf{W}^{k_{2}}_{jv}\right)$ (22) $$

if the discrete convolution

$$ $\sum_{p=0}^{k}f^{(1)}_{p}f^{(2)}_{k-p}=\alpha_{k}$ (23) $$

holds.
We do not necessarily require that $f^{(1)}=f^{(2)}$; this is a special case.
Another special case is $f^{(2)}_{p}=\{1\textrm{ if }p=0,\hskip 1.99997pt0\textrm{ otherwise}\}$, whereupon we need $f_{p}^{(1)}=\alpha_{p}\hskip 1.99997pt\forall\hskip 1.99997ptp$.
Sampling GRFs in this manner will still give an unbiased estimate of $\mathbf{M}_{\alpha}(\mathcal{G})_{ij}$, but with different concentration properties.

Computational benefits.
Supposing that $f^{(2)}_{p}=\{1\textrm{ if }p=0,\hskip 1.99997pt0\textrm{ otherwise}\}$, $\phi_{\mathcal{G}}^{(2)}(v_{i})_{n}=\mathbb{I}(i=n)$, a one-hot vector only nonzero at the coordinate corresponding to $v_{i}$.
In other words, we do not need to simulate any random walks to obtain GRFs for the keys; $\mathbf{\Phi}_{\mathcal{G}}^{(2)}=\mathbf{I}_{N}$ with $\mathbf{I}_{N}\in\mathbb{R}^{N\times N}$ the identity matrix.
It follows that

$$ $\widehat{\mathbf{M}}_{\alpha}(\mathcal{G})=\mathbf{\Phi}_{\mathcal{G}}^{(1)}{ \mathbf{\Phi}_{\mathcal{G}}^{(2)}}^{\top}=\mathbf{\Phi}_{\mathcal{G}}^{(1)} \eqqcolon[\phi^{(1)}_{\mathcal{G}}(v_{i})]_{i=1}^{N}.$ (24) $$

But $\phi^{(1)}(v_{i})$ is only nonzero at nodes visited by the ensemble of walkers beginning at $v_{i}$.
This makes masked attention straightforward to efficiently compute, obviating the outer products and vectorisation.
Simply,

$$ $(\mathbf{A}\odot\widehat{\mathbf{M}}_{\alpha}(\mathcal{G}))\mathbf{V}=(\mathbf {A}\odot\mathbf{\Phi}_{\mathcal{G}}^{(1)})\mathbf{V}=\left[\sum_{j=1}^{N} \mathbf{A}_{ij}\phi^{(1)}_{\mathcal{G}}(v_{i})_{j}\bm{v}_{j}\right]_{i=1}^{N}$ (25) $$

where $\mathbf{A}_{ij}=\exp(\bm{q}_{i}^{\top}\bm{k}_{j})$ (softmax) or $\mathbf{A}_{ij}=\bm{q}_{i}^{\top}\bm{k}_{j}$ (linear) and $\{\bm{v}_{j}\}_{j=1}^{N}$ are the value vectors.
Since $\phi_{\mathcal{G}}^{(1)}(v_{i})$ is sparse, the sum in Eq. 25 is computed in constant time so attention is computed in $\mathcal{O}(N)$.
Even more explicitly, for an ensemble of walks $\{\omega_{k}^{(i)}\}_{v_{i}\in\mathcal{N},\hskip 1.39998ptk\in[\![1,n]\!]}$, we can rewrite the above as

$$ $\left[\sum_{k=1}^{n}\sum_{\omega\textrm{ p.s. }\omega_{k}^{(i)}}\mathbf{A}_{i \omega[-1]}\frac{\widetilde{\omega}(\omega)f_{\textrm{len}(\omega)}}{p(\omega) }\bm{v}_{\omega[-1]}\right]_{i=1}^{N}$ (26) $$

where $\omega[-1]$ stands for the final node of the walk $\omega$ and p.s. means ‘prefix subwalk’ (see App. A.2).
The number of FLOPs to compute attention for node $v_{i}$ depends on the length of all the geometrically distributed walks beginning at node $v_{i}$, but we have seen that this is independent of graph size (Sec. 3.4).
Computing the entire matrix is only $\mathcal{O}(N)$.
We give pseudocode for this approach in Alg. 2.

Benefits and limitations of asymmetric GRFs.
The chief benefit of asymmetric GRFs is that Eq. 26 makes efficient masked attention very straightforward to compute, without relying on sparse matrix operations that may not be well-optimised in machine learning libraries.
We use it for the HD-VPD experiments because it can be easily integrated with the existing code: instead of finding each node’s (approximate) $k$-nearest neighbours and implementing the complicated message passing-type algorithm by Whitney et al. [], we find the $3$-hop neighbourhood and straightforwardly compute weighted attention with a learnable function of walk length.
The drawback is that the variance of the mask estimate with asymmetric GRFs tends to be greater, though proving this is an open problem.

Figure: Algorithm 2 $\mathcal{O}(N)$ topological masking with *asymmetric* GRFs

## Peer Reviews

confidence: 5
contribution: 3
initial_rating: 3
presentation: 4
questions: While the empirical results might be interesting, the fundamental theoretical claims that form the paper's main contribution are incorrect. A major revision would be needed to:
1. Correct the theoretical analysis
2. Properly characterize complexity and mention which category of graphs it can address
3. Either prove better bounds under specific assumptions or acknowledge limitations
4. Frame results in terms of parameterized complexity

The authors should consider addressing these fundamental issues during rebuttal, if possible.
soundness: 3
strengths: 1. The paper provides strong theoretical foundations with proven concentration bounds and complexity guarantees for GRFs.
2. The method shows concrete performance improvements on real-world tasks and scales to large problems (>30k nodes) that would be intractable with quadratic approaches.
3. The approach can be implemented with both symmetric and asymmetric GRFs, offering different trade-offs between computational efficiency and variance in mask estimation.
4. The experiments cover diverse applications (images, point clouds, videos) and include detailed ablation studies.
summary: This paper addresses the challenge of incorporating graph structural information into transformer attention mechanisms while maintaining computational efficiency. Their main focus is on topological masking, especially under the low-rank assumption of the attention matrix. The authors use graph random features (GRFs) to approximate topological masks for attention via importance sampling, which are parameterized as learnable functions of the weighted adjacency matrix. They propose a method to control transformer attention using graph node kernels based on random walks via power series of the adjacency matrix, with a random halting probability at each step. They provide concentration bounds in Theorem 3.1. Additionally, their empirical evaluation is carried out on diverse tasks like vision transformers on ImageNet, iNaturalist2021, Places365 and point cloud dynamics prediction for robotics applications.

While the experimental results show good promise, the paper’s theoretical complexity analysis and claims about O(1) GRF sparsity doesn’t hold true for all general graphs. Despite these theoretical issues, the paper introduces interesting ideas about using graph structure in attention mechanisms and provides novel empirical results, particularly in the robotics domain.
weaknesses: The paper's central claim of O(N) complexity relies critically on the assertion that Graph Random Features (GRFs) have O(1) sparsity. This claim is mathematically incorrect for several reasons:
In Lemma 3.2, while the result doesn’t show an N term, it is still implicitly dependent on the size of the graph. O(1) complexity implies that your non-zero entries per row vector $\hat{\phi}G(vi)$ are bounded by a constant independent of input size. The bound in Lemma 3.2 is still dependent on multiple parameters like $n$, $p_halt$ and $\delta$. So, one can say that your complexity is like the complexity of a “parameterized algorithm”, i.e., $O(f(n,p_halt,\delta))$, where f is some function of the parameters. 
Let's consider a family of complete graphs ${G_N}{N \geq 1}$ where $G_N$ has N vertices and each vertex has degree N-1. Then all edge weights are equal, i.e., $1/ sqrt( (N-1)(N-1) )$. 
At any step, the walk can move to another vertex with probability 1/(N-1). 
For a given walk starting at an arbitrary vertex, if you assign a r.v. to count the number of “distinct vertices” visited, even with the inclusion of geometric termination, you will find that this r.v grows with N because it has (i) more possible vertices available at each step to visit, (ii) the probability of visiting a new vertex at each step increases with N, and (iii) each successful step before halting can easily reach O(N) vertices. Therefore, making O(N) non-zero entries in the row vector and hence your attention matrix. With more independent walks starting from v, you can fill up even more non-zero entries. Hence, the O(1) bound doesn’t hold here.

As a demonstrative simple counter-example, consider the following two cases with fixed parameters. Let’s fix the parameters as n=10 (walks), p_halt = 0.5 and \delta=0.1 

Case 1: Small complete graph
- N = 10 nodes
- Each node has degree 9
- Even a 1-hop walk can reach 9 other nodes
- A 2-hop walk can reach all nodes

Case 2: Larger complete graph
- N = 1000 nodes
- Each node has degree 999
- A 1-hop walk can reach 999 other nodes
- A 2-hop walk can reach all nodes

While their bound might be the same in both cases, as it depends only on $n$, $p_halt$ and $\delta$, but the number of non-zero entries in both cases ends up being very different from one another. In the second case, you are much more likely to get O(N) non-zero entries due to the reasons I mentioned earlier about each hop having many more options, more distinct nodes, thus more reachability and coverage of the underlying graph (or at least exploring a large portion of the graph before terminating). 

This demonstrates that actual sparsity heavily depends on the structure of the graph. 

This rigorous analysis shows that the number of non-zero entries cannot be independent of graph size without additional constraints on the graph structure. The analysis shows that the results proposed by the author hold only with some assumptions on the graph structure, for example sparse graphs of bounded-degree graphs. 

The authors make a significant assumption in their discussion of Theorem 3.1. (Lines 278-280), where they state that "assuming that c remains constant (i.e. we fix a maximum edge weight and node degree as the graph grows)..." 
This reveals that their theoretical analysis works only for graphs with a bounded maximum degree and hence contradicts the claims about working for general graphs (line 82 in Introduction). This algorithm cannot handle dense graphs, complete graphs (or almost complete graphs), graphs where node degrees grow with N and many real world graphs where there can be very high-degree nodes and no bounds on degrees. 

The authors should make this bounded-degree assumption explicit upfront in the introduction and modify their claims about general graphs. 

A complete graph (or almost complete graph) isn’t completely unusual especially in the context of attention mechanisms where in the final layers pretty much end up having all tokens attending to each other in a pairwise manner. The experiments done in the paper are done on very low-degree graphs like grid graphs, which doesn’t demonstrate the applicability of their method to general graphs, especially large dense ones. 

Inconsistent and confusing notation:
- $(f_k)_{k=0}^\infty$ is sometimes treated as a sequence of reals, sometimes as a function
- Incorrect set notation: claiming (f_k)_{k=0}^∞ ⊂ R when sequences are functions from N to R
- Weighted adjacency matrix definition issues:
  * Claim W is weighted but then suggest $w_{ij} = 1/sqrt(d_i d_j)$
  * This normalization discards meaningful edge weights in attention context

ii) Missing crucial assumptions:
- No explicit assumptions about graph structure
- No discussion of how graph density affects complexity
- No proper analysis of how maximum degree impacts sparsity

The paper's main contribution is focused specifically on making topological masking efficient, rather than improving linear attention in general. The paper makes empirical comparisons to only basic linear attention models and focuses solely on topological masking efficiency. 

In the experiments section, it would be interesting to compare against other major linear attention variants like (i) Performers (Choromanski et al., 2020) Favor+ which uses kernel approximations and (iii) Nyströmformer which uses Nystrom’s method to approximate attention.

---

confidence: 2
contribution: 3
initial_rating: 6
presentation: 3
questions: I appreciate the authors’ valuable contributions in this area. As I am less familiar with applications in image, point cloud, or robotics contexts, I am particularly interested in understanding how Graph Random Features (GRFs) benefit graph neural networks on traditional graph datasets.

1. Could the authors provide examples or case studies that apply GRFs to commonly used graph datasets, such as Cora or Citeseer?
2. Could the authors also include a comparison of computational times between your methods and baseline approaches?

Thanks for this important work, and apologies for my gaps in my background knowledge.  I would also kindly request that the Area Chair reduce the weight of my review in the final evaluation.
soundness: 3
strengths: 1. The prposed method shares $O(n)$ time complexity and suitable for the relatively large scale input.
summary: This paper introduces a novel, efficient topological masking approach for transformers on graph-structured data, using learnable functions of the weighted adjacency matrix to adjust attention based on graph structure. By approximating with graph random features, this method supports linear attention, offering strong performance gains across diverse data types and large graphs.
weaknesses: see question.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 1
questions: 1. Eq. (4): The power series is generally not guaranteed to converge. It is better to specify clearly the underlying assumptions on W and alpha that guarantee convergence. Is $\alpha_0$ assumed here to equal 1, as in (Reid et al. 2024b)?
2. Remark 3.1: 
   - This remark is used as a lemma. Better state it as such.
   - It is in general not guaranteed that alpha has a deconvolution. Is it an assumption of Remark 3.1? Or is it guaranteed by some other assumption? Better clarify.
3. Ln. 115-116: "$\Phi_{Q,K} \in$" should probably be "$\Phi_Q, \Phi_K \in$"
4. Ln. 184-185: Statement is unclear. Why should it necessarily be faster?
soundness: 3
strengths: 1. Their method is the first to achieve $\mathcal{O}(N)$-time complexity for computing masked attention for general graphs, $N$ being the number of vertices.
2. The paper provides the first known concentration bounds for GRFs and rigorous sparsity guarantees. These theoretical insights are valuable, potentially extending beyond transformers to other domains that rely on scalable graph-based representations. 
3. Their method demonstrates improved predictive performance in various learning tasks.
summary: This paper presents a method for integrating topological graph information into graph transformers through a learnable topological-masking mechanism, using graph random features (GRFs). The authors propose to approximate topological masks via Monte Carlo estimation via GRFs to represent structural biases while ensuring linear-time computation.
weaknesses: 1. Dense and unclear presentation: 
   - While the method is theoretically sound, the presentation is mathematically dense and lacks clear explanations. This may pose a barrier to readers, particularly those less familiar with GRFs. In particular, the technical exposition in lines 184–254 is notation-heavy and unclear.
   - Algorithmic descriptions, such as those in Algorithm 1, are highly abstract and may be difficult to follow. 
   
   Without clearer explanations, the accessibility of the paper is reduced.
2. The paper lacks discussion of the method limitations. For example:
   - The practical applicability of this method depends heavily on the specifics of the graph structure and the task requirements, since it relies on approximations with random walks. In  graphs where relevant information is distributed over long distances or requires traversing multiple nodes, random walks may fail to capture the full structure efficiently.
   - For dynamic or evolving graphs, precomputing random walks is not feasible, and recomputing them on the fly could reduce efficiency. 
   - Since random walks introduce stochasticity, their effectiveness can vary based on the number of walks and the chosen halting probability. This means that the quality of topological masking may be sensitive to hyperparameters like the number of walks and the stopping probability, making it challenging to generalize the method across different graph structures.

   These limitations should be acknowledged and discussed for a more balanced perspective.

---

confidence: 2
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. While the author shows the test accuracies in Table 1, can the author also present other results, like the total training time or total flops, to validate the proposed method’s efficiency?

2. In Figure 5, it seems the accuracy improvement achieved by GRF Interlacer is on the starting timestep 0. After several timesteps, it’s becoming similar to MP Interlaced. Can the author explain the reason behind the accuracy improvement at the beginning and the drop? 

3. Besides GRF, are there other methods to do implicit graph masking in equation (8)? How's the performance?

If all my concerns are resolved properly, I will be happy to increase my score.
soundness: 3
strengths: 1. The paper has a good motivation for introducing linear topological masking of low-rank attention. 

2. The author explains well from introducing the topological mask, using the graph feature to achieve low-rank attention, and leveraging GRF to approximate the graph feature.

3. The explanation is clear, the figures are illustrative, and the writing is well-structured.
summary: This paper proposes a topological masking method when training transformers on graph-structured data. By decomposing and approximating the graph mask with graph random features, the proposed method achieves linear time and space complexity w.r.t input size. The author shows that their masking algorithm is efficient and high-performance using experiment results.
weaknesses: 1. While the author emphasizes a lot about the efficiency of the proposed method, the evaluation and experimental parts mainly show the accuracy achieved and lack the corresponding efficiency results like time and memory.

---

confidence: 3
contribution: 3
initial_rating: 8
presentation: 4
questions: * While the true time complexity of the proposed algo is $\mathcal{O}(Nmd)$, most part of the paper omit $m$ and $d$. I think the authors should make this point clear since in some cases where $m$ and $d$ can be large enough.
* The citation in line 161 (Borgwardt et al., 2005) is about graph kernels, not graph node kernels if I understand it well.
* Though the paper shows experimentally the impact of the number of random walks $n$ on the performance, I would also like to see its impact on the computation time.
* In section 4, maybe it's better to use \textit{subsection} for each experiment instead of \textit{paragraph}.
* This question is not about the contribution of the paper, but the general idea of using structural graph masking for transformer. Isn't it a reinvention of the graph neural network by integrating a matrix representing the structure information into Transformer? Can author make a (possible) link between the two?
soundness: 3
strengths: * The paper is well-organized and well-written, it's a pleasure to read.
* The reasoning of the idea is clear. We can understand very well why we need to use kernel gram matrix as the masking matrix and why we need the graph random features.
* The experiments are well-presented, which show the relevance of the proposition for the situation where $N$ is large.
summary: The paper proposes to use the gram matrix of the graph node kernels as the masking matrix for the attention mechanism of Transformer. The fact that each element of the gram matrix can be written as the inner product of two feature vectors (kernel def), we can write the **attention mechanism with masking**  in the form of Equation 3 of the low-ranking setting by redefining $\Phi_Q$ and $\Phi_K$. graph random features are also further applied to reduce the complexity. The experiments are conducted on ViTs and the prediction of the particle dynamics for robotics.
weaknesses: Please see the questions raised.

---

confidence: 2
contribution: 2
initial_rating: 6
presentation: 3
questions: See the weakness.
soundness: 3
strengths: 1. The research topic is highly important and intriguing and the authors provide experimental results on some larger scale image and point-cloud datasets.
2. Generally, the notation and theorem in the paper is well defined and illustrated. 
3. The authors provide the theoretical analysis and ablation study to show the effectiveness of the proposed method.
summary: This paper presents a novel approach by parameterizing topological masks as a learnable function of a weighted adjacency matrix. This method incorporates a strong structural inductive bias with rigorous concentration bounds, improving both time and space complexity.
weaknesses: #### Major problems
1. The background and related work section is limited. To my knowledge, this is not the first work to incorporate the inductive bias inherent in graph nodes into GNN and transformer-like models.
2. Additionally, improved time and space complexity have been discussed in existing papers such as [A, B]. Therefore, the benefits and differences between linear attention and the proposed topological masking should be thoroughly explained.
3. It's better to include more experiments and evaluations on traditional graph datasets.

[A] Wu, Qitian, et al. "Nodeformer: A scalable graph structure learning transformer for node classification." NIPS2022.  
[B] Wu, Yi, et al. "KDLGT: A Linear Graph Transformer Framework via Kernel Decomposition Approach." IJCAI 2023.
#### Minor issues
1. Confusing notations in line 148, the size of the graph N ($\mathcal{N}$)?

Meta-review: The authors consider the transformer attention with graph structured data. In particular, the authors focus on the topological masking of low-rank attention. The authors propose to leverage the graph random features to approximate topological masks and parameterize it as a learnable function of a weighted adjacency matrix. The authors also derive the concentration bounds and show the linear time and space complexity for the proposed approach. The authors illustrate its empirical advantages on images and point cloud data. All Reviewers agree that it is a good submission for ICLR'2025. We urge the Authors to incorporate the Reviewers' comments and discussion in the rebuttal to the updated version, especially regarding time consumption.

Decision: Accept (Poster)