What Can We Learn from State Space Models for Machine Learning on Graphs?

We thank the reviewer for the insightful comments. Here is our response.

Weakness 1: Layer Complexity vs Expressivity: The paper states that the complexity of a GSSC layer is “… O(nmd) where n is the number of nodes and m, d are hidden and positional encoding dimension.” (239-240) This means that a GSSC layer has O(|V|) complexity. Consequently, GSSC is in general incapable of examining every edge in the graph, unlike MPNNs with O(|V|+|E|) complexity, such as GINE. Although the authors prove that GSSC is “more powerful than MPNNs” in terms of the WL hierarchy, the expressivity implications of not being able to examine every edge seems to be overlooked. Even preprocessing the graph to incorporate edge features into nodes, which itself requires O(|E|) time, would still necessitate m∈O(|E||V|) to store these features without information loss, thus exceeding O(|V|) complexity overall.

Note that GSSC is practically linear-time as we set d (number of eigenvectors) to be a constant in actual implementation. In theory, GSSC with constant d many eigenvectors can also distinguish non-isomorphic graphs that 1-WL cannot distinguish. It is true that in our proof, n many eigenvectors are required as a sufficient condition to losslessly simulate MPNNs, which means the complexity of eigendecomposition preprocessing is O(nE). But practically it is unnecessary, especially given that we also incorporate GSSC with MPNNs in our actual implementation.

Weakness 2: Preprocessing Complexity: The paper claims that finding the top d eigenpairs with Lanczos methods has O(nd2) complexity (286-287). However, since sparse matrix-vector multiplication with the Laplacian matrix is necessary for Lanczos methods, the complexity per iteration would be at least O(|E|), resulting in an overall complexity of at least O(d|E|) to find d eigenpairs. This exceeds the paper’s claim of O(nd2) preprocessing, and thus requires clarification.

Thank you for pointing this out. We made a typo and the complexity of Lanczos methods is indeed O(Ed), which is still considered efficient for sparse graphs. We clarify that in the revised manuscript.

Thank the authors for their response.

Sorry for the misunderstanding, I meant figure 3, not figure 4. The paper claims (L498-500):

…to obtain the top-dresults, which is both fast and GPU memory-efficient (linear in n), shown by the orange lines (LapPE GPU-topEigen) in the figures with $d=32$.

