Linear Transformer Topological Masking with Graph Random Features
transformer;linear;attention;graph;random walk;Monte Carlo;encoding;topological masking;point cloud;Performer
(Part 2/2)

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

*Minor points*:
1. $f: \mathbb{N} \to \mathbb{R}$ is a function mapping from the natural numbers to real numbers. $(f_k)^\infty_{k=0}$ is a sequence of reals, intended to denote the evaluations of some particular $f$ for the natural numbers rather than a sequence of functions. To make this clearer, we have updated the latter to $(f(k))_{k=0}^\infty$. Thanks for the suggestion.
2. Normalisation of $\mathbf{W}$. Please see earlier comments. This is standard practice to ensure that the kernel converges.

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

__________
[1] Kernels and regularization on graphs, Smola and Kondor, COLT 2003, https://people.cs.uchicago.edu/~risi/papers/SmolaKondor.pdf 

[2] Spectral graph theory, Chung, 2007, https://mathweb.ucsd.edu/~fan/research/revised.html

[3] General graph random features, Reid et al., ICLR 2024, https://arxiv.org/pdf/2310.04859 

[4] Rethinking attention with Performers, Choromanski et al., ICLR 2021, https://arxiv.org/pdf/2009.14794
Thanks for the review (2/2)
(Part 1/2)

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

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

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

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

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

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

CONTINUED BELOW.
Thanks for the review (1/2)
Yes
No ethics review needed.
While the empirical results might be interesting, the fundamental theoretical claims that form the paper's main contribution are incorrect. A major revision would be needed to:
1. Correct the theoretical analysis
2. Properly characterize complexity and mention which category of graphs it can address
3. Either prove better bounds under specific assumptions or acknowledge limitations
4. Frame results in terms of parameterized complexity

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

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

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

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

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

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

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

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

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

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

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

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

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

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

In the experiments section, it would be interesting to compare against other major linear attention variants like (i) Performers (Choromanski et al., 2020) Favor+ which uses kernel approximations and (iii) Nyströmformer which uses Nystrom’s method to approximate attention.
Yes
No ethics review needed.
I appreciate the authors’ valuable contributions in this area. As I am less familiar with applications in image, point cloud, or robotics contexts, I am particularly interested in understanding how Graph Random Features (GRFs) benefit graph neural networks on traditional graph datasets.

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

Thanks for this important work, and apologies for my gaps in my background knowledge.  I would also kindly request that the Area Chair reduce the weight of my review in the final evaluation.
1. The prposed method shares $O(n)$ time complexity and suitable for the relatively large scale input.
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.
see question.
Yes
No ethics review needed.
1. Eq. (4): The power series is generally not guaranteed to converge. It is better to specify clearly the underlying assumptions on W and alpha that guarantee convergence. Is $\alpha_0$ assumed here to equal 1, as in (Reid et al. 2024b)?
2. Remark 3.1: 
   - This remark is used as a lemma. Better state it as such.
   - It is in general not guaranteed that alpha has a deconvolution. Is it an assumption of Remark 3.1? Or is it guaranteed by some other assumption? Better clarify.
3. Ln. 115-116: "$\Phi_{Q,K} \in$" should probably be "$\Phi_Q, \Phi_K \in$"
4. Ln. 184-185: Statement is unclear. Why should it necessarily be faster?
1. Their method is the first to achieve $\mathcal{O}(N)$-time complexity for computing masked attention for general graphs, $N$ being the number of vertices.
2. The paper provides the first known concentration bounds for GRFs and rigorous sparsity guarantees. These theoretical insights are valuable, potentially extending beyond transformers to other domains that rely on scalable graph-based representations. 
3. Their method demonstrates improved predictive performance in various learning tasks.
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.
1. Dense and unclear presentation: 
   - While the method is theoretically sound, the presentation is mathematically dense and lacks clear explanations. This may pose a barrier to readers, particularly those less familiar with GRFs. In particular, the technical exposition in lines 184–254 is notation-heavy and unclear.
   - Algorithmic descriptions, such as those in Algorithm 1, are highly abstract and may be difficult to follow. 
   
   Without clearer explanations, the accessibility of the paper is reduced.
