Graph Transformers Dream of Electric Flow

Authors: Xiang Cheng, Lawrence Carin, Suvrit Sra
Primary area: learning on graphs and other geometries & topologies
Keywords: Transformer, Graph Neural Network

Abstract:
We show theoretically and empirically that the linear Transformer, when applied to graph data, can implement algorithms that solve canonical problems such as electric flow and eigenvector decomposition. The Transformer has access to information on the input graph only via the graph's incidence matrix. We present explicit weight configurations for implementing each algorithm, and we bound the constructed Transformers' errors by the errors of the underlying algorithms. Our theoretical findings are corroborated by experiments on synthetic data. Additionally, on a real-world molecular regression task, we observe that the linear Transformer is capable of learning a more effective positional encoding than the default one based on Laplacian eigenvectors. Our work is an initial step towards elucidating the inner-workings of the Transformer for graph data. Code is available at https://github.com/chengxiang/LinearGraphTransformer

Decision: accept (poster)
Reviewer scores: 5, 6, 6, 6

Reviews and discussion:
<thread>
creation date:
1730460257024

thread_type:
official_review

<thread_object>
creator: reviewer_pygd

<content>
summary:
The papers show theoretically and empirically that the linear transformer can approximate the hypothesis function space of several graph problem (electric flow and Laplacian eigenvector decomposition). The authors also conduct a real-world experiment on molecular regression, in which positional encoding outputs by transformers are adopted.

strengths:
* The organisation of the paper is clear, it's easy to understand the main contribution of the paper. And mathematically, it's generally well-written.
* The idea of solving the traditional graph problem with a Transformer is novel. 
* For the proposed lemmas, the paper provides the relevant experiments which are quite nice.

weaknesses:
* My main concern is that since the neural network is a universe approximator, it's not surprising that it (in this paper, even it concerns linear Transformer, which is more like in a polynomial function) can solve approximate the hypothesis function space of several graph problems. I wonder how the lemmas shown in the paper can provide us with more useful information.
* The paper lacks a conclusion section.
* The real-world experiment concerns only the Laplacian eigenvector decomposition, it would be nice to have another for the electric flow.
* See the questions.

questions:
* While the paper concerns two graph problems: electric flow and Laplacian eigenvector decomposition, why the title mentions only the first one?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730666060062

thread_type:
official_review

<thread_object>
creator: reviewer_bwga

<content>
summary:
The paper shows that Linear Transformers (and their variant) using the incidence matrix can implement canonical problems, such as electric flow and subspace iteration. While the contribution is mainly theoretical, the authors also empirically validate their findings, and further show that on existing molecular datasets the linear transformer is capable of learning better performing positional encodings than standard laplacian eigenvectors.

strengths:
The results are interesting. I particularly appreciated the structure of the paper, in which experimental results are shown after their corresponding theoretical claims instead of having them all the end, which highlights the connection between the presented theory and the experimental results.

weaknesses:
W1. Generally speaking, I think the impact of the paper tends to be a bit limited. This is especially true because the considered architecture (linear Transformers), and in particular its variant which includes an L2 normalization, is not widely used. The main implication I can see is to use the variant of the linear transformer in place of existing predefined positional encodings, and therefore as an additional component of a bigger architecture.


W2. I find a bit confusing that the authors claim that the input to the Transformer is only the graph incidence matrix, without any other positional encoding (Lines 11-13), while for instance for the subspace iteration $\Phi_0$ is required to be some arbitrary column-orthogonal matrix. I think -- depending on the graph -- $\Phi_0$ might need to be a positional encoding to ensure all columns are orthogonal (note that if $\Phi_0$ is a random matrix then it is a positional encoding [Srinivasan and Ribeiro, 2020]). I think the authors need to clarify this point and adjust accordingly in the entire paper.


Srinivasan and Ribeiro, 2020. On the Equivalence between Positional Node Embeddings and Structural Graph Representations.

questions:
Q1. Can you elaborate on the practical implications of your findings (see W1)?

Q2. I think the eigenvectors you converge to have sign/basis ambiguities derived from the choice of $\Phi_0$. How do you deal with these ambiguities? Do you perform sign flipping at every forward pass?

