Primphormer: Leveraging Primal Representation for Graph Transformers

Thank you for your appreciation of the novelty of our primal representation for graph transformers. We also appreciate your recognition of the extensive experiments conducted in our work. We address your concerns as below:

R4.1 Clarity.

Thank you again for pointing out this issue. To address this, we will include clear definitions of the used variables in the Notations part at the beginning of Section 2.

R4.2 Discussion about long-range dependencies.

We agree that long-range interactions (LRI) may have minimal impact in certain scenarios. Nevertheless, there is a growing interest in tasks that necessitate LRI for optimal performance, where long-hop attention should not be neglected. This interest has led the development of the Long Range Graph Benchmark dataset [1], on which we also conduct experiments in the manuscript.

Regarding local subgraphs or PPR sampled neighborhood schemes, we can adjust the formulation of the data-dependent projection $f\_X$ to accommodate them without disrupting the primal-dual relationship. For instance, we could enhance $f\_X$ to adapt to each node $i$, allowing $f\_{X,i}$ to aggregate local information. We believe this is a promising and interesting avenue for future research.

R4.3 Universal approximation property.

We sincerely appreciate your comment. The universal approximation property is a fundamental theoretical concept in deep learning theory. It is widely recognized that models that exhibit this property potentially possess strong generalization capabilities to unseen data or tasks. As we developed our Primphormer in the primal space, we were eager to determine whether our Primphormer still retains this advantageous characteristic. To explore this question, we introduced Theorems 3 and 4, offering theoretical assurances regarding the approximation abilities of our approach. This discussion will be included in the final revision. While directly demonstrating this property through experiments poses challenges, we hope the good performance of Primphormer presented in our manuscript's experiments could provide indirect evidence of this property.

R4.4 Residual connections.

In our manuscript, we maintain the same model architecture and residual connection scheme from [2, 3], substituting solely the attention module with our primal representation to ensure a fair comparison. We hope this clarification could address your concern.

[1] Dwivedi V P, Rampášek L, Galkin M, et al. Long range graph benchmark[C]. NeurIPS, 2022.

[2] Rampášek L, Galkin M, Dwivedi V P, et al. Recipe for a general, powerful, scalable graph transformer[C], NeurIPS, 2022.

[3] Shirzad H, Velingker A, Venkatachalam B, et al. Exphormer: Sparse transformers for graphs[C], ICML, 2023.

Thank you for your careful reading and insightful suggestions. We address your concerns as below:

R1.1 & R1.2 Novelty.

The key to accelerating the attention from the primal-dual perspective is to find good approximation for $o$: $o(\boldsymbol{x})=\sum\nolimits\_{i} v(\boldsymbol{x}\_i)\kappa(\boldsymbol{x}, \boldsymbol{x}\_i)=\sum\nolimits\_{i} v(\boldsymbol{x}\_i)\langle\phi\_q(\boldsymbol{x}), \phi\_k(\boldsymbol{x}\_i)\rangle$. Here, $v(x_i) \in \mathbb{R}^{d_o}$ provides the basis and $\kappa(x,x_i) \in \mathbb{R}$ are the weight in the corresponding basis. The approximation error depends on both basis and weights, but the effect of basis is more significant.

To achieve better approximation, there are two ways to introduce data:

• Data Adaptive Weight this is the way of [1]: the weight is set to be $\langle f\_X\phi\_q(\boldsymbol{x}), f\_X\phi\_k(\boldsymbol{y})\rangle$

• Data Adaptive Basis this is the way proposed in our paper: the basis is set to be $F\_X\boldsymbol{h}\_{e}$, $F\_X\boldsymbol{h}\_{r}$, where $\boldsymbol{h}\_{e}, \boldsymbol{h}\_{r}$ are the dual variables.

We hope this comparision could highlight the fundemental difference between the two papers. The advantage of making the basis data-adaptive is evident from the universal approximation capability (Theorem 3 and Theorem 4). Only when sufficient flexibility is introduced, we could prove such theorems. Additional evidence comes from the numerical experiments, see Table 3 for Primphormer (ours) vs Prim-atten ([1]).

