Linear Transformer Topological Masking with Graph Random Features

Abstract: When training transformers on graph-structured data, incorporating information about the underlying topology is crucial for good performance. Topological masking, a type of relative position encoding, achieves this by upweighting or downweighting attention depending on the relationship between the query and keys in a graph. In this paper, we propose to parameterise topological masks as a learnable function of a weighted adjacency matrix -- a novel, flexible approach which incorporates a strong structural inductive bias. By approximating this mask with graph random features (for which we prove the first known concentration bounds), we show how this can be made fully compatible with linear attention, preserving $\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 tasks on image and point cloud data, including with $>30$k nodes.

• The paper introduces the use of graph node kernels to parameterize topological masks, offering a flexible and inductive bias for transformers.
• It employs graph random features for efficient mask approximation, reducing computational complexity to O(N) while maintaining performance.
• The work provides concentration bounds and sparsity guarantees, enabling efficient matrix-vector multiplications for large graph data.

Linear Transformer Topological Masking with Graph Random Features

This paper presents a novel approach to efficiently incorporate graph-structured data into transformers through topological masking. The proposed method parameterizes topological masks as a learnable function of a weighted adjacency matrix. Using graph random features (GRFs), the approach achieves $\mathcal{O}(N)$ time and space complexity, significantly improving upon the previous $\mathcal{O}(N \log N)$ alternatives.

Key Contributions

1. Parameterization with Graph Node Kernels: The paper introduces the idea of using graph node kernels to parameterize the topological masks. This method captures structural information through a power series function of the adjacency matrix, offering flexibility and a strong inductive bias.
2. Efficient Masking Using Graph Random Features: To address the challenges of low-rank transformer settings, the authors propose using GRFs to approximate the masks efficiently. GRFs provide an unbiased estimate of a function of the adjacency matrix, supporting $\mathcal{O}(N)$ complexity while maintaining strong performance gains.
3. Concentration Bounds and Sparsity Guarantees: The paper provides the first concentration bounds for GRFs, ensuring stable and reliable approximations of the topological masks. It proves that GRFs are inherently sparse, allowing for efficient matrix-vector multiplications within the transformers.

Theoretical Insights

The introduction of GRFs revolutionizes the ability to incorporate structural information into transformers without increasing computational burden. This not only improves efficiency but also facilitates the handling of larger graphs, where full-rank approaches would be prohibitive.

Experimental Evaluations

• Image and Video Data: The method was tested on diverse datasets, such as ImageNet and iNaturalist, where GRF-masked transformers showed consistent improvements over unmasked and previously proposed masked alternatives.
• Robotics Applications: The paper also demonstrates impressive gains in modeling high-density particle dynamics for robotics applications, where the complexity and scale of point cloud data necessitate efficient computation.

Practical and Theoretical Implications

The incorporation of graph-structured data into transformers using GRFs enhances model expressivity, making it a promising tool for a wide range of applications, from vision to robotics. The efficiency and learnability of the approach open new avenues for scaling transformers with respect to graph size and complexity.

Future Directions

Future work could explore more diverse graph kernels and further optimization of sparse linear algebra libraries to maximize efficiency gains. The potential adaptation and integration of these methods into broader AI frameworks appear promising, particularly in real-time and large-scale applications.

In sum, the paper presents a comprehensive and efficient method for utilizing graph structures within transformers, offering significant improvements in scalability and performance across multiple domains.