If it’s referring to figure 3 (# nodes vs GPU memory usage, preprocessing), it still does not look linear as claimed. The slope of the orange line increases as # nodes increases.

The normalized Laplacian $L$. It has the similar level of sparsity as $A$.

Theoretically, logbpcg can be implemented in a matrix-free manner, which means it does not need to store the adjacency matrix $A$ , but instead accesses $A$ by evaluating matrix-vector products . Therefore, the space complexity is simply storing the eigenvectors $x$ in memory, which has space complexity $O(n)$.

In that case, is $A$ meant to be $L$ in your response above? It does not change the analysis but I’m slightly confused by the inconsistency with the paper (the paper decomposes $L$ while $A$ is considered in your response).

Thank you for the follow-up questions.

the slope of the orange line increases as $n$ increases. Does this claim agree with figure 4?

Figure 4 shows the training time/inference time, which is not related to preprocessing (eigendecomposition). Note that this evaluation is on random graphs with $E=n^2\cdot0.0001$. It does not seem linear to $n$ because GSSC incorporates MPNN layers in practice, the latter of which grows linear w.r.t. number of edges. The GSSC block itself is linear w.r.t. $n$, unlike $O(n^2)$ for graph transformers.

does torch.lobpcg (L498 of your paper) work with a matrix not stored in memory?

In principle, we can implement matrix-free by definining the matrix to be a linear operator object (i.e., a object which can take a input vector and return a output vector). Packages like scipy.sparse.linalg.lobpcg can support this. Torch.lobpcg does not support this function for now (https://github.com/pytorch/pytorch/issues/89255). In our implementation we just stored the matrix into memory as the graphs are not large.

which matrix does GSSC require: the adjacency matrix $A$ or the Laplacian matrix $L$?

The normalized Laplacian $L$. It has the similar level of sparsity as $A$.

Thank the authors for their response.

I can better understand the complexities of relevant algorithms now, but this reflects major issues in the clarity of the paper. Here are some issues that came up in our conversation:

1. The paper incorrectly claims $O(nd^2)$ time complexity for preprocessing (286-287), because the complexity of Lanczos iteration was misrepresented (fixed in revision).
2. The paper claims $O(n)$ space complexity for preprocessing (L498-500), but the implementation has a space complexity of $O(|E|)$. This is not explicitly discussed in the paper.
3. To support the above claim, the paper presents Figure 3, and claims that memory consumption is linear to $n$. Meanwhile, the orange line of Figure 3 is clearly not linear, due to the actual implementation requiring $O(|E|)$ memory.

These inconsistencies would likely mislead a casual reader, and prompt a more serious reader to question the complexity claims of the paper. Therefore, I strongly recommend that the authors address these issues, and preferably, check all complexity claims in the paper for similar issues.

My overall rating still reflects my evaluation of the paper in its current form. However, I have increased my soundness score to acknowledge the authors’ efforts in clarifying their claims.

We thank the reviewer for the valuable comments. Here is our response.

Weakness 1: Experimental Design: My main concern is that the model evaluated in the experiments is a hybrid of GSSC and MPNN, without showing the performance of GSSC alone.

Thank you for the suggestion. Note that the general goal of SSM (and the proposed GSSC) is to replace vanilla attention, which does not have a good interface to incorporate edge (token-pairwise) features. This is an inherent limitation of SSMs. To address this limitation, SSMs (and their graph extension GSSC) need to be paired with modules that can handle token-pairwise features. Indeed, most of the graph transformer [1,2] and graph SSMs [3, 4] baselines are also built upon a hybrid with MPNNs to handle edge features.

Weakness 2: Baseline Selection: In Section 5.1, the authors compare only with MPNN, without providing baseline comparisons with SSM or GT, which seems insufficient.

Thank you for the constructive suggestion. Previously we have comparisons to graph Mamba and graph transformers in Section 5.2, Tables 1 and 3. Now we also add graph transformer (GraphGPS), sparse/linear graph transformer (Exformer, GraphGPS+performer) and graph SSMs (Graph-mamba-I) in Section 5.1, counting experiments. Please see Table A and Table B in the General Response. Note that High-order MPNNs like I2-GNN were already included in Section 5.1 in the original manuscript. Notably, we find that though these new baselines have relatively low training error, their test error is many times larger than GSSC. This is because these models do not preserve permutation equivariance, causing a poor generalization ability. In contrast, GSSC properly handles positional encodings in a permutation equivariant and stable manner, and thus has better generalization.

Weakness 3: Inconsistency in Model Architecture: In Section 5.1, the authors use selective GSSC, but it is not used in other benchmarks. If different architectural variants of the model are used in the experiments, I advise clearly indicating this in the tables.

Thank you for the suggestion. We now highlight in the Table caption if the selective mechanism is applied. For the counting task, we applied the selective mechanism because it is crucial for a higher expressive power. For real-world tasks, we found the selective mechanism does not bring significant gain and we guess expressive power is not a main bottleneck there. For model consistency, we will include the results of selective mechanisms on real-world benchmarks in the final version of the manuscript.

Question 1: Compared to other graph SSM works, this study indeed offers a novel perspective. In your opinion, what advantages does this new viewpoint provide over other methods? How do your experiments support these advantages?

Note that we have already discussed the difference between GSSC and other graph SSMs in Section 4, Line 350. Compared to other graph SSMs, the biggest difference is that we do not tokenize the graph and apply existing SSMs for sequences. Instead, we provide a way to design SSMs specific for graph data. As a result, our method inherently preserves permutation equivariance, an important inductive bias for graphs, while previous graph SSMs fail to do so. This advantage leads to a better generalization, supported by results of our experiments.

Question 2: The title is "What Can We Learn from State Space Models for Machine Learning on Graphs." Could you elaborate a bit more on other graph SSM models / SSM model on other domins, and compare them with GSSC?

We have a comprehensive comparison to other graph SSMs in Related Works, Line 350.

[1] Rampášek, L., Galkin, M., Dwivedi, V. P., Luu, A. T., Wolf, G., & Beaini, D. (2022). Recipe for a general, powerful, scalable graph transformer. Advances in Neural Information Processing Systems, 35, 14501-14515.

[2] Chen, D., O’Bray, L., & Borgwardt, K. (2022, June). Structure-aware transformer for graph representation learning. In International Conference on Machine Learning (pp. 3469-3489). PMLR.

[3] Wang, C., Tsepa, O., Ma, J., & Wang, B. (2024). Graph-mamba: Towards long-range graph sequence modeling with selective state spaces. arXiv preprint arXiv:2402.00789.\

[4] Behrouz, A., & Hashemi, F. (2024, August). Graph mamba: Towards learning on graphs with state space models. In Proceedings of the 30th ACM SIGKDD Conference on Knowledge Discovery and Data Mining (pp. 119-130).

We thank the reviewer for the valuable suggestions. Here is our response.

Weakness 1: While there is a careful analysis of the different design decisions/performance tradeoffs, I feel that there is only a limited understanding about what are the properties of the Architecture that lead to these decisions/performance differences.

The motivation to design the GSSC architecture is by generalizing the key components (absolute positions, factorizable kernels and global sum) in SSMs to graph data, as we argued in Section 3.1, Line 211 to 215. Such design leads to many theoretical advantages, as shown in Remark 3.1, that: (1) GSSC is permutation equivariant, while most linear graph transformers and graph SSMs are not; (2) GSSC is stable to graph perturbation, leading to a better generalization ability; (3) GSSC can capture long-range dependencies (Proposition 3.1). We believe these merits contribute to the performance gain.

Weakness 2: Scalability Concerns: The paper acknowledges the challenge of scalability for larger graphs. To address this, the authors could explore methods for optimizing the computational complexity of the GSSC architecture.

The main computation bottleneck is from preprocessing of eigendecompositions. Ways to alleviate the costs of eigendecompositions are discussed in Section 3.2, Line 282, such as using top K eigenvectors with advanced eigenvalue decomposition algorithms. Comparisons of actual preprocessing time was shown in Figure 3. Notably, our empirical evaluation shows that the practical preprocessing time is relatively small or negligible compared to total training time (Section 5.3, Line 509).

Weakness3: Weak experimental study: The paper lacks experimental evaluation of the Graph State Space Convolution (GSSC) method on heterophilic datasets. This omission is significant as such studies are crucial to assess GSSC's performance with heterophilic data and its robustness against issues like over-squashing and over-smoothing, which are common challenges in graph data analysis.

We respectfully disagree with the reviewer on the “Weak experimental study”. Though we agree that heterophilic data (as well as over-squashing/over-smoothing problems) is important, it is not widely considered/adopted by our relevant works/baselines [1,2,3,4]. On the other hand, our work studies generalizing SSMs to graph data, with a focus on achieving efficiency and long-range dependencies on graph data. To demonstrate these claimed benefits, we have extensive experiments covering substructure counting (to examine expressive power and verify Proposition 3.3), long-range benchmark (to verify Proposition 3.1), well-accepted graph benchmarks such as ZINC, OGB-MOLHIV, and runtime evaluation. How to deal with heterophilic data has unique challenges and is out of the scope of our main focus.

Question 1: For the scalable version of the architecture when you proposed using the first k eigenvectors, how does the method address the issue of missing eigenvectors in the middle and end of the spectrum?

Top-K eigenvectors is a well-adopted technique in positional-encodings-based methods. One justification is that many graph distances, e.g., diffusion distance, biharmonic distance, are dominated by low-frequency eigenvectors [5]. In other words, graph distances can be well approximated by top K eigenvalues and eigenvectors.

Question 2: What are the main challenges in extending GSSC to heterophilic graphs, and do you have any preliminary insights on how these challenges could be addressed?

We believe that heterophilic graphs are another challenging scenario with many unique problems to consider, which are out of scope of our work.

[1] Shirzad, H., Velingker, A., Venkatachalam, B., Sutherland, D. J., & Sinop, A. K. (2023, July). Exphormer: Sparse transformers for graphs. In International Conference on Machine Learning (pp. 31613-31632). PMLR.

[2] Wang, C., Tsepa, O., Ma, J., & Wang, B. (2024). Graph-mamba: Towards long-range graph sequence modeling with selective state spaces. arXiv preprint arXiv:2402.00789.

[3] Rampášek, L., Galkin, M., Dwivedi, V. P., Luu, A. T., Wolf, G., & Beaini, D. (2022). Recipe for a general, powerful, scalable graph transformer. Advances in Neural Information Processing Systems, 35, 14501-14515.

[4] Ma, L., Lin, C., Lim, D., Romero-Soriano, A., Dokania, P. K., Coates, M., ... & Lim, S. N. (2023, July). Graph inductive biases in transformers without message passing. In International Conference on Machine Learning (pp. 23321-23337). PMLR.

[5] Kreuzer, D., Beaini, D., Hamilton, W., Létourneau, V., & Tossou, P. (2021). Rethinking graph transformers with spectral attention. Advances in Neural Information Processing Systems, 34, 21618-21629.

We thank the reviewer for the insightful comments. Here is our response.

Weakness 1: The distinction between the proposed model (equation 4) and linear graph attention is minimal.

First, though conceptually related, GSSC is technically different from existing linear graph attention. Linear attention aims to approximate the softmax attention kernel by random features. Particularly, the linear attention cannot handle the symmetry of positional encodings properly as they are treated as regular node features, which breaks permutation equivariance. In contrast, GSSC uses flexible and learnable positional encodings as the features in a permutation equivariant and stable manner, and they are not tied to specific kernels.

Finally, to see the empirical difference between these architectures, we add extra comparisons to graph linear attention (e.g., GraphGPS+performer, Exphormer) on substructures counting and molecular graph benchmark. Please kindly see Tables A and B in our General Response. GraphGPS+performer has constantly worse performance than GraphGPS, as it is an approximation of the latter. In contrast, GSSC has a consistently superior performance over GraphGPS, GraphGPS+performer and Exphormer.

Weakness 2: The presentation requires improvement, and some claims need to be more precise.

We thank the reviewer for the detailed comment. For bullet point (a), we add discussions on the FFT algorithm in the revised manuscript. For bullet point (b), it seems the reviewer did not provide the mentioned reference [2] about cross correlation. We would like to include it if the reviewer can provide the reference. For bullet points (c-d), we find these are two typos and we now fix them. For bullet point (e), we use Deepsets as the phi function to model interaction between eigenvalues and we will elaborate it in the main text, Experiment Section.

Weakness 3.1 : Since the proposed model is compared with GSC and Linear Graph Transformers in the paper, it is essential to include their performance comparison in the experimental section (particularly in Table 2 for graph substructure counting, where the proposed model is expected to show superior expressiveness).

Thank you for the suggestion. We add extra comparisons to graph linear attention (e.g., GraphGPS+performer, Exphormer) on substructure counting and molecular graph benchmark. Please kindly see Tables A and B in our General Respons above.

Weakness 3.2: Comparisons with newer models like Spatial-Spectral GNN and Polynormer are also necessary.’

Thank you for introducing the models. Spatio-Spectral GNN [1] also has results on peptides-func and peptides-struct, which we will include in Table 1 in the next version of the manuscript. However, since it is not open-sourced, we would not be able to do further comparison given the limited time of rebuttal. Polynormer [3] has no overlapping datasets with us and we are working on adapting its code in our pipeline. If time allows we will update the results before the end of rebuttal.

Weakness 3.3: Additionally, comparisons against other SSM-based models (Graph-Mamba I and II), GSC, Linear Graph Transformers and recent models in Tables 2 and 3 are important to demonstrate the proposed model empirically.

Thank you for your suggestion. We add extra comparisons to graph linear attention (e.g., GraphGPS+performer, Exphormer) and graph SSMs (Graph-Mamba I) on substructure counting and molecular graph benchmark. Please see Table A and Table B in our General Response.

Question2: about typos.

Thank you for pointing them out. The parentheses in Equation 5 (line 201) is to emphasize how the efficient computation is achieved. We fix the other two typos in the revised manuscript.

[1] Geisler, S., Kosmala, A., Herbst, D., & Günnemann, S. (2024). Spatio-Spectral Graph Neural Networks. arXiv preprint arXiv:2405.19121.

[3] Deng, C., Yue, Z., & Zhang, Z. (2024). Polynormer: Polynomial-expressive graph transformer in linear time. arXiv preprint arXiv:2403.01232.