2. The paper lacks discussion of the method limitations. For example:
   - The practical applicability of this method depends heavily on the specifics of the graph structure and the task requirements, since it relies on approximations with random walks. In  graphs where relevant information is distributed over long distances or requires traversing multiple nodes, random walks may fail to capture the full structure efficiently.
   - For dynamic or evolving graphs, precomputing random walks is not feasible, and recomputing them on the fly could reduce efficiency. 
   - Since random walks introduce stochasticity, their effectiveness can vary based on the number of walks and the chosen halting probability. This means that the quality of topological masking may be sensitive to hyperparameters like the number of walks and the stopping probability, making it challenging to generalize the method across different graph structures.

   These limitations should be acknowledged and discussed for a more balanced perspective.
Yes
No ethics review needed.
1. While the author shows the test accuracies in Table 1, can the author also present other results, like the total training time or total flops, to validate the proposed method’s efficiency?

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

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

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

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

3. The explanation is clear, the figures are illustrative, and the writing is well-structured.
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.
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.
Yes
No ethics review needed.
* While the true time complexity of the proposed algo is $\mathcal{O}(Nmd)$, most part of the paper omit $m$ and $d$. I think the authors should make this point clear since in some cases where $m$ and $d$ can be large enough.
* The citation in line 161 (Borgwardt et al., 2005) is about graph kernels, not graph node kernels if I understand it well.
* Though the paper shows experimentally the impact of the number of random walks $n$ on the performance, I would also like to see its impact on the computation time.
* In section 4, maybe it's better to use \textit{subsection} for each experiment instead of \textit{paragraph}.
* This question is not about the contribution of the paper, but the general idea of using structural graph masking for transformer. Isn't it a reinvention of the graph neural network by integrating a matrix representing the structure information into Transformer? Can author make a (possible) link between the two?
* The paper is well-organized and well-written, it's a pleasure to read.
* The reasoning of the idea is clear. We can understand very well why we need to use kernel gram matrix as the masking matrix and why we need the graph random features.
* The experiments are well-presented, which show the relevance of the proposition for the situation where $N$ is large.
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.
Please see the questions raised.
Yes
N/A
No ethics review needed.
See the weakness.
1. The research topic is highly important and intriguing and the authors provide experimental results on some larger scale image and point-cloud datasets.
2. Generally, the notation and theorem in the paper is well defined and illustrated. 
3. The authors provide the theoretical analysis and ablation study to show the effectiveness of the proposed method.
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.
#### Major problems
1. The background and related work section is limited. To my knowledge, this is not the first work to incorporate the inductive bias inherent in graph nodes into GNN and transformer-like models.
2. Additionally, improved time and space complexity have been discussed in existing papers such as [A, B]. Therefore, the benefits and differences between linear attention and the proposed topological masking should be thoroughly explained.
3. It's better to include more experiments and evaluations on traditional graph datasets.

[A] Wu, Qitian, et al. "Nodeformer: A scalable graph structure learning transformer for node classification." NIPS2022.  
[B] Wu, Yi, et al. "KDLGT: A Linear Graph Transformer Framework via Kernel Decomposition Approach." IJCAI 2023.
#### Minor issues
1. Confusing notations in line 148, the size of the graph N ($\mathcal{N}$)?
An efficient algorithm to incorporate structural information about the underlying graph topology into linear attention transformers
@inproceedings{
anonymous2024linear,
title={Linear Transformer Topological Masking with Graph Random Features},
author={Anonymous},
booktitle={Submitted to The Thirteenth International Conference on Learning Representations},
year={2024},
url={https://openreview.net/forum?id=6MBqQLp17E},
note={under review}
}
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.
I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.
I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.
transformer linear attention graph random walk Monte Carlo encoding topological masking point cloud Performer
I certify that there is no acknowledgement section in this submission for double blind review.
/pdf/6fad354a76e7d7bea4cc2d94b1ca64d71ecfc54b.pdf
learning on graphs and other geometries & topologies
I understand the reciprocal reviewing requirement as described on https://iclr.cc/Conferences/2025/CallForPapers. If none of the authors are registered as a reviewer, it may result in a desk rejection at the discretion of the program chairs. To request an exception, please complete this form at https://forms.gle/Huojr6VjkFxiQsUp6.
I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.
/attachment/6bd66a6477459bd9c606d005e1a4ebcde3781140.zip
Linear Transformer Topological Masking with Graph Random Features
ICLR 2025 Conference Submission
ICLR.cc/2025/Conference/Submission