Linear Transformer Topological Masking with Graph Random Features
Authors: Isaac Reid
Kumar Avinava Dubey
Deepali Jain
William F Whitney
Amr Ahmed
Joshua Ainslie
Alex Bewley
Mithun George Jacob
Aranyak Mehta
David Rendleman
Connor Schenck
Richard E. Turner
René Wagner
Adrian Weller
Krzysztof Marcin Choromanski
Venue: ICLR 2025 Poster (2025)
Keywords: transformer
linear
attention
graph
random walk
Monte Carlo
encoding
topological masking
point cloud
Performer
Primary area: learning on graphs and other geometries & topologies
TL;DR: An efficient algorithm to incorporate structural information about the underlying graph topology into linear attention transformers

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 the 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 $\mathcal{O}(N)$ time and space complexity with respect to the number of input tokens. The fastest previous alternative was $\mathcal{O}(N \log N)$ and only suitable for specific graphs. Our efficient masking algorithms provide strong performance gains for image and point cloud data, including with $>30$k nodes.

Published as a conference paper at ICLR 2025
LINEAR TRANSFORMER TOPOLOGICAL MASKING WITH GRAPH RANDOM FEATURES
Isaac Reid Avinava Dubey?, Deepali Jain', Will Whitney", Amr Ahmed?, Joshua Ainslie,, Alex Bewley', Mithun George Jacob', Aranyak Mehta', David Rendleman", Connor Schenck3, Richard E. Turner', René Wagner", Adrian Weller Krzysztof Choromanski
'University of Cambridge, GGoole Research, GGoogle DeepMind, *Alan Turing Institute, CColumbia University
ABSTRACT
When training transformers on graph-structured data, incorporating information about the underlying topology is crucial for good performance. Topological mask- ing, a type of relative position encoding, achieves this by upweighting or down- weighting attention depending on the relationship between the query and keys in a graph. In this paper, we propose to parameterise topological masks as a learn- able function of weighted adjacency matrix novel, flexible approach which incorporates a strong structural inductive bias. By approximating this mask with graph random features (for which prove the first known concentration bounds), we show how this can be made fully compatible with linear attention, preserving (N) time and space complexity with respect the number of input tokens. The fastest previous alternative was ONN log 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 > 30k nodes.
INTRODUCTION
Across data modalities, transformers have emerged as a leading machine learning architecture (Vaswani, 2017: Dosovitskiy et al., 2020; Arna et al., 2021). They derive their success from mod- elling complex dependencies between the tokens using 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; c.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 el- ements of the sequence (Yang al., 2019). Meanwhile, relative position encoding (RPE) captures the spatial relationship between the tokens (Shaw et al., 2018). In some cases, the queries and keys can be arranged in a graph G and function M(G) can be used to modulate the transformer atten- tion. This is referred to as topological masking (Choromanski et al., 2022). It provides a powerful avenue for incorporating structural inductive bias from G into transformers, in some cases closing the performance gap with the best-customised graph neural networks (Ying et al., 2021).
Topological masking of low-rank attention. Topological masking is simple when the attention is full-rank. In this case, the entire attention matrix A E N (with N the number of input tokens) materialised memory, so its individual entries Aij can pointwise multiplied by (fixed or learnalle) mask M(G)s)- But full-rank attention famously incurs quadratic O(N2) space and time complexity cost, which may become prohibitive on large inputs. This has motivated a number of low-rank decompositions of A. These improve scalability by avoiding computing 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 M and restricted classes of graphs G (e.g. trees and grids), and are often O(N log N) time complexity rather than (N) (App. B; Choromanski et al., 2022; Liutkus et al.,
Work completed student researcher tSenior lead.

---

Published as a conference paper at ICLR 2025
Unmasked
Masked
g
Regular
M
A
V
N
A
O
M(G)
V
O(N3)
m
N
Niw (sp.)
N
Linear ((N)
k
V
Q.9
V
Figure 1: Schematic overview. Regular attention is ( with N the number of input tokens. Topological masking modulates A by graph function M(G), improving predictive performance. Linear attention reduces the time complexity ((N) by leveraging a low-rank decomposition. Our contribution (blue) is the first algorithm to achieve both 1oth-0(1) O(N) topological masking of low-rank at- tention - by approximating M(G) with graph random features (GRFs). GRFs are sparse vectors (de- noted 'sp.") computed by sampling random walks, constructed so that IE P with strong concentration properties (Thm. 3.2).
(K.G) A C M(g)
2021). Addressing these shortcomings is the chief focus of this manuscript. Our central question is: how can we implement efficient yet flexible topological igraph-based masking of transformer attention, in the low-rank setting where A never explicitly instantiated?
Graph random features. To answer the question posed above, we propose to use graph random features (GRFs) (Choromanski, 2023; Reid et al., 2024b). GRFs sparse, randomised vectors whose dot product gives an unbiased estimate of arbitrary functions of a weighted adjacency matrix, class that captures the structure of G and is popular in classical graph-based machine learning (Smola and Kondor, 2003; Kondor and Lafferty, 2002). GRFs have been used for approximate ker- nelised node clustering and in scalable graph-based Gaussian processes (Reid al., 2023a; 2024c), but never before in transformers. Advancing understanding of the convergence properties of GRFs another key goal of this work.
Our contributions. We present the first known O(N) algorithm for topological masking of linear attention transformers with general 9, using graph random features (Fig. 1).
1. We propose to modulate the transformer attention mechanism with learnable function of weighted udjacency matrix a simple, general topological masking trick which brings a strong inductive bias and substantial improvements to predictive performance.
To scale, we replace the exact, deterministic mask with a sparse, stochastic estimate constructed using graph random features (GRFs). We use novel exponential concentration bound, requir- ing only a mild technical assumption on G as N grows, to prove that our method linear time and space complexity with respect input size: the first known algorithm with this property.
3. We demonstrate strong accuracy gains on image data, well as for modelling the dynamics of massive point clouds (>> 30k particles) in robotics applications where efficiency essential.
PRELIMINARIES
Masked attention. Denote by N the number of input tokens and d the dimensionality of their latent representations. Consider matrices Q, K, V E TR xd whose rows are given by the query (fq: |. ) key ((k; } and value vectors (f ((v,) ) respectively. The atlention mechanism, a basic processing
unit of the transformer, can be written:
(1) Here, the function K. assigns a similarity score to each query-key pair and 1N = [1 E RRV For regular softmax attention, one simply takes KC(Q, K);; = exp(g, k,), optionally also normalising by d. To incorporate masking, we make the modification
Att K. V) = D- A K(Q,KK D A1x.
Am M @ K)Q,K),
2

---

Published as a conference paper at ICLR 2025
where M E R NNN wise) product. For softmax, this is equivalent to adding a bias term by == log Myj to q. k, before exponentiating. As remarked in Sec. 1, this can be used to enforce causality or relative position encoding (RPE) (Shaw et al., 2018). Taking M(9) 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.
is a (hard soft) masking matrix and @ denotes the Hadamard (element-
Improving efficiency with low-rank attention. The time-complexity of Eq. I is O(N2), which famously limits the ability of regular transformers scale long inputs. A leading approach to mitigate this and recover ((N) complexity is linear or low-rank attention (Katharopoulos et al., 2020; Choromanski et al., 2020). This takes a feature mapping & TRd -> Ihe with m < N and constructs the matrices ΦQ= == (((g:)] and &x = [o(k))]== where QQ.K E CRNxm. defines LR Q, K) = TQ Ø and computes
xm, One then
ALLLR K, =) ( (DKV)), D QQ (*KIN),
where the parentheses show the order of computation. Exploiting associativity, the time complexity is reduced to (NNdd), albeit usually with some sacrifice in performance compared to full-rank softmax attention. Since we care about scaling with graph size, we will henceforth write this as O(N). Different choices of feature maps é exist; for example, one take random feature map that provides a Monte Carlo estimate of the softmax kernel (Performer; Choromanski et al., 2020), or deterministic nonlinearity like elu(x) + x (linear transformer; Katharopoulos al., 2020).
Efficient masking of low-rank attention. As Eq. 3 suggests, low-rank attention derives its speed from avoiding computing A K.. Q, K) explicitly, preferring to evaluate faster products like V and ØKlN first and use them to weight the query features ΦQ. This makes the direct Hadamard product implementation of masking in Eq. 2 impossible: if the matrix is never explic- itly materialised in memory, one cannot pointwise multiply each of its elements. Algorithms have been proposed mask implicitly (without instantiating A or M) in very simple cases, e.g. causal masking (Choromanski et al., 2020) one-dimensional RPE (Liutkus et al., 2021; Luo al., 2021). Subquadratic implicit masking is also possible when M is very structured: e.g. functions of shortest path distance on grid graphs and trees (Choromanski et al., 2022), of the degrees of the pair of input nodes (Chen et al., 2024). See App. B for a detailed review. However, an (NV algorithm capable of masking with flexible functions general graphs has, until now, remained out of reach.
Remainder of manuseript. In Sec. 3.1 we parameterise topological masks M(G) as Icarnable 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 A M in memory. Sec. 3.3 offers faster stochastic alternative using GRFs, which we prove is 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 G(N where N == ( UN) } is the set of nodes and & is the set of edges, with (v), E & if and only if there exists an edge between and uy in G. The size of the graph N corresponds to the number of tokens. Denote by M(G) E IR N the topological masking matrix.
3.1 GRAPH NODE KERNELS AS TOPOLOGICAL MASKS
As remarked above, common choice for M(9), function of the shortest path distance be- tween nodes and Uy (Ying et al., 2021). Though effective at improving performance, this requires computation for every pair of nodes so cannot straightforwardly be applied to low-rank attention in scalable way pay-eexce except for special structured & where the mask becomes simple (Choromanski et al., 2022). Shortest path distances are also very sensitive to changes like the addition or removal of edges, and need access the whole graph upfront.
new topological mask: graph node kernels. Given that we intend M(9);; to quantify some sense of topoloiiaa similarity between v and ny. a more principled choice may be graph node
3

---

Published as a conference paper at ICLR 2025
kernels: positive definite, symmetric functions k : N + mapping from pairs of graph nodes to real numbers. These are widely used in bioinformatics (Borgwardt et al., 2005), community detection (Kloster and Gleich, 2014) and recommender systems (Yajima, 2006), with more recent applications in manifold learning for deep generative modelling (Zhou et al., 2020) and solving single- and multiple-source shortest path problems (Crane al., 2017). They are also used model the covariance function in geometric Gaussian processes (Zhi et al., 2023; Borovitskiy et al., 2021).
Graph node kernels are often parameterised follows. Let W RN be weighted adjacency matrix, so that W := [wiykkjew with Wij nonzero if and only if Up) E &. A common choice is Waj = 1/ dddj, with d, the degree of U;. Then consider the power series
20 W
M. (9)
Σ Og
W
(4)
A=0
where (au) R is a set of real-valued Taylor coefficients. This is positive definite if and only if Ex=g AR A4 A E A(W), where A A(W) denotes the set of eigenvalues of W. With different choices for (ox)) 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 Or (Reid et al., 2024b; Zhi et al., 2023).
Eq. provides 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 unmasked attention. However, explicitly computing (9) e R N N and storing it in memory is incompatible with (NN) 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(q, k,) with 'featurised' alterna- tive (qq:) (k,). Can we do the same for Mg(g)? The following is true.
Lemma 3.1 (Explicit N-dimensional features for graph node kernels (Reid et al., 2024b)). Let
(fx) C R denote the deconvolution of (aw)) i.e. ΣØ=o fpfk-p = ax V k. There exists feature matrix &g E -NN satisfying (G) = &g LJ given by &g = - fa Wk
This is can be shown in a few lines of algebra; see App. A.I. In general, one requires some extra technical assumptions to ensure that Ma(9) and &g converge, which we also describe in App A.1. We refer to the rows of the feature matrix Dg as graph features, letting Dy =: [6g (v:),;,;;- with (g(()) E R-N the graph feature of node v.- For every kernel there trivially exists a feature map (Gretton, 2013). but unusually in this case is fiiie-dimensional and admits a simple closed form.
Let 69 R x N - RMxN denote the outer product which maps pair of vectors to a matrix, and vec RmxN + RRMN denote the vectorising operation that flattens a ma- trix into a vector. Then o(q.) @ @c(e;) is an m x N matrix whose jktth element is ((g))) and vec (ó((.) @ (e)) an mN din--iinnsiioal vector whose (N(j 1) + k)th
entry is ((q:)oo (vi):. Observe the following:!
mN
vec ( @ (o(())) T veo (Ø(ks) @ oo((;)) = Σ vec (ó(q:) @ do(v:)), (olks) @ da(1))).
N
= (o(q:) To(k)) (do(u:) oa(1;)) = Ma
(5)
Hence, taking (vec (o(fq, k;)) & BREN C RmN as features, we can compute masked attention implicitly, without materialising A and M.. Concretely, we have that:
$ T K.g KLR o M. if [QQK) [vec (Ø((q; k;)) & (o(r:)) NxNm.
(6)
Masked low-rank attention can then be computed using 4 (Q.Kk. as the feature matrices:
AttER,M (Q, K, V,9) D
= D-1 ( ( V)). D = &Q. (0K,oln)
(7)
Recall the old adage: the dot product of outer products is the product of dot products.
4

---

Published as a conference paper at ICLR 2025
Time complexity. Eqs 6 and give recipe for computing masked attention implicitly in feature space, without needing to materialise A or M in memory. However, this approach relies on com- puting the feature matrix $ (Q E IN Nm and evaluating products like & k glv, which still incurs O(N²) space and time complexity. For a linear masking algorithm, we must reduce this to O(N). This can be achieved by approximating the features.
3.3 FASTER IMPLICIT MASKING WITH GRAPH RANDOM FEATURES
Our goal is now to replace $ (Q.K), by Monte Carlo estimates T (Q.K). that satisfy
E (Qockks) P Ba. KLR O Ma.
2k ((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, 2023; Reid et al., 2024b).
At a high-level, GRFs can be understood as follows. From Lemma 3.1, the graph feature for node Us is given by (v,) |E f& W 1 Since W is adjacency matrix, W counts the number of graph walks of length & between nodes and j. weighted by the product of traversed edge weights. weighted sum over all walks of any length between nodes u; and ey. with longer contributions suppressed since edge weights are at most 1 (by normalisation of W)- GRFs apply importance sampling, approximating this infinite sum by a Monte Carlo estimate ōg using random walks.
So (v:); can be interpreted
Using GRFs. Concretely, to build the GRF Tg (1z) one samples n random walks {4/2) [k=1 C Ω
out of node v,, where Ω = | n2 E (M,, v(+1) E E, E N} is the set of all walks. Denote by p(w) the probability of walk which is known from the sampling mechanism. p(w) is typically chosen so that walks randomly halt, prioritising sampling lower powers of W that contribute more to the kernel. Let w 12 > IR be function computing the product of traversed edge weights, len(w)
(ww)= = II
W [[( if len(w) 1, otherwise
(9)
where len(w) is the number of hops and wli] is the node visited at the ith timestep. Then:
((wir) ((ellaa) M[win prefix subwalk of ("),,, 8= I, N, (10)
(o( : 71
p(wig)
Way EDig
where Qig denotes the set of all walks between nodes v; and ug. of which wig is member. 'Prefix subwalk' means that the walk w beginning at node i; initially follows the sequence of nodes wig then optionally continues. This complicated notation belies a very simple algorithm: one simulates random walks out of node u; 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
GRFs give unbiased approximation of graph node kernels: E((g(())) T TTo(1;)) = Maij- This fol- lows because [[[ prefix subwalk of w = p(w:,)) Unbiasedness is one of our required prop- erties; we also need & (Q.K).c to support ((N) matrix-vecto multiplication. We will show this by providing the first known proof that GRFs are sparse: ((v,) has (l) nonzero entries && E W.
3.4 NOVEL THRORETICAL RESULTS FOR GREs: SHARP ESTIMATORS WITH SPARSE FEATURES
It is intuitive that GRFs are sparse. The walk sampler p(w) is chosen so that walks are short, e.g. taking simple random walks with geometrically distributed lengths. From Eq. 10, this means that (g(4)) is only nonzero at nodes 'close' to Most of its entries are O. However, to prove that (Q,Kj supports O(N) matrix-vector multiplication we need to rigorously relate this to the quality of approximation of (G). This requires concentration inequalities for GRFs, but no such bounds were previously known. The following result novel.
Theorem 3.2 (GRF exponential concentration bounds), Consider graph G with adjacency matrix W and node degrees {ds oo Suppose we construct GRFs (pg(v:))ree by sampling random
5

---

Published as a conference paper at ICLR 2025
walks (wk ( JU,EN, kE[1.n] with geometrically distributed lengths, terminating with probability Phote
at each timestep. Let = - [aal, " Pass yb. Suppose is finite, which is guaran- teed e.g. for bounded f and suitably normalised W, or if there exists some imax E N such that fa vanishes for all i > Amua (see App. A.2). Then:
[Wsylds
Pr ( (o((y) - -Massl >1) S 2exp
(11)
2(2n 1)2c4
Proof sketch. Full details are in App. A.2; here, we provide an overview. Treat each pair of ran- dom walks Xg == (wk) _ ())
E f x Ω simulated out of nodes and "y respectively, as a random variable. Let E [1, n] so we sample n such pairs of walks in total. Consider modifying one of the pairs of walks. If is finite, the maximum load it is capable of depositing is finite we can bound the -nom of the contribution made to (c(U;) and og(U;) (vz). This allows us to bound the change in the estimator (v:) (1y) if Xg is modified. Applying McDiarmid's inequality, the result follows.
Thm 3.2 provides exponential bound on the probability of topological mask estimate deviating from its true value. Moreover, assuming that e remains constant (i.e. we fix a maximum edge weight and node degree 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 good mask estimate with high probability, independent of graph size. Now observe the following.
Lemma 3.3 (GRF sparsity). Consider a graph feature (g(()) constructed using m random walkers each terminating with probability Phalt at every timestep. With probability at least 1 - 6, @g (1:) will have log(1 (1 - 8/) log(1 log
Phoit) or fewer nonzero entries.
Proof. Since the walk lengths are geometrically distributed, the probability that single walk is of (1 )"att) n independent walkers are all b or shorter with probability (1 (1 - "h Let this be equal to 1 - ō and solve for b. At most bn entries of (1:))) can then be nonzero, so the sparsity guarantee follows.
length or shorter
Lemma 3.3 means that, if fix the number of walkers (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.4 (GRFs implement O(N) topological masking). Suppose we construct set of graph random features (Øg(v;) (og())))e,e)) Je,e, choosing the number of walkers n as described above with c, Photo 8, t and the deviation probability Pr (IÓGRF(U:) (GRF(U,) MMM]l t) fixed hyperparameters. Then the number of nonzero entries in & (Q,K), N is linear in N, which means that topo- logical masking is implemented with O(N) time complexity.
Proof. Lemma 3.3 shows that the number of nonzero entries in (c is independent of the number of nodes N. This must also be true of vee(ő(k,) @ (g(())) so the number of nonzero entries in KK.G == (vec(((k))) @ dg (ox))))) is proportional to N. The same holds for the equivalent matrix for the queries, @Q.g. The time complexity of the sparse matrix-matrix matrix multiplication & KK.g V is ((Ndd) and Φ IN is O(N)). Both these time complexities are linear in N. This is also the k
case for the subsequent computations involving &
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., 2024c; Borovitskiy et al., 2021). Thm. 3.2 is remarkably general, requiring only the modest assumption about G and f that c is finite. To guarantee (N) scaling assume that is fixed as G changes, which is guaranteed by e.g. fixed maximum node degree and edge weight.
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 overview). Graph node kernels, which compute infinite weighted sum of walks

---

Published as a conference paper at ICLR 2025
2o(1:) @o(())
Ag R (o(a)))((k))))
Mas = (ōotv) (o((y)))
A. May veo (o(a) & 6o(vs)]
vec ((K)) @ So(ez))
Ma E 9 2
Figure 2: Visual overview. A graph G (left) has weighted adjacency matrix W (centre left). A learnable power series (G) Ei ag W is an effective topological mask or graph RPE (centre right). (G) be efficiently approximated using graph random features (centre top), which perform importance sampling of halting random walks. The feature øg (v:) is only nonzero at entries visited by the ensemble of walks beginning at Us In Thm. 3.2, we prove that the number of such entries is ((1) whilst still accurately estimating M, (W)ss with high probability, so GRFs are sparse. This unlocks (N) topological masking. Note that @g (v:) (g(1)) only nonzero 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.
between pairs of imput 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 O(N) with low performance loss. Eq. 11 gives remarkably tight bound on the quality of mask estimation. Physically, masking with GRFs means that nodes v; and vy 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 G. It upweights crucial interactions between nearby nodes, whilst also sampling longer-range range attention with lower probability.
Algorithm ((N) topologically-mann masked attention for general graphs
query matrix Q E IR xd key matrix K E Nxd value matrix V R graph G with weighted adjacency matrix W E NxN, learnable mask parameters (fi).me, hm number of random walks to sample n E N, query/key feature map Φ(-) Rd - R Output: masked attention Att LR.Mi (Q, K, V,G) (. in O(N) time.
Input:
1: Simulate n terminating random walks {w/ (2) In Jkm1 out of every node is E
2: Compute sparse graph random features (oG(U:) learnable mask parameters (fi)
R. using the walks (w()0
(0) yizz and
3: Compute the query and key features ) i-1 nonlinearity (e.g. elu(x:) + (Katharopoulos al., 2020) Monte Carlo estimate of softmax (Choromanski et al., 2020))
N with Φ the chosen linear-attentin attention
4: Combine the query/key features and graph features into sparse topology-enhanced feature matrix, $ qq.K).g = [vec(o({q,, k;)) (c(e)))) E IR
5: Compute low-rank attention with topological masking in O(N) time by
AU QQ = D (ΦQ9 (okoV)). D = TQ. (KK,OIN).
(12)
Distributed computation, expressivity and GATs. Alg. 1 is simple to distribute for massive graphs that cannot be stored on single machine; to simulate random walks one only needs a node's im- mediate neighbours rather than all of G upfront. This desirable property is not shared by masking algorithms that rely on the fast Fourier transform (Luo et al., 2021: Choromanski et al., 2022). More- over, 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
expressive than standard graph

---

Published as a conference paper at ICLR 2025
neural networks and unmasked transformers, which notoriously fail this test (Morris et al., 2019; Xu et al., 2019). Finally, we remark that our algorithm can be interpreted as a stochastic extension of graph attention networks (GATs; Veličković et al., 2018), with greater focus linear attention and scalability. Tokens still attend to a subset of other tokens informed by the topology of g, but importance sampling random walks permits interactions beyond nearest-neighhour. Incorporating learnable parameters (fi) lets us estimate principled function of W.
4 EXPERIMENTS
In this section, we test our algorithms for topological masking with GRFs, We consider data modal- ities with different graph topologies: images and point clouds.
4.1 TIME COMPLEXITY
For hardware- aggnosiic comparison, first com- pute the total number FLOPs for evaluating (i) un- masked 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 theoreti- cal results in Sec. 3.4 are independent of (provided is upper bounded). We take per node with termination probability Phali 0.5. The latent dimension is d 8 and the feature size m 8. In the masked case we show the mean number of FLOPs over 10 seeds since GRF sparsity itself a random variable. As anticipated. the time complexity of our mechanism is (NN); there is constant multiplicative cost for topological masking. We clearly improve upon the time complexity of ((N2) softmax attention.
Time complexity scaling
10
10'
Lineur
FLOP% 6700%)
GRF
10"
10"
random walkers
10"
Number of graph modos, N0
Figure 3: Number of FLOPs number of graph nodes for softmax attention, linear at- tention, and linear attention with GRF topo- logical masking (ours).
4.2 VrTs: GRID GRAPHS AND IMAGE DATA
Next, we use our topological masking algorithm to improve the vision transformer (ViT; Doso- vitskiy et al., 2020). The underlying graph G is a square grid of size num x x_patchhs x num num_y_ p_ppythhpe,s,. with nodes connected by edge if they neighbours in y. Since the graph is fixed, random walks can pre-computed and frozen for each attention head and layer be- fore training and inference. Full architectural and training details in App. C. All our ViT models also endow tokens with learned additive absolute position embeddings.
Table shows the final test accuracies for ImageNet (Deng al., 2009), iNaturalist2021 (Horn et al., 2018) and Places365 (Zhou et al., 2018). For the slow (superlinear) variants, we include (i) unmasked softmax and (ii) Toeplitz-masked linear (Choromanski et al., 2022) attention. Toeplitz can only be used for this specific, structured graph topology. We also include (iii) M. (G) mm)-keeed linear, which takes an explicit power series in W (Eq. to directly modulate the N N attention matrix. This is equal to our Monte Carlo GRFs algorithm in the xi limit, so it indicates the asymptotic performance GRFs with very many random walkers. For the fast (linear) variants,
Variant
Time comp.
ImageNet IM
NNauuralis20211
Places365 (Small)
2.7M
1.8M
Unmasked softmax
O(N2)
0.741
0.699
0.567
Toeplitz-masked linear
O(N log N) O(V2)
0.733
0.674
0,550
(G -maakedd linear
0.741
0.692
0,551
BAST
Unmasked linear
O(N) ((N)
0.693
0.667
0,543
GRF-masked linear
0.730
0.689
0.548
Table 1: Final transformer test accuracies across attention masking variants and datasets. Our algo- rithms are shown in bold. Topological masking with GRFs is the best O(N) mechanism, matching or beating more expensive alternatives.
8

---

Published as a conference paper at ICLR 2025
include (iv) unmasked linear and (v) GRF-masked linear attention. Our algorithm gives strong improvements above the unmasked baseline, by +3.7%, +2.2% and +0.5% respectively. Remark- ably, it is competitive with much more expensive variants: GRFs often match or beat Toeplitz. For ImageNet and 220, it even substantially narrows the gap with ((22) softmax.
4.3 ABLATION STUDIES
In App. C.2, we perform ablation studies for the ViT experiments. We find that predictive perfor- mance improves with the number of walkers n because the accuracy of the Monte Carlo estimate (G) increases. Moreover, importance sampling walks is crucial; if we fail to upweight improb- able walks, the gains vanish. Lastly, the inductive bias from parameterising (G) as a function of W key. Fully learnable graph features of the same dimensionality N, an impractical and expensive alternative to our method, give much smaller gains.
4.4 PCTs: PREDICTING PARTICLE DYNAMICS FOR ROBOTICS
As a final task, we use our algorithm model high-density visual particle dynamics (HD-VPD; Whitney et al., 2023; 2024), predicting the physical evolution of scenes involving bi-manual robot using a point cloud transformer (PCT). Images from a pair of cameras are unprojected 3D particles. Their interactions, conditioned on specific robot actions, are modelled by dynamics network. We predict updates the latent point cloud, render the result to an image with condi- tional NeRF (Mildenhall et al., 2021), 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.
set of
For the dynamics network, we which alternates unmasked linear attention and message passing-style updates between neighbouring particles (Whit- ney et al., 2024). This architecture is twice as fast as GNNs with the same prediction quality, and can handle 4x bigger point clouds. To test the ability of GRFs to efficiently cap- ture local topological structure, we replace message passing by GRF-masked linear attention, taking query/key features <(+) ReLU(.). Implementation details for HD-VPD, in- cluding the robotic arm and camera configurations, are the same as reported by Whitney et al. (2024); see App. C.3 for discussion. To disambiguate,
Interlacer: PCT
Accuracy vs. rollout timestep
Image SSIM
MP Interlace
0.78
0.76
0.74
Rollout timestep
Figure 4: Accuracy comparison. Structural similarity index measure (SSIM) between ground truth camera frames and predictive NeRF render- ings after 100k training steps, plot- ted against rollout timestep. Higher is better. GRFs improve the accuracy of dynamics predictions.
will refer to our model the GRF Interlacer and Whitney et al.'s as the message pass- ing (MP) Interlacer.
In contrast to our previous experiments, for HD-VPD the underlying graph G is no longer fixed walks cannot be pre-computed. Instead, we construct them on the fly, using an approximate k-nearest neighbours solver get adjacency lists of length k = 3 for each node. We use repelling random walks for variance reduction (Reid al., 2024a), sampling hops. For easy integration with existing Interlacer code, we use asymmetric GRFs (App. D; Reid et al., 2024b). We consider point clouds with N 32768 particles, which is too many explicitly instantiate A or M in memory.
Fig. 5 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., 2020) is left to future work. Fig. 4 shows predicted video frame accuracy (image SSIM) vs. rollout timestep for each model. GRFs outperform the vanilla transformer and MP Inter- lacer baselines, modelling the particle dynamics more faithfully. We include examples of predictive videos in the supplementary material.
In addition to modelling cloud dynamics more accurately, our method is faster. While the MP Interlacer trains at 0.925 steps/second, the GRFs Interlacer trains at 0.982 steps/second. This makes it faster by 6% a gain our method enjoys because it avoids expensive GNN-style message passing. The unmasked baseline, which struggles to capture local structure so gives poor accuracy, trains at 1.49 steps/second.

---

Published as a conference paper at ICLR 2025
Groundtruh video frame
GRF Interlacer prediction
Groundruth video frame
GRF Interlacer prediction
Timestep
2 4
1+ 2
1+ AL
44 24
Figure 5: 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 can, moving 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.
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 perfor- mance gains across data modalities, including on massive point clouds with > 30k particles for which computational efficiency is essential. Future work might include improving hardware and libraries for high- nigh-dimensional sparse linear algebra to fully benefit from the guaranteed complex- ity improvements, and investigating whether dimensionality reduction techniques like the Johnson- Lindenstrauss transform (Dasgupta al., 2010) can be applied to GRFs for further speedups.
10

---

Published as a conference paper at ICLR 2025
ETHICS AND REPRODUCIBILITY
Ethics statement: Our work is foundational with no immediate ethical concerns apparent to us. However, increases scalability provided by improvements efficient transformers could exacer- bate 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. Apart from the robotics dataset, all datasets are standard and freely available online. Exhaustive experimental details about the training and architectures are reported App. C; see especially Table 2.
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 Stu- dentship and Google PhD Fellowship. This work was completed whilst employed as student researcher at Google. IR sincerely thanks Silvio Lattanzi for making this possible. RET is sup- ported by the EPSRC Probabilistic AI Hub (EP/Y028783/1). AW acknowledges support from Turing Al fellowship under grant EP/V025279/1 and the Leverhulme Trust via CFI.
11

---

Published as a conference paper at ICLR 2025
REFERENCES
Anurag Arna, Mostafa Dehghani, Georg Heigold, Chen Sun, Mario Lučić, and Cordelia Schmid. Vivit: A video vision transformer. Proceedings of the IEEE/CVF international conference on computer vision, pages 6836-6846, 2021. URL 0ttp://///.l...110..8885/00//000
Zhaojun Bai, James Demmel, Jack Dongarra, Axel Ruhe, and Henk van der Vorst. Templates for the solution of algebraic eigenvalue problems: practical guide. SIAM, 2000. URL https: / /www Cs ucdavis. wetaaaiisseuuubaaaaaaa www.:eeieeaaaaaeeeeeaaaaeaaaaaae
html.
Karsten M Borgwardt, Cheng Soon Ong. Stefan Schönauer, SVN Vishwanathan, Alex Smola, and Hans-Peter Kriegel. Protein function prediction via graph kernels. Bioinformatics, 21(suppl__) I): i47-i56, 2005. URL attps::/eeet /tt:://diiiie// http:///....////oooooooaaaatt//a
Viacheslav Borovitskiy, Iskander Azangulov, Alexander Terenin, Peter Mostowsky, Marc Deisen- roth, and Nicolas Durrande. Matérn gaussian processes on graphs. In International Confer- ence on Artificial Intelligence and Statistics, pages //doi. org 10. 28550/arXivy 2010. 15538.
PMLR, 2021. URL https:
Huiyuan Chen, Zhe Xu, Chin-Chia Michael Yeh, Vivian Lai, Yan Zheng, Minghua Xu, and Hang- hang Tong. Masked graph transformer for large-scale recommendation. Proceedings of the 47th International ACM SIGIR Conference on Research and Development in Information Retrieval, pages 2502-2506, 2024. URL https::////. //i..icooooooo http://cc..oo.oo.o...o 214444////222772377777777 3657971
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. arXiv preprint arXiv:2009. 14794, 2020. URL https: //doi//do.////...o org/
48550/arXiv.200 2009. 14794.
Krzysztof Choromanski, Han Lin, Haoxian Chen, Tianyi Zhang, Arijit Sehanobish, Valerii Likhosh- erstov, Jack Parker-Holdes Tamas Sarlos, Adrian Weller, and Thomas Weingarten. From block- toeplitz matrices to differential equations on graphs: towards general theory for scalable masked transformers. In International Conference on Machine Learning, pages 3962-390 -3833. PMLR, 2022. 110.
URL https:tpp/////cc.t. /ow. org 10. 48550/ arwiv..ii.v//.////iii 07709.7999.
Krzysztof Marcin Choromanski. Taming graph kernels with random features. In International Conference Machine Learning, pages 5964 5977. PMLR, 2023. URL https::: /oi. orq/
10. 4 08455855aaXaa 2305 00156.
Fan RK Chung. Spectral graph theory, volume 92. American Mathematical Soc., 1997.
Keenan Crane, Clarisse Weischedel. and Max Wardetzky. The heat method for distance computation. Communications of the ACM, 60(11))90-99, 2017. URL https: / al..am..... org owt/o/ooo/////.////////////////. 1145/3131280.
Anirban Dasgupta, Ravi Kumar, and Tamás Sarlós. A sparse johnson: Lindenstrauss transform. In Proceedings of the forty-second ACM symposium Theory of computing, pages 341-350, 2010. URL https::tp///s.c...p.m /t////ddo. doi http///....1//110000
806737.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Fei-Fei. Imagenet: A large-scale hi- erarchical image database. In 2009 IEEE conference computer vision and pattern recognition, pages 248-24882 255. leee, 2009. URL https: : doi. org 111...1111//.... 2009. 5206848.
Joseph L Doob. Regularity properties of certain families of chance variables. Transactions of the American Mathematical Society, 47(3):455-48 1940.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenbom,, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010. 1929, 2020. URL URL httts:t:: oooii.or//in 1tteg/11.88558 48550/ arxiv 2010. 11929.
12

---

Published as a conference paper at ICLR 2025
Arthur Gretton. Introduction to rkhs, and some simple kernel algorithms. Adv. Top. Mach. Learn. Lecture Conducted from University College London, (6.23)).2 2013. URL https: //www. www.ecuure_iintooottoooo introToRKHS pdf.
gatsby www.uccccc. wtt.a//geeeettoooooorseeeeeeeeee
Grant Van Hom, Oisin Mac Aodha, Yang Song, Yin Cui, Chen Sun, Alexander Shepard, Hartwig Adam, Pietro Perona, and Serge J. Belongie. The inaturalist species classifi- cation and detection dataset. tern Recognition, CVPR 2018, Salt Lake City, UT. USA. June 18-22, 2018, pages 8769- 8778. Computer Vision Foundation / IEEE Computer Society, 2018. doi: 10.1109/CVPR. 2018.00914. URL URL /ttp:////pppccccseseeeeee theovf. com/ content_cvpr.e. 2018 2018/htmm h/tmm// 2tt___20111_ppehemm_mm_a html.
2018 IEEE Conference on Computer Vision and Put-
Var__Horn__hh__aaaturaalail
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are mns: Fast autoregressive transformers with linear attention. In International conference on machine learning, pages 5156-5165. PMLR, 2020. URL https::////// /ttp///oiiiii/ http:////oiooo//./////oo 48550/ arXiv. 2006 16236.
Will Kay, Joao Carreira, Karen Simonyan, Brian Zhang, Chloe Hillier, Sudheendra Vijaya- narasimhan, Fabio Viola, Tim Green, Trevor Back, Paul Natsev, al. The kinetics human ac- tion video dataset. arXiv preprint 48550/ arxiv. 1705.010000 06950.
06950, 2017. URL https: //w.///.///////// org/ 10.
Kyle Kloster and David Gleich. Heat kernel based community detection. In Proceedings of the 20th ACM SIGKDD international conference on Knowledge discovery and data mining, pages 1386-1395, 2014. URL https: //w.i//oii./////////////. org/10.......ooo....o...........
3148.
Risi Imre Kondor and John Lafferty. Diffusion kernels graphs and other discrete struc- tures. In Proceedings of the 19th international conference on machine learning, volume 2002, pages 315-322, 2002. URL https: htmp:://www.... www. ml cmu. www.uueduureeeeeeeeeeee www..//eereeacccaaaaaaaeeee kordor-diifusion-- keroon-eeenlll...e pdf.
Valerii Likhosherstov, Krzysztof M Choromanski, Kumar Avinava Dubey, Frederick Liu, Tamas Sarlos, and Adrian Weller. Dense-exponent eopeeexxponettal random features: sharp positive estimators of the gaussian kernel. Advances in Neural Information Processing Systems, 36, 2024. URL https: //dl. acm. wttp/////////////oo//////t../ret 5555 66262.6622222 05513/33666222.6615555 3666168.
Antoine Liutkus, Ondfej Cıfka, Shih-Lun Wu, Umut Simsekli, Yi-Hsuan Yang, and Gael Richard. Relative positional encoding for transformers with linear complexity. In International Confer- ence on Machine Learning, pages 7067-7079. PMLR, 2021. URL https: //ooi. /10. 48550/ arrivv/ 2839.
Loshchilov, Decoupled weight decay regularization. arxiv preprint arXiv: 2017. URL 48550 /tt//arivvv./. 1711. 05101.
bttts://doi.orrg http://o..oorgoo
Shengjie Luo, Shanda Li, Tianle Cai, He, Dinglan Peng, Shuxin Zheng, Guolin Ke, Liwei Wang. and Yan Liu. Stable, fast and accurate: Kernelized attention with relative positional encoding. Advances Neural Information Processing Systems, 34:22795-22807, 2021. URL https:
/ 100
Ben Mildenhall, Pratul P Srinivasan, Matthew Tancik, Jonathan T Barron, Ravi Ramamoorthi, and Ren Ng. Nerf: Representing scenes as neural radiance fields for view synthesis. Communications of the ACM, 65(1):99-106, 2021. URL https: /stps:///oi... http//dii.ogg/////////// org/10rgg/1.. wwtp/............... 48550/arXiv.... 2003. 08934.
Christopher Morris, Martin Ritzert, Matthias Fey, William L Hamilton, Jan Eric Lenssen, Gaurav Rattan, and Martin Grohe. Weisfeiler and leman go neural: Higher-order graph neural networks. In Proceedings of the AAAI conference on artificial intelligence, volume 33, pages 4602-4609, 2019. URL https: doi. www.ooooooorgoo/ ooo//0../.....o.oo/ooo// 48550/ wttp/aa.iiv... 1810 02244.
Isaac Reid, Krzysztof Choromanski, and Adrian Weller. Quasi-monte carlo graph random features. arXiv preprint arXiv:2305.12470, 2023a. URL https: /oi. org /10./10.44 48550 arxiv. 2305.12470 12470.
13

---

Published as a conference paper at ICLR 2025
Isaac Reid, Krzysztof Marcin Choromanski, Valerii Likhosherstov, and Adrian Weller. Simplex random features. In Internationnl Conference on Machine Learning, pages 28864-288 28888. PMLR, 2023b. URL https::/doi..ooteeeeeeooooo
10. 48550/ arxiv .tl.i......111
Isaac Reid, Eli Berger, Krzysztof Choromanski, and Adrian Weller. Repelling random walks. In International Conference on Learning Representations, 2024a. URL htts:: http::/ddooooggoeooo orggiiorggiii 48550/ arwivv
Isaac Reid, Krzysztof Choromanski, Eli Berger, and Adrian Weller. General graph random features. In International Conference on Learning Representations, 2024b. URL https: /dwii/dii..///// org/ 10. 48.48850// arXiv. 2310. 04859.
Isaac Reid, Stratis Markou, Krzysztof Choromanski, Richard E Turner, and Adrian Weller. Variance- reducing couplings for random features: Perspectives from optimal transport. arXiv preprint 16441 2024c. URL https:/:://// //oii /ttg//......... w8w00/ar....240.......
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomed- ical image segmentation. In Medical image computing and compuler-ussisted intervention- MICCAI 2015: 18th international conference, Munich, Germany, October 5-9, 2015, proceed- ings, part 18, pages 234 Springer, 2015. URL https: //o.: wogg/100-. 48550/ arXiv. 1505. 04597.
Peter Shaw, Jakob Uszkoreit, and Ashish Vaswani. Self-attention with relative position repre- sentations. arXiv preprint arXiv: 02000000255555 155, 2018. URL https:: org/ 110. 48550/ 1eXiv.1803. 02155.
Alexander J Smola and Risi Kondor. Kernels and regularization graphs. In Learning The- ory and Kernel Machines: 16th Annual Conference on Learning Theory and 7th Kernel Work- shop, COLT/Kernel 2003, Washington, DC, USA, August 24-27, 2003. Proceedings, pages 144 158. Springer, 2003. URL petps://peopliae htt.pppp......c.ccc.aa..eeeeeees wwtp:/eei/./.ssssprsrrrr.e
uch wachicaag.coo wetp:..duuurrirrr
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. basee generative modeling through stochastic differential equations. arXiv preprint orXiv:2011.13456, 2020. URL https: /tt.iooor///ooooooooo//..o/ooo /00 48550 arxiv 2011.13456.
A Vaswani. Attention is all you need. Advances in Neural Information Processing Systems, 2017. URL http:: /doi..d.oo/ orgg 100. 10. 48550/ wwt.///axvv...m 1706. 037 62.
Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Liò, and Yoshua Bengio. Graph Attention Networks. International Conference on Learning Representations, 2018. URL https: /tw.///oi..o/////////o. org. 48550/ wttp://ariviv..o Diiiv.11111 1710.1090:
William Whitney, Tatiana Lopez-Guevara, Tobias Pfaff, Yulia Rubanova, Thomas Kipf, Kimberly Stachenfeld, and Kelsey R Allen. Learning 3d particle-based simulators from rgb-d videos. arXiv
preprint 2321v:231..05555, 2023. URL https: / /ooi 05359.
William F Whitney, Jacob Varley, Deepali Jain, Krzysztof Choromanski, Sumeet Singh, and Vikas Sindhwani. Modeling the real world with high-density visual particle dynamies. arXiv preprint 24Xiv:246.198800 2024. URL https: / /oii org/10 10. 11..1.4855555 48550/ arxiv.
Yi Wu, Yanyang Xu, Wenhao Zhu, Guojie Song, Zhouchen Lin, Liang Wang, and Shaoguo Liu. Kdlgt: A linear graph transformer framework via kernel decomposition approach. URL https:
//doi.org// 10. 24 9 ////////aaa/// tttpjjaaa...22222.. 2023/263.
Keyulu Xu, Weihua Hu. Jure Leskovec, and Stefanie Jegelka. How powerful are graph neural networks? arXiv preprint arXiv: 2019. URL httpps///do org 10. 48550/ arxiv. 1810. 1811600826......
Qiangeng Xu, Zexiang Xu, Julien Philip, Sai Bi, Zhixin Shu, Kalyan Sunkavalli, and Ulrich Neu- mann. Point-nerf: based neural radiance fields. In Proceedings of the IEEE/CVF con- ference on computer vision and pattern recognition, pages 5438-5448, 5448, 2022. URL https: //doi. ooboi.org/100 10. 48550/ rxiv. 2201 08845.
14

---

Published as a conference paper at ICLR 2025
Yasutoshi Yajima. One-class support vector machines for recommendation tasks. In Pacific-Asia Conference on Knowledge Discovery and Data Mining, pages 230-239. Springer, 2006. URL 28.
https: http//....../....//1111111111111
Zhilin Yang, Zihang Dai, Yiming Yang, Jaime Carbonell, Russ R Salakhutdinov, and Quoc V Le. Xlnet: Generalized autoregressive pretraining for language understanding. Advances in neural information processing systems, 32, 2019. URL bctps://catt org /tgg/..... 411.//.1..5./////// 405500/55/r/i 1906. 08237.
Chengxuan Ying, Tianle Cai, Shengjie Luo, Shuxin Zheng, Guolin Ke, Di He. Yanming Shen, and Tie Liu. Do transformers really perform badly for graph representation? Advances in Neural Information Processing Systems, 34:28877-28888, 2021. URL https: //oii. org/ 10. 4 8550 arxiiv 2106. 05234.
Yin-Cong Zhi, Yin Cheng Ng, and Xiaowen Dong. Gaussian processes on graphs via spectral kernel learning. JEEE Transactions on Signal and Information Processing over Networks, 9:304-314, 2023. URL https: //tp////oii.////. org/ 10 200.5///rxiv.. 2006.07361.
Bolei Zhou, Âgata Lapedriza, Aditya Khosla, Aude Oliva, and Antonio Torralba. Places: A 10 million image database for scene recognition. IEEE Trans. Pattern Anal. Mach. Intell., 40 (6):1452-1464, 2018. doi: 10.1 10101T/AM..2111277300000 URL https: //t...ogg/1....// 11C9/TPANN 2017. 0TT.::..12.222200000.0 2723009.
Yufan Zhou, Changyou Chen, and Jinhui Xu. Learning manifold implicitly via explicit heat-kernel learning, Advances in Neural Information Processing Systems, 33:477-487, 2020. URL https: / 48550/ aaaiiiiiiiv....i. 2010 01761.
15

---

Published as a conference paper at ICLR 2025
A PROOFS
A.1
PROOF OF LEMMA 3.1
This result was presented by Reid et al. (2024b); a short demonstration included below for the reader's convenience. Observe that
N
M. (9) = Σ
W
W
O
k=U
2m K=U
vmm
k=O pm
(13)
fky Wh
Σ
Skg W
k1=
where $ = En=o Sa W Here, step (1) splits the W into product of W and W with E N; step (2) replaces og by the convolution og: - indices; step (4) arrives at the desired form.
Splk-p: step (3) changes the summation
Convergence. The brief analysis above has made some implicit assumptions. First, in the definition (g( = Ex=o og W we have taken that the Taylor series converges. This occurs if all the eigenvalues of W lie within the series' radius of convergence, meaning p(W) < (lim supn++o V[ann)- Here, p(W) == max{]A]], ANl} denotes the spectral ra- dius. This is common problem in the graph node kernel literature. It is typically ensured by regularising W to control its spectral radius (Chung, 1997). We have also taken that (g(9) is pos- itive definite, meaning all its eigenvalues are greater than than O. This can be ensured e.g. by taking Or > O, again regularising W to control its spectral radius, and appealing to Weyl's 's perturbation inequality (Bai al., 2000).
We have also assumed that (f) exists which satisfies the discrete convolution relation, and that the Taylor series &g = -i= W' also converges. Recall that
00
(1 - (-=2+ --=)+=>x Σ
2 E"
(14)
which is well-known to converge if |||||| < 1. Letting 2 = 1 1-MM((9) (G)
M. (G)+ - E
2
(1 (L = (-1)"
2
(1
a w')".
(15)
7==
R=0
1=0
To converge, we require that p(1 (g) 1. Since M. (g) is positive definite, it sufficient
that (9) < 2. Again, this is achieved by regularising W to control its spectral radius and choosing suitable (a.) Inspecting Eq. 15, it is clear that M. (G) can be expanded in powers of W, so we let
M. (G) =
Σ f.wk K=U
(16)
Substituting this back in, we recover the convolution relationship a; =: Ex-n One find the deconvolution (f given a sequence (at) using an iterative formula, given in Eq. 6 by Reid et al. (2024b). We take & = 0e fa w then Ma( = D since W is symmetric.
The case for implicit kernel learning. We have seen that, to guarantee convergence and positive definiteness of Σx=o 0g W* one must in general control the spectrum of W. This can be achieved by normalising it by suitably downscaling the weights (Chung, 1997; Reid al., 2024b), or choosing (ax) to stay within the radius of convergence. To avoid tricky constrained optimisation, we therefore prefer to instead directly optimise the set of coefficients (f+) up to some maximum order kmux above which ff vanishes. In doing so, we automatically guarantee both convergence and positive definiteness: the learned Dg must be finite and they implicitly defines kernel in feature space (Gretton, 2013). This is why we choose optimise mask parameters in Alg. 1, elegantly sidestepping any convergence and positive definiteness issues.
16

---

Published as a conference paper at ICLR 2025
A.2 PROOF OF THM. 3.2 (NOVEL)
The following concentration bounds are new. Referring back to Eq. 10, recall that graph random features computed by
(G(1))a = (a
(1) )a: q -=1..,,
(17)
where we define the projection vector 15.) S2 RN such that
((vrr)/ven(ura)) T(wig prefix subwalk of (()).
(w 09
Σ
(18)
p(wrr)
Wey Elii
Dig is the set of walks between nodes and eg: p(wig) is the probability of particular walk Wigs known from the sampling mechanism; ((wig) is function that returns the product of weights of edges traversed by win (Eqq 9); len(wig) is the length of Wayi and f : N + R is sequence of a real numbers that determines the graph node kernel to be estimated. Following convention, will assume that we sample random walks with geometrically distributed lengths. halting with probability Phalt 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, Ma (9) We do so by taking n
Maij == ( = n2
1
Σ Σ
((w )
(19)
k1= 1g=1
where E(M aay) Maij since the estimate is unbiased. We sample independent random walks starting from every node of the graph.
Let us now consider single pair of random walks from i, and by respectively. Denote this by the random variable
X2 (w (1) ) c s
(20)
where n such random variables {X are used in total. Suppose we modify Xg. Under mild assumptions, we can prove that the change in the estimator us to use derive concentration inequalities.
AMaijl is bounded, which will enable
This can be understood as follows. To begin, we will bound the -nom of (() Observe that
(kn)W) p(w)
Wayld, & = c, (21)
(1) )11
I max EQ
Σ
S f
max IIRREN
Phalt
O===
which corresponds to the case of an infinite length walk that happens almost never. Here, w p.s. w means that w a 'prefix subwalk' of w/ so the first len(w) nodes w!()) are exactly the nodes of w.
(i)
Suppose that is finite. This can be guaranteed for bounded f by normalisation of W, multiplying W.,d
by a scalar so that max I-p < 1. More pragmatically, it can be achieved by stipulating that there exists some integer mmax E N such that f. = 0 > imax so we only consider walks up to certain length. This a natural choice in experiments since we want 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 graph of finite size.
We have seen that 0 < [|
(8)
()) s e (k1, K2, (k,, i,j), || < It follows that - s (( (e(w
whereupon the change in May if we modify X₁ obeys
2n n²
Maa S
2
(22)
With straightforward application of McDiarmid's inequality (Doob, 1940), it follows that
n 2(2n - 1)
Pr (1)g(v:) (o(())) - Maij > 1) S 2exp
(23)
17

---

Published as a conference paper at ICLR 2025
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.
B EXTENDED DISCUSSION ON RELATED WORK
In this appendix, we give detailed explanation of the relationship to previous work.
1. Stochastic position encoding (SPE) (Liutkus al., 2021). To our knowledge, this is the earliest RPE strategy compatible with (N)) transformers, incorporating information about the 'lag' between tokens in the linear attention setting by drawing samples from random processes with prescribed covariance structure. The time complexity is ((N) ) but the method is only applicable to sequence data, i.e. the very special case that 9 is a 1-dimensional grid.
2. Toeplitz masks and the fast Fourier transform (FFT) (Luo et al., 2021). Again consid- ering sequence data, suppose that RPE between a query qi and a key Kü is impemented by adding bias term b-- to the dot product before exponentiating. Then the corresponding mask My = ehs i is Toeplitz (constant on all diagonals). From numerical linear algebra, multiplying a Toeplitz matrix and a vector can be achieved with O(N log N) time complexity rather than O(N2), using the FFT. This is subquadratic, though not linear, and again only specific considered.
Graphormer (Ying et al., 2021). 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., 2022). This paper proves that subquadratic matrix-vector multiplication is sufficient for masking to be applied efficiently, and suggests number of algorithms to achieve this in special cases.
Block Toeplitz matrices. This method extends the work of Luo et al. (2021) to show that the analogous mask on a more general d-dimensional grid is d-level block Toeplitz. Hence, this mask parameterisation also supports O log N) masking, but now for a slightly more general class of graphs
Affine mappings on forests. If is forest and Mij exp(r(dist(v,, U2) ))), where dist(n;, 22) is the shortest path distance between nodes n: and by and + an affine map- ping, then topological masking can be implemented in O(V). 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 (2002) derived a closed form for the heat kernel on hypercube, which also turns out give Toeplitz mask that can be applied in (N log N). This special instantiation of our power series masks (W) for specific choice of 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 tak- ing the dot product between two 'frequency vectors', recording the nodes visited by the random walks beginning at the two nodes of interest. This method the closest to our approach, but with some crucial differences: (i) we importance sampling of random walks, which we find to be crucial for good performance (see App. C); (ii) we use learn- able 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 en- joys certain regularisation properties (Reid et al., 2024b). More importantly, we are able to derive strong concentration bounds that prove that our method is O(N), whereas Choro- manski et al. (2022) do not explicitly guarantee any time complexities.
Centrality degree encoding (Chen et al., 2024). This paper takes the mask Mý [n({[((())) + (d,)))). where d(s,) is the degree of node m(e.) and z N (0, 1) is learnable
18

---

Published as a conference paper at ICLR 2025
ImageNct
iNaturalist2021
Places365 (Small)
1M
2.7M
1.8M
Num. layers
12
12
Num. heads
12
12
12
Num. patches
16 x 16
16 x 16
16 x 16
Hidden size
768
768
768
MLP dim.
3072
3072
3072
Optimiser
Adam
Adam
Adam 5
Epochs
90
5
Base learning rate
x 10
x 10-a
3 10 3
Final learning rate
1 x 10
x 10-5
1 x 10-5
Learning rate schedule
Linear warmup (10" steps), constant, cosine decay
Batch size
4096
64
64
(-))
ReLU(()
ReL.U(-)
ReLUU())
Pretrained?
/
Num. walks
100
20
20
Phalt
0.1
0.1
0.1
Max. walk length (6aax)
100
10
10
Table 2: Architecture, hyperparameters and training details for ViT experiments.
function thereof. This can be implemented with linear transformers on general graphs in ((N) time complexity, but since it is completely insensitive the relative positions of the nodes in the graph it cannot be considered a general topological masking mechanism.
Kernel decomposition linear graph transformer (KDLGT; Wu et al.) This paper also tempts to efficiently incorporate topological information into graph transformers by defining graph node kernel with a low-rank decomposition. It is based off a 'structure extractor' for subgraph centred at each respective node the 'shortest anchor path distance'. However, bias is incorporated as an additive term, rather multiplicatively which is our chief goal (Eq, 8).
Graph attention networks (Veličković et al., 2018). 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 efficient masking of linear attention transformers. Our mask parameterisation is more general than this hard nearest-neige neighbous example, but shares the finding that modulating attention by local topological structure boosts performance.
C FURTHER EXPERIMENTS AND DETAILS
Here, 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.
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 termination probability Phat 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 fea- tures (pg(t))) (r:)), less sparse increases computational cost. Since the estimator is unbiased, the n + oo limit gives dense features {Øg (v;) b.ee whose dot product exactly Taylor mask (g).
Importance sampling. The crux of our approach is modulating attention by an unbiased Monte Carlo approximation of function of weighted adjacency matrix W. Using GRFs, this involves
19

---

Published as a conference paper at ICLR 2025
sampling random walks which reweight by (i) their probability under the sampling mechanism, (ii) the product of traversed edge weights ((w), and (iii) a learnable function f (see Eq. 10). To investigate the importance of this principled sampling mechanism, ablate (i) and (ii) and compare with an ad-hoc empirical random walk kernel defined by the features
Pod-hoc (1i)g
1 Ti
(intwne)((wp prefix subwalk of w q q= I N. E0
Σ
(24)
Eq. 24 still simulates random walks and records their destinations, giving empirical measure of the overlap between respective nodes' neighborhoods. However, in absence of dependence on (ww) and p(w) this cannot be interpreted approach employed by Choromanski et al. (2022) but with extra learnability via f.
MC estimate of a function of W. This is similar to the
Comparing ViT performance with = 100 random walks and Phalt 0.5, the results are striking; see Table 3. contrast to our method, the mask defined by Eq. 24 worse than the unmasked base- line. This shows the crucial importance of the theory in Sec. 3.4. Intuitively, without upweighting long, improbable walks by p(w) it difficult to capture long-range dependencies in the topologi- cal mask.
Fully learnable features. As a final check, we endow each graph with a fully learnable, dense N N-dimensional feature ddene, learned (e:) ) IN,, relaxing the parameterisation
a function of weighted adjacency matrix (g)... This uses many more parameters and is O(N) rather than ONN. but it technically includes GRFs all structural inductive bias from G. allowing the model to implicitly learn a masking kernel without any constraints (except feature dimension). The model performs badly, providing only very slight improvement over the unmasked linear baseline. See T.ble 3.
special case. In making this choice have removed
Figure 6: Final test accuracy number of walkers n. Monte Carlo mask variance drops with the number of walks predictions im- prove.
Table 3: Ablation studies. Importance sam- pling of random walks and mask parameteri- sation as function of W are crucial for good performance.
Predictive performance vs. sample size
0.74
Variant
Test acc.
nocumicy
GRF-masked
GRF-masked
0.730
Unmasked
Unmasked
0.693
M.
Ad-hoc masked (Eq. 24)
0.689
Fully learnable features
0.696
10
Number of walkers,
20

---

Published as a conference paper at ICLR 2025
Time complexity scaling with number of walkers. Readers may also be interested in how the total number of FLOPs scales with the number of walkers for given graph 9 and termination probability Phalt- There are two main contributions consider. First, must sample the random walks, the cost of which is (n). This is done once the start of training, so is a big consideration in practice. Second, the number of walks determines the sparsity of the GRF, and therefore the time complexity of implicit masking via Eq. 7. We have see that this is at most (n)) but the details will vary depending on the graph topology (see I Lemma 3.3). Fig. 7 gives exam- ple for a ID grid with n = 64, taking Phalt = 0.5 and d 8. It reflects the behaviour described above.
Computation vs. walkers
0.25
Incar GRFF
0.20
FLOPs (/309
0.10
150 200
250
Number walklesss
Figure 7:
Time complexity scaling wrt number of walkers
Linear attention variant. As a brief note, our topological masking algorithm is compatible with different choices for the nonlinearity d(.). In the text have generally taken Φ(-) ReLU(-) so that m = d, but other choices are possible. For instance, one could use a randomised mapping that approximates the softmax kernel (Choromanski et al., 2020; Likhosherstov et al., 2024), which is commonly referred to as the Performer.
To demonstrate, we replace the ReLU features by FAVOR+ Monte Carlo attention approximation, considering the same VIT architecture and training as reported in Sec. 4 (see Table 2). For ImageNet, the unmasked Performer achieves a test accuracy of 0.693. Meanwhile, the GRF-masked performer achieves a test accuracy of 0. 0.734 - 3.9% absolute improvement. This shows how Alg. can be expected to provide gains across a range of linear attention mechanisms.
C.3 HD-VPD EXPERIMENTAL DETAILS
Here, we describe the high density visual particle dynamics (HD-VPD) experiment from Sec. in more detail. The setup closely follows that (Whitney et al., 2024), 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 video prediction. The HDDVVPD 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, 1. 2. Each RGB image is encoded into per-pixel features using U-Net (Ron- neberger et al., 2015), and unprojected using known camera intrinsics and extrinsics to a point cloud (æ;; ft E Nguec with E R and f E R For every timestep. the particles across different cameras are merged and subsampled uniformly at random to get N = 32768. The robot actions at timesteps ( 2, 10-1,,1 1,, representing the robot joints and where they plan to move, are also encoded as kinematic particles.
Predict the point cloud dynamics. With an Interlacer (see below), use the point kinematic particles to predict (Av;, Af) and hence the updated state (xr) (x:--+ Aa; S! -!-1+
t clouds and
3. Render the new state to an image. Use a ray-based renderer similar to Point-NeRF (Xu et al., 2022) 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 Le loss between the predicted and observed RGB values.
Interlacers as PCTs. The Interlacer is type of point cloud transformer (PCT) designed to scale to tens of thousands of particles. It alternates linear-ateention transformer layers that capture global dependencies between tokens with local nnighbour-attender layers that model fine-grained geomet- ric structure (see Fig. by Whitney et al. (2024)). This is reported to outcompete both graph neural
21

---

Published as a conference paper at ICLR 2025
networks and vanilla linear- attentnn transformers. For the neighbour-attende the authors use message passing-type algorithm that updates the features 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. (2024), 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 3 = 27 repelling random walks (Reid al.. 2023b). 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 ((-) = ReLU(()) and compute GRF-masked attention with single head, projecting the queries and keys with dense layer.
3. To construct G, find k neighbours for every node. This is done the fly since the neighbours can change at every timestep.
4. For Fig. 4, we train for 100k steps, using a learning rate of x 10 until step 1000 then 1x 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, 2017) with weight decay 10-3, clipping the gradient norm to 0.01; models are trained with 6 step rollouts, with losses computed constrained to [-0.25, 0.25] by a scaled tanh nonlinearity so particles do not move more than 25cm per timestep. For the rendering, like Whitney et al. (2024) we use four concentric annular kernels of radii [o, 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 1m and the far plane is 2m, with the background set to solid white.
128 sampled rays; Aæ; are
C.4 VIVIT: TOPOLOGICAL MASKING FOR VIDEO DATA
Here, we present preliminary results for incorporating topological masking with graph random features into video vision transformers (ViVIT; Arnab et al., 2021). Prompted by the original paper, we use a factorised en- coder, where the spatial encoder follows the VIT archi- tecture from earlier in the paper (see Sec. and App. C.I) with patch size 16 x 16 x 2. The final index mixes pairs of successive frames. The temporal encoder, which subse- quently models interactions between the frame-level rep- resentations, 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., 2017). Table shows the results for (i) unmasked softmax attention and (ii) softmax attention with GRFs. Topological masking boosts predictive performance by +0.4%. To our knowledge, this is the first application topological masking to the video modality.
Table 4: VIVIT test accuracies the Kinetics benchmark. Topological mask- ing with GRFs boosts performance.
VARIANT
TEST ACC.
SOFTMAX
0.754
SOFTMAX GRFs
0.758
D ASYMMETRIC GRAPH RANDOM FEATURES
In this Appendix, we present 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., 2024b), 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., 2023) because it can be straightforwardly swapped into existing code.
22

---

Published as a conference paper at ICLR 2025
Asymmetric GRFs. is straightforward that Eq. 13 of App. A.1 generalises to
(.
W
Σ
w
(25)
N =0
if the discrete convolution
f(1) (2) Jp kkp ak
(26)
holds. We do not necessarily require that f(1) = f(2); this is a special case. Another special case foa {1 if p O, 0 otherwise), whereupon we need f/ ap Vp. Sampling GRFs in this manner will still give an unbiased estimate of M. (G),,, but with different concentration properties.
Computational benefits. Supposing that fo {1 O, 0 otherwise), og (f (t:)). = I[ n), one-hot vector only nonzero at the coordinate corresponding to ti. In other words, we do not need simulate any random walks obtain GRFs for the keys; o Iy with Iy RNN the identity matrix. It follows that
(9) 2
. = 10g (x:)]
(27)
But (( (v:) is only nonzero at nodes visited by the ensemble of walkers beginning Vi. This makes masked attention straightforward to efficiently compute, obviating the outer products and vectorisation. Simply.
N Σ (
( ( 2
(28)
N where Aij = exp(q, kj) (softmax) or A.j = q. ky (linear) and (u, j== are the value vectors. Since 0g (w;) is sparse, the sum in Eq. 28 is computed in constant time so attention computed in ((N). Even more explicitly, for an ensemble walks {w JoEN, ke[1.n]. we can rewrite the above as
N
A F(w) p(w)
ΣΣ k=1 w.w.s.
(29)
where w[-1] stands for the final node of the walk w and p.s. means 'prefix subwalk' (see App. A.2). The number of FLOPs to compute attention for node v; depends the length of all the geometrically distributed walks beginning at node U, but we have seen that this is independent of graph size (Sec. 3.4). Computing the entire matrix is only O(N). We give pseudocode for this approach in Alg. 2.
Benefits and limitations asymmetric GRFs. The chief benefit of asymmetric GRFs is that Eq. 29 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 for the HD- VPD experiments because it can be easily integrated with the existing code: instead of finding each node's 's (approximate) k-nearest neighbours and implementing the complicated message passing-type algorithm by Whitney et al. (2024), we find the 3-hop neighbourhood and straightforwardyy compute weighted attention with a learnalle function of walk length. The drawback that the variance the mask estimate with asymmetric GRFs tends to be greater, though proving this is an open problem.
E ADDENDUM: ADDITIONAL. GRF RESULTS
Thm. 3.2 provides novel bound on the GRF estimator concentration, specifying the number of walkers n and termination probability Paalt needed to ensure sharp mask estimates with high proba- bility. Under mild assumptions graph (namely, that does not change N grows), this can be used to upper bound GRF sparsity and thus the time complexity of the algorithm. The bound de- rived in Lemma 3.3 implicitly assumes the worst case that walkers never backtrack and do not visit
23

---

Published as a conference paper at ICLR 2025
Algorithm 2 ((N) topological masking with asymmetrii GRFs
Input: query matrix E R Nxd key matrix K E I xxd value matrix V R Nxd graph with weighted adjacency matrix W RNxN learnable mask parameters (fi number of random walks to sample n E N, query/key feature map o(-) I I!". Output:
masked attention All, LR,M (Q, K, V.9) in ((N) time, without using sparse linear algebra libraries.
1: Simulate n terminating random walks twt H Fi- out of every node v; E
2: for U; E EN do
3: Initialise output feature, f. < Oad
4: Initialise attention normalisation, S. -
5: for (1) in {wx jk- J do for 15) in enumerate(w do
6:
f. += ((q)) (kk, fi Md Vj
7:
[(w [+]]
( [
8:
S, (q.)) ((kk) fe
p( [t]
9: end for
10: end for
11: Normalise output feature, fi/ = (S, x n))
12: end for
13: Return low-rank attention with topological masking in O(N) time by
Att LR.M (Q, K, V, V) [f]]
(30)
any of the same nodes, so in practice how tight this bound may depend on the particular graph being considered. In particular, it will tend to be tighter for denser, bigger graphs. Our algorithm is even more efficient for smaller, sparser graphs. This behaviour, beyond asymptotic time complexity results, may be of interest to practitioners.
To supplement with some empirical results, we consider using GRFs estimate the heat kernel cx)(o"W), where E R a lengthscale parameter and Wi = 1/Vd,d, as usual. We consider Erdós-Rénn graphs, where edges are independent Bernoulli random variables with probability p.
Scaling with graph sparsity. The left pane of Fig. shows the Frobenius norm error on the mask estimate, |M- M| /|| where M is the estimate, vs. the edge generation probability p. The take
4 walks per node with 0.5. We average the results over 100 trials for standard errors. As expected, the approximation error becomes a little worse for denser graphs, though it remains excellent throughout. The right hand pane of Fig. 8 shows the RF sparsity (proportion of nonzero entries in E (QQK)).)) over the same parameters. Again, the number of nonzero entries grows as graphs become denser because visiting the same node multiple times is less likely, but GRFss are sparse in every case. We emphasise once more that the bounds still hold in every case; this is simply a malter of how tight they are.
number of nodes is N = 100, and
Scaling with graph size. Fig. performs the same analysis for Erdös-Rényi graphs with fixed edge probability p 0.5, now varying the size of the graph N. The approximation error grows slightly for bigger graphs but soon platcaus, and again remains excellent for all graphs considered. The GRF sparsity drops because the number of visited nodes is ((1) in every instance, meaning the proportion on nonzero GRF entries must scale as (1/N).
We have included this section for the convenience of the interested reader. For more exhaustive details, we invite them to read the previous works of Reid et al. (2024b) and Choromanski (2023), which first introduced GRFs and explored their properties.
24

---

Published as a conference paper at ICLR 2025
sparsity probubility
Figure GRF approximation quality and sparsity for Erdós-Rény. graphs of different sparsities.
Approxiiation ertor vs. nlb vertioss
RF sparsity vs. nb. veriices
Figure 9: GRF approximation quality and sparsity for Erdōs (en6is-ééyiiao graphs of different sizes.
25


Reviews:
[{"review_id": "4VUgSfDxFS", "reviewer": "Reviewer_Ptpy", "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.\n\nWhile 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.", "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:\n1. Correct the theoretical analysis\n2. Properly characterize complexity and mention which category of graphs it can address\n3. Either prove better bounds under specific assumptions or acknowledge limitations\n4. Frame results in terms of parameterized complexity\n\nThe authors should consider addressing these fundamental issues during rebuttal, if possible.", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. The paper provides strong theoretical foundations with proven concentration bounds and complexity guarantees for GRFs.\n2. The method shows concrete performance improvements on real-world tasks and scales to large problems (>30k nodes) that would be intractable with quadratic approaches.\n3. The approach can be implemented with both symmetric and asymmetric GRFs, offering different trade-offs between computational efficiency and variance in mask estimation.\n4. The experiments cover diverse applications (images, point clouds, videos) and include detailed ablation studies.", "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:\nIn 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. \nLet'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) )$. \nAt any step, the walk can move to another vertex with probability 1/(N-1). \nFor 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.\n\nAs 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 \n\nCase 1: Small complete graph\n- N = 10 nodes\n- Each node has degree 9\n- Even a 1-hop walk can reach 9 other nodes\n- A 2-hop walk can reach all nodes\n\nCase 2: Larger complete graph\n- N = 1000 nodes\n- Each node has degree 999\n- A 1-hop walk can reach 999 other nodes\n- A 2-hop walk can reach all nodes\n\nWhile 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). \n\nThis demonstrates that actual sparsity heavily depends on the structure of the graph. \n\nThis 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. \n\nThe 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)...\" \nThis 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. \n\nThe authors should make this bounded-degree assumption explicit upfront in the introduction and modify their claims about general graphs. \n\nA 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. \n\nInconsistent and confusing notation:\n- $(f_k)_{k=0}^\\infty$ is sometimes treated as a sequence of reals, sometimes as a function\n- Incorrect set notation: claiming (f_k)_{k=0}^∞ ⊂ R when sequences are functions from N to R\n- Weighted adjacency matrix definition issues:\n  * Claim W is weighted but then suggest $w_{ij} = 1/sqrt(d_i d_j)$\n  * This normalization discards meaningful edge weights in attention context\n\nii) Missing crucial assumptions:\n- No explicit assumptions about graph structure\n- No discussion of how graph density affects complexity\n- No proper analysis of how maximum degree impacts sparsity\n\nThe 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. \n\nIn 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "FYPnJffqsU", "reviewer": "Reviewer_UHQw", "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.", "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.\n\n1. Could the authors provide examples or case studies that apply GRFs to commonly used graph datasets, such as Cora or Citeseer?\n2. Could the authors also include a comparison of computational times between your methods and baseline approaches?\n\nThanks 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.", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The prposed method shares $O(n)$ time complexity and suitable for the relatively large scale input.", "weaknesses": "see question.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "pAJkEYeuI3", "reviewer": "Reviewer_BeTo", "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.", "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)?\n2. Remark 3.1: \n   - This remark is used as a lemma. Better state it as such.\n   - 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.\n3. Ln. 115-116: \"$\\Phi_{Q,K} \\in$\" should probably be \"$\\Phi_Q, \\Phi_K \\in$\"\n4. Ln. 184-185: Statement is unclear. Why should it necessarily be faster?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 1, "contribution": 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.\n2. 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. \n3. Their method demonstrates improved predictive performance in various learning tasks.", "weaknesses": "1. Dense and unclear presentation: \n   - 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.\n   - Algorithmic descriptions, such as those in Algorithm 1, are highly abstract and may be difficult to follow. \n   \n   Without clearer explanations, the accessibility of the paper is reduced.\n2. The paper lacks discussion of the method limitations. For example:\n   - 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.\n   - For dynamic or evolving graphs, precomputing random walks is not feasible, and recomputing them on the fly could reduce efficiency. \n   - 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.\n\n   These limitations should be acknowledged and discussed for a more balanced perspective.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "7ihEymdN8h", "reviewer": "Reviewer_w6hQ", "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.", "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?\n\n2. 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? \n\n3. Besides GRF, are there other methods to do implicit graph masking in equation (8)? How's the performance?\n\nIf all my concerns are resolved properly, I will be happy to increase my score.", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper has a good motivation for introducing linear topological masking of low-rank attention. \n\n2. 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.\n\n3. The explanation is clear, the figures are illustrative, and the writing is well-structured.", "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "NkBOGcys2f", "reviewer": "Reviewer_QFRp", "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.", "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.\n* The citation in line 161 (Borgwardt et al., 2005) is about graph kernels, not graph node kernels if I understand it well.\n* 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.\n* In section 4, maybe it's better to use \\textit{subsection} for each experiment instead of \\textit{paragraph}.\n* 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?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "* The paper is well-organized and well-written, it's a pleasure to read.\n* 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.\n* The experiments are well-presented, which show the relevance of the proposition for the situation where $N$ is large.", "weaknesses": "Please see the questions raised.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "t5ZGesj3sN", "reviewer": "Reviewer_uZmX", "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.", "questions": "See the weakness.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "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.\n2. Generally, the notation and theorem in the paper is well defined and illustrated. \n3. The authors provide the theoretical analysis and ablation study to show the effectiveness of the proposed method.", "weaknesses": "#### Major problems\n1. 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.\n2. 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.\n3. It's better to include more experiments and evaluations on traditional graph datasets.\n\n[A] Wu, Qitian, et al. \"Nodeformer: A scalable graph structure learning transformer for node classification.\" NIPS2022.  \n[B] Wu, Yi, et al. \"KDLGT: A Linear Graph Transformer Framework via Kernel Decomposition Approach.\" IJCAI 2023.\n#### Minor issues\n1. Confusing notations in line 148, the size of the graph N ($\\mathcal{N}$)?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)