SFi-Former: Sparse Flow Induced Attention for Graph Transformer

TL;DR

SFi-Former introduces Sparse-Flow-induced (SFi) attention, an energy-based sparse attention mechanism for Graph Transformers that minimizes a flow energy with an $\ell_1$ penalty to produce selective, sparse information transfer. By interpreting attention as network flows on a resistor-like graph and incorporating friction terms, the method derives sparse attention patterns and iteratively computes them via proximal optimization, integrating them with adjacent graph structure through an adjacency-enhanced architecture. The approach achieves state-of-the-art results on long-range graph benchmarks (LRGB) and competitive performance on standard GNN benchmarks, while also exhibiting reduced train–test gaps indicative of better generalization. The framework provides a flexible, extensible way to design attention mechanisms by changing the energy function and supports efficient training via proximal methods and a GraphGPS-based platform, making it practical for graphs with complex, long-range dependencies.

Abstract

Graph Transformers (GTs) have demonstrated superior performance compared to traditional message-passing graph neural networks in many studies, especially in processing graph data with long-range dependencies. However, GTs tend to suffer from weak inductive bias, overfitting and over-globalizing problems due to the dense attention. In this paper, we introduce SFi-attention, a novel attention mechanism designed to learn sparse pattern by minimizing an energy function based on network flows with l1-norm regularization, to relieve those issues caused by dense attention. Furthermore, SFi-Former is accordingly devised which can leverage the sparse attention pattern of SFi-attention to generate sparse network flows beyond adjacency matrix of graph data. Specifically, SFi-Former aggregates features selectively from other nodes through flexible adaptation of the sparse attention, leading to a more robust model. We validate our SFi-Former on various graph datasets, especially those graph data exhibiting long-range dependencies. Experimental results show that our SFi-Former obtains competitive performance on GNN Benchmark datasets and SOTA performance on LongRange Graph Benchmark (LRGB) datasets. Additionally, our model gives rise to smaller generalization gaps, which indicates that it is less prone to over-fitting. Click here for codes.

Figures (4)

• Figure 1: Overview of the SFi-Former architecture. Our design enables node features to be aggregated from the features of their adjacent nodes and selectively from distant nodes based on our Sparse-Flow-induced attention mechanism, achieving robust performance on downstream tasks.
• Figure 2: Illustration of the SFi-attention, where the energy-minimized flow $z_i^* = \frac{\operatorname{Soft}_{\lambda{f}_i}(\mu^*)}{r_i}$ serves as the attention score from a given query node to the key node ${\bm{v}}_i$. Here, frictions serve as learnable node-wise noise filters, allowing only strong signals to pass through. The resistance $r_i$ represents a dissimilarity measure of the query vector and the key vector of node ${\bm{v}}_i$. The optimal flows $\{ z_i^* \}$ correspond to the attention pattern from the given query node to all key nodes.
• Figure 3: Differences between the training and testing metrics for the GraphGPS and SFi-Former models throughout the entire training process across three datasets. Models with smaller differences between these metrics indicate better generalization.
• Figure 4: Demonstration of SFi-attention and its iterative process. We utilize a logarithmic transformation on the original attention values, represented with a viridis colorbar, where yellow areas indicate values near 1 and blue areas signify values close to 0 but above the threshold of 1e-8. Values exceptionally close to (below the threshold) appear white in this representation.