Q3. I think denoting by $d$ the number of edges (Line 128) is a bit confusing, because $d$ is usually the embedding dimension. Also, in Equation 1, since $j= 1, \ldots, d$ is the first index of $B$, I think $B \in \mathbb{R}^{d \times n }$, but you wrote $B \in \mathbb{R}^{n \times d }$ (Line 132). Please clarify.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730706154263

thread_type:
official_review

<thread_object>
creator: reviewer_faup

<content>
summary:
The authors introduce a simplified (linear) Transformer architecture (no softmax activation in self-attention layers) and use it to show (theoretically and empirically) that it can implement graph algorithms. In particular they present a series of lemmas providing bounds in approximating the solution to problems like flow and heat kernel computation, finding resistive embeddings and performing eigenvalue decomposition (assuming specific weight matrices). Their input encodings include the incidence matrix representation of the graph and then the approximate solution will land in a subset of columns in the output encodings' matrix,  after passing through a a number L of transformer layers (where L is decided by the target approximation error). A parameter efficient variant is also presented. The linear transformer is tested on a series of small graph instances of the problems (e.g. of the order of 10 nodes and 20 edges) and for learning a positional encoding (PE) for a molecular regression task.

strengths:
- The idea is novel and appealing: using a transformer to compute the solution to a graph problem as part of the latent vectors in its node output representations adds another level of applicability of transformers well beyond language understanding or generation.

- Lemmas are well organized, follow similar themes and the narrative is smooth and clear.

weaknesses:
- Removing the nonlinear softmax terms from standard transformer architecture, facilitates analysis but severely impacts the power of the model.

- Complexity of the approach is prohibitive: it can be O(n^4) and this explains their experimentation with very small synthetic graphs. Parameter efficient implementation is promising, but still the original idea is far from being scalable and thus practically testable beyond a couple of tens' of graph nodes.

questions:
- How would the analysis be impacted if nonlinearity in self-attention was re-introduced? In order to "understand, at a mechanistic level, how the Transformer processes graph-structured data" (lines 32-33), we are expected to remove only non-essential elements of Transformer architecture (and nonlinearity seems to be a critical one). The linearity in the Transformer, means that if its operations are expanded for a number of layers we'll reduce to a simple matrix operator.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730723639424

thread_type:
official_review

<thread_object>
creator: reviewer_ryva

<content>
summary:
The paper investigates the capabilities of the linear Transformer when applied to graph data, particularly its ability to implement algorithms for core graph tasks without explicit positional encodings. The authors demonstrate that the linear Transformer, which uses the graph incidence matrix as input, can solve canonical problems like electric flow computation and eigenvector decomposition. Key contributions include explicit configurations for weight matrices, error bounds for each task, and empirical results validating theoretical findings on synthetic data.

strengths:
**Originality**: This paper introduces a novel use of a linear Transformer, to perform core graph algorithms like electric flow and eigenvector decomposition without explicit positional encodings. 

**Quality**: The paper offers rigorous theoretical analysis with explicit weight constructions and error bounds for each algorithm.

**Clarity**: The paper is well-organized, clearly written, and enjoyable to read.

weaknesses:
1. Since the paper’s contribution is primarily theoretical, providing a proof sketch under the main results would be highly beneficial. Additionally, a more detailed description of the weight matrices would enhance clarity.

2. The theoretical results are not particularly surprising given the use of a linear Transformer. Could these results also apply to GNNs?

3. The practical impact of the proposed approach is unclear, as the empirical results are limited compared to numerous existing works. For instance, the performance on the ZINC dataset is significantly lower than the state-of-the-art, with an error approximately twice as large.

questions:
Please address the points raised in the previous section.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732293211420

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Summary of  Reviews and Our Response

comment:
We once again thank all the reviewers for their comments and feedback, and for the numerous insightful discussions on various aspects of our paper. Below, we summarize the reviews and our responses

We begin with a summary of the strengths of our paper

## **Strengths**