The similarity in formulation comes from the fact that both [1] and our method are built on the same least square framework [2], where the dual variables are from equation constraints, i.e., Eq. (6) in [1] and Eq. (2.4) in our paper.

Lastly, the design of data-dependent projection $f\_X$ is also important. In this paper, we use a different, data-adaptive scheme:

The $f_X$ that we used can also be found in [3, 4] and is known as permutation-equivariant projection in deepsets. As we discussed in the manuscript, graph structures are determined by edges and the arrangement or ordering of nodes is not explicitly specified. Therefore, this formulation is more suitable for our tasks.

R1.3 Definition Citations.

In fact, we have cited several papers that use this definition in the line above Definition 1. But we do agree with you that using an inline citation may be better. We will certainly modify it in the final version.

R1.4 Theorems between two works.

The theorems of the zero-valued objective are quite important for our method and the one in [1]. This property is essential for making the alternative optimization approach applicable to solving the primal problem.

As explained before, both of them are in the same least square framework so that the proofs of Theorem 2 and Lemma 4.2 in [1] are quite similar (actually, they are both following Corollary 1 in [2]). For our new method, we still need to prove it, although there is little technical contribution in the proof and we did not list this as a contribution.

Actually, the universal approximation (Theorem 3 and Theorem 4) is the main property we would like to highlight. Our Primphormer's ability to approximate any continuous function on a compact domain (see Appendix C.4 and C.5) is not found in the architecture of [1]. This is a significant advantage of our new representation, which introduces a data-adaptive basis, as explained in R1.1 and R1.2.

[1] Chen Y, Tao Q, Tonin F, et al. Primal-attention: Self-attention through asymmetric kernel svd in primal representation[C]. NeurIPS, 2024.

[2] Suykens J. A. K. SVD revisited: A new variational principle, compatible feature maps and nonlinear extensions[J]. ACHA, 2016, 40(3): 600-609.

[3] Zaheer M, Kottur S, Ravanbakhsh S, et al. Deep sets[C]. NeurIPS, 2017.

[4] Cai C, Hy T S, Yu R, et al. On the connection between mpnn and graph transformer[C]. ICML, 2023.

Thank you for your recognition of comprehensiveness and detailed feedback on our work. We address your concerns as below:

R3.1 Equation clarification.

We sincerely appreciate your valuable feedback. It is indeed a meaningful suggestion for readers.

• In Eq. (2.2), $\alpha\_j$ is a scalar and $\boldsymbol{\omega}$ is a vector. This equation comes from the original representer theorem, which gives an element-wise explanation of attention output. As you said, for a multi-dimensional output $\tilde{\boldsymbol{g}}$, we indeed need the vector form, naturally generalized from the element-wise definition: $\boldsymbol{\alpha}\_j\in\mathbb{R}^s$ are vectors, where $j\in[N]$, $N$ is the number of nodes, and the feature mapping $\phi(\boldsymbol{x}):\mathbb{R}^d\rightarrow\mathbb{R}^p$,

$$\tilde{\boldsymbol{g}}=\sum\nolimits\_{j} \boldsymbol{\alpha}\_j\langle\phi(\boldsymbol{x}\_i),\phi(\boldsymbol{x}\_j)\rangle =\sum\nolimits\_{j} {\rm vec}(\boldsymbol{\alpha}\_j\phi(\boldsymbol{x}\_i)^\top\phi(\boldsymbol{x}\_j)),$$ $$\overset{(a)}{=} \sum\nolimits\_{j} \left(\phi(\boldsymbol{x}\_j)^\top\otimes\boldsymbol{\alpha}\_j\right)\phi(\boldsymbol{x}\_i)=\left\langle \sum\nolimits\_{j} \phi(\boldsymbol{x}\_j)\otimes\boldsymbol{\alpha}\_j^\top,\phi(\boldsymbol{x}\_i) \right\rangle:=\langle\boldsymbol{W}, \phi(\boldsymbol{x}\_i)\rangle,$$