1. Using a Transformer to solve a graph problem is novel, and adds a level of applicability of Transformers beyond LLMs (reviewers ryVa, fauP and PYgD)
2. Theoretical analysis is rigorous, with explicit weight constructions and error bounds for each algorithm (reviewer ryVa). 
*[The authors wish to emphasize that the constructions are **parameter-efficient**, and the error bounds **contract exponentially with layers**]*
3. Synthetic experiments validate our theory. (reviewers ryVa, BwGA, PYgD)
4. Molecular regression experiment shows that the linear Transformer is capable of learning better positional encodings than standard laplacian eigenvectors on a real-world problem. (reviewers BwGA, PYgD)

## **Weaknesses and Our Response**
> Removing the softmax nonlinearity
> - severely impacts the power of the model (reviewer fauP)
> - limits impact of paper as linear Transformer is not widely used, L2 normalization is also not widely used (reviewer BwGA).
   
   **Our Response**:
   - Linear attention $Attn^{linear}(Z)$ is a cubic polynomial in $Z$ (**thus not linear in $Z$**). Expressive power of linear Transformer does increase with layers, as seen in all our experiments.
   - There is evidence that **linear Transformers are better suited to the problems considered in this paper**. (For this application, the expressive power of linear Transformer is **not lower** than softmax Transformer.)
   (i) Figures 1 and 2 of [additional plots](https://anonymous.4open.science/r/iclr_2025_rebuttal_plots-3C6/iclr_rebuttal_plots.pdf) show that **linear Transformer outperforms softmax Transformer** in this application.
   (ii) Theory and experiments of [Cheng et al 2023](https://arxiv.org/abs/2312.06528) suggest that softmax activation performs suboptimally compared to linear activation on tasks related to solving linear systems.
   - Appendix A.9 of [von Oswald et al 2023](https://arxiv.org/pdf/2212.07677) suggest that **our constructions for the linear Transformer may be approximated by a multi-head softmax Transformers.**
   - L2 normalization in our paper is similar to **BatchNorm, which outperforms LayerNorm in many graph Transformer experiments.** (see experiments in [(Dwivedi and Bresson 2020)](https://arxiv.org/pdf/2012.09699)).
   
> Transformers are universal function approximators
> - Theoretical results are not particularly surprising given the use of a linear Transformer (reviewer ryVa)
> - Since the neural network is a universal approximator, it is not surprising that it can solve approximately the hypothesis function space of several graph problems (reviewer PYgD)
   
   **Our Response**
   - Universal approximation results generally require a very large number of parameters, exponential in dimension or $1/\epsilon$ (see Corollary 4.1, [(Chong 2020)](https://openreview.net/pdf?id=rkevSgrtPr))
   - In contrast, for graphs of $n$ nodes and $m$ edges, our standard construction uses only $O(n^2+m^2)$ parameters per layer. The efficient implementation further reduces this to O(#query^2). Number of layers needed is only $\log(1/\epsilon)$, and sometimes as few as $\log\log(1/\epsilon)$.



> Complexity of the approach is prohibitive, can be $O(n^4)$.

**Our Response**
   - Theoretically, all constructions in our paper can be implemented using the efficient implementation (Section 5), which **requires only O(#query^2) parameters per-layer, independent of the size of the graph**.
   - Experimentally, we show in Figure 3 of [additional plots](https://anonymous.4open.science/r/iclr_2025_rebuttal_plots-3C6/iclr_rebuttal_plots.pdf) that **the efficient implementation performs as-well, and often better-than the standard implementation** on all the problems considered.
   - Table 2 of our paper demonstrates the good performance of the efficient implementation on a real-world molecular ZINC molecular regression problem; **our implementation only used 488 parameters**.

(continued below)
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper explores the application of linear Transformers to graph-structured data, demonstrating both theoretical and empirical capabilities in solving canonical graph problems such as electric flow and eigenvector decomposition. The authors provide explicit weight configurations, error bounds, and validate their theoretical claims through synthetic and real-world experiments, particularly focusing on learning positional encodings for molecular regression tasks. The strengths include originality and novelty, theoretical rigor, as well as synthetic experiments consistent with the theory. Suggestion for improvement include conducting more experiments against other PEs and on larger and more diverse datasets, and finding more graph applications for linear Transformers to broaden the impact.