where $\boldsymbol{W}\in\mathbb{R}^{p\times s}$ and $(a)$ comes from the vectorization (${\rm vec}$) property of the Kronecker product $\otimes$. We hope these equations could address your concerns.

• Thank you for pointing out this issue. We recognize that the lack of clear definitions for the variables has caused confusion. To address this, we will include clear definitions of the used variables in the Notations part at the beginning of Section 2. Additionally, we will introduce the clear connection to the Transformer architecture in a proper position.

R3.2 Theorems.

We would like to first answer the question: the theories in our method has not been already established and we would like to clarify it as follows,

• Theorem 1. Our method uses a new primal representation for attention mechanisms, thereby introducing a new primal optimization problem. Theorem 1 establishes the corresponding primal-dual relationship for this new primal representation.

• Theorem 2. Given that we establish the primal-dual relationship for our primal representation within the least-squares framework, we are intrigued to ascertain if our method retains the property of a zero-valued objective. Theorem 2 validates this, thereby enabling the application of an alternative optimization approach to address our primal problem.

• Theorem 3. In this theorem, we establish that our Primphormer serves as a universal approximator for any permutation-equivariant function. We establish this theorem by firstly demonstrating its representational capacity to Sumformer [1], and subsequently utilizing Sumformer as an intermediary to control the overall error through the application of the triangle inequality (refer to Appendix C.4). A notable distinction between our method and Exphormer[2] pertains to the primal and dual spaces. While Exphormer demonstrated its property in the dual space (sparse attention), our method operates in the primal space. Moreover, our focus centers on exploring worst-case scenarios, employing the supremum norm, in distinction to Exphormer, which utilized the $L_p$ norm to measure the accuracy. These disparities lead to divergent proof strategies.

• Theorem 4. Next, we introduced our theorem for any continuous function with positional encodings on compact supports. In contrast, Exphormer [2] necessitated both expander edges and virtual nodes in its sparse attention model. To ensure sufficient node interactions and the universal approximation theorem, [2] required approximately $\mathcal{O}(\log N)$ attention layers, where $N$ denotes the number of nodes. Through our proof, we have streamlined this to a single attention layer by leveraging virtual nodes in our primal representation.

[1] Alberti S, Dern N, Thesing L, et al. Sumformer: Universal approximation for efficient transformers[C]. Topological, Algebraic and Geometric Learning workshops, 2023.

[2] Shirzad H, Velingker A, Venkatachalam B, et al. Exphormer: Sparse transformers for graphs[C]. ICML, 2023.

Thank you very much for finding our idea interesting and providing insightful comments. We also appreciate your recognition of the practicality of our work. We address your concerns as below:

R2.1 Information flow via virtual nodes.

Virtual nodes (VNs) create shortcuts for sharing information between graph nodes, facilitating global information exchange and enhancing information flow, as shown in some previous works [1,2,3]. In cases where important information needs to be preserved across distant nodes, VNs can improve information flow.

Your suggestion about complex hierarchical structures hits the idea of a very recent paper [4] that introduces an extension of VNs to enhance information exchange. We believe this is an interesting and important scenario. This technique could also be integrated into our method, which we plan to explore in future work.

R2.2 Applicability of our method.

The transition to the primal space does not require any preconditions. Regarding your question, the discussion about the asymmetric kernel actually covers the symmetric one. We hope this clarification could address your concern: Primphormer can handle both asymmetric and symmetric attension score matrices.

We will make the necessary modifications in the manuscript to address the issues you have pointed out in the manuscript.

[1] Hu W, Fey M, Zitnik M, et al. Open graph benchmark: Datasets for machine learning on graphs[C]. NeurIPS, 2020.

[2] Hwang E J, Thost V, Dasgupta S S, et al. An analysis of virtual nodes in graph neural networks for link prediction[C]. LoG, 2022.

[3] Cai C, Hy T S, Yu R, et al. On the connection between mpnn and graph transformer[C]. ICML, 2023.

[4] Vonessen C, Grötschla F, Wattenhofer R. Next level message-passing with hierarchical support graphs[J]. arXiv preprint arXiv:2406.15852, 2024.