Thanks for your valuable review. Here are our responses. Please also refer to the rebuttal PDF for some extra tables and figures.

Innovation

As mentioned in lines 97-102, we use Exphormer to interpolate between MPNNs and full-attention Transformers. Lower expander degrees make the model similar to MPNNs (particularly, eg, GAT), and the highest degree (n-1) will be equivalent to full Transformers. None of our theory relies on the expander degree or even the existence of the expander graph, and the experimental results include expander degrees from 30-200. The main idea of our work is to make a sparse and scalable method, and it made sense to start from an already sparse method such as Exphormer, and build on it to sparsify further. Sparsification from the full Transformer can be a little tricky, because the attention scores can be very small and unreliable (because of the large number of nodes to attend).

In addition, our paper has other contributions of additional interest, such as the batching technique (a particular fit for our sparsification, but usable in other settings), and our theoretical results regarding the compressibility of the network.

Comparison with GNNs

Thanks for pointing to these papers. The second paper mentioned here was published after our submission, and so we were not aware of it before submitting. However, there are a few important things to consider:

• The networks they have used for their experiments are much larger than the ones used in our work, in the sense of the number of parameters. For example, for Minesweeper they use 15 layers with hidden dimension 512, while we use four layers with hidden dimension 32. We also have focused on the memory efficiency, and most of the small datasets are trained with less than 4GB GPU limitation. Their models have been trained on a workstation with 8 RTX 3090 GPUs (each of which would typically have 24GB of memory).

• While our work is focused on memory efficiency, we also do batching based on the task and the requirements of the problem (relying on the attention scores from the initial network). However, these works need to train their method based on random subset selection batching, which as shown in our Figure 4 and Appendix E can sometimes behave very poorly. As the ratio of the batch size to the whole graph goes to zero for a fixed-degree graph, the probability of ever seeing any edge in the training process also goes to zero. By contrast, our batching is independent of the ratio of the graph size/batch size, and allows trading off time and memory without biasing the gradients of the problem parameters.

• These GNNs are well studied for many years, and Graph Transformers are relatively new. Thus, there is probably also potential to also find new tricks that improve GT performance. Many key advantages of the GTs can also be brought to the GNN world, as most of these things appear in the papers you have mentioned. The long-range dependencies can be handled by deeper GNNs, and oversmoothing/oversquashing and simplistic message functions can be handled by residual connections, jumping knowledge, layer normalization, and larger hidden dimension. Even so, much smaller GTs “naturally” handle many of these problems, and so the potential for better performance by finding appropriate tricks seems solid.

Ablation studies

Thanks for this suggestion; we have added an ablation study including this in our rebuttal pdf.

The theoretical analysis

Indeed, this is a limitation of our theory at time of submission. Since submission, we’ve continued to work on that, and have been able to extend the result to deeper layers for a scheme similar to (but not the same as) the narrow network we use in practice: a narrow attention dimension, but dimensions for the other parts of the network (e.g., the MLP layers) the same as the target network. (This “in-between” scheme saves in the most expensive part of the network.)

Because of the space limitation in this reply, please refer to our response to reviewer UFNp for the details.

This expanded version of the theorem helps extend our guarantees about the possibility of approximation to deeper layers. Thank you for pointing out this important issue.

Minor Issues

Thanks for mentioning this. We will fix the mistake and make the writing more clear for the next revision, as well as adding pseudocode to clarify.

Speed improvement

To give you some examples, the training time per epoch for the Tolokers dataset when training on the sparsified graph improves from 0.56s to 0.47s, and for the Photo dataset, it improves from 0.43s to 0.36s on average. This is a significant improvement, considering that many other parts of the network, such as Q, K, and V mappings and the feed-forward networks, are fixed and have similar overhead in both cases, with only the attention score estimation part changing.

The time improvement is even more considerable when comparing the Exphormer model to our sparsified regular graph. For example, the whole epoch time (neighbor sampling + training + validation) improves from 6.2s/epoch to 1.1s/epoch on the Tolokers dataset, and from 1.7s/epoch to 0.5s/epoch on the Photo dataset.

Thank you for your response. I have carefully reviewed all the content, particularly the comments from other reviewers. While I acknowledge the theoretical contributions of the paper, I remain unconvinced by the experimental results. When I closely examined Tables 1 and 2 in the original paper, I noticed that SGFormer [1] was not included. This inconsistency raises concerns about the validity of the experimental findings. The authors should conduct a thorough comparison with the latest scalable transformers [1,2], including performance, memory usage, and runtime across a broader range of datasets in the original paper (performance across all datasets). Given the current results, I will be maintaining my initial score.

[1] SGFormer: Simplifying and Empowering Transformers for Large-Graph Representations.

[2] Polynormer: Polynomial-Expressive Graph Transformer in Linear Time.

Thanks for your very thorough review; sorry for brevity (character limit). Please also see rebuttal pdf for updated results.

quadratic memory complexity

FlashAttention computes either full-attention or sparse-block attention mechanisms, with memory-aware loading of relevant parts of data. The use of full-attention, block-sparse attention (Big Bird), or linear approximate attention mechanisms (Performer) has been explored in the GraphGPS paper; Exphormer usually gets superior accuracies, probably thanks to its inductive bias.

A full-graph transformer with FlashAttention should work (with somewhat-worse accuracy expected), but would highly rely on positional encodings, standard versions of which are difficult to compute for large graphs. This is a feasible model to explore, but would require significant innovation and be very different from our work.

The block-sparse version of FlashAttention without considering the structural biases of the graph seems a poor fit for graph structures (cf BigBird results from GraphGPS).

Doing FlashAttention on an Exphormer-type attention pattern seems inefficient, since the block structure seems important to FlashAttention's efficacy, while Exphormer-style graphs can have nodes of wildly varying degrees (if the original graph’s degree varies, as usual).

Something like FlashAttention might be more feasible on regular graphs, as from Spexphormer. But this potential computational improvement to our approach seems to require deep thought about GPU memory structures and difficult implementation work, as did FlashAttention; we don't think this should be necessary for an initial paper exploring a new approach.

evaluation focuses solely on memory

Where we can train the original network, our approach is faster than Exphormer in general. For Photo, Exphormer takes 1.7s/epoch, our initial network 1.1s/epoch, final Spexphormer 0.5s/epoch. For Tolokers, Exphormer takes 6.2s/epoch, our initial network 4.6s/epoch, final Spexphormer network 1.1s/epoch. This advantage is larger for denser graphs. We also have added some experiments on the trade-off between RAM and GPU, as we can trade-off memory and time without any advanced GPU-level implementations (Figure 2 in rebuttal pdf).

potential bias in the evaluation

Even using a smaller degree Exphormer uses much more memory than our higher degree Spexphormer, see Fig 1 in our pdf for some numbers. Since Spexphormer can effectively use these higher degrees, it is more likely to be able to handle long-range dependencies. For example, in the ogbn-proteins dataset, an expander degree of 30 barely achieved an AUC of around 0.78 with a thorough hyperparameter search. Increasing the degree to 200 yielded an average AUC of 0.806.

Table 2 ... stronger baselines are needed

Thanks for highlighting this. We mostly followed the experimental setup + setup code from SGFormer; while they state they use the standard 80-10-10 train-val-test split, we had not previously noticed that their code (and ours) actually uses a 50-25-25 split. The numbers in Table 2 are correct for this split, but exacerbate the difference vs Polynormer due to smaller training set; eg SGFormer's GCN has accuracy 62%, while Polynormer's is 75%. We only just realized this, and are rerunning with the standard split; we'll try to provide numbers before the discussion period ends.

It's also worth noting Polynormer uses many more parameters than our models: 7 local + 2 global layers, each with hidden dimension 512, compared to our 2-layer network with hidden dimension 64. They train with batching based on a random subset of nodes, which as mentioned in Figure 4 and Appendix E can fail catastrophically; on Pokec in particular, though, the neighborhood summary is the most important feature, and the graph structure is not especially relevant, so this does not hurt much.

The idea of polynomial layers is also possible to directly drop in most common architectures, including ours; a SpexPolynormer combination model would be straightforward to write down. Their added complexity and hyperparameters, however, including a warm-up phase, mixture of global and local layers, etc. makes this cumbersome to include when exploring other new ideas.

Also, all of our current large-graph experiments now run with under 4GB of GPU memory, showcasing that our approach works well with very small batch sizes – unlike random-subset approaches which do not scale to small batches. (See the rebuttal pdf for more.)

Moreover, our first phase of training uses high CPU memory, but later training and evaluating the learned model on a given node can be done with low memory. Pollynormer, SGFormer, even GCN all batch for low memory during training, but evaluate with high CPU memory (loading the whole graph).

Baseline with random sparsification; baseline distributions appear overly simplistic

Also, uniform neighbor sampling is a common approach, used famously eg in GraphSAGE which often gets competitive results. In Table 1 of the pdf, we see using our network instead of random sparsification helps a lot in some datasets.

We’re not sure what you mean by “distributions based on the actual models’ behavior,” other than exactly the distributions that we use in Spexphormer. Do you have a specific baseline in mind to compare to?

Ablation studies

Good idea; also in new Table 1.

Appendix B

Thanks for noticing this; sorry for the mistake which we overlooked in submission. We will rewrite this section.

l162

Our claim is not that it is always impossible (though it is at initialization) for a model to decide to e.g. ignore expander edges, but that bias variables make it much easier. We’ll clarify.

l242

You are correct that we are making an intuitive argument for why we chose sampling over top-k selection; we'll clarify.

Other minor points

Thanks; we'll fix these.

Questions

Please see the separate author response above; thank you for your detailed review!

We again thank the reviewer for the detailed review. Here are a few new confirmations:

Sampling Importance: Using the maximum-ranked edges instead of sampling can lead to poor results, especially seen on heterogeneous datasets. Here are some numbers:

This confirms that sampling instead of getting the maximum is usually a better approach. The size of the effect can be very different among datasets, though.

Subgraph Growth: For the efficiency of batching, as we argued earlier:

When $c$ is smaller than the expander degree $k$, as usual in our work, the neighborhood expansions are strict subsets of those from Exphormer. We haven't yet done explicit experimentation on this but will try to do so during the discussion period, or else for the revised version.

It is easy to see our expansion should be much smaller than for example batching with considering the whole neighborhood. For example, for the Photo dataset if we start with 5 nodes, averaging over 10 times of sampling 5 initial nodes and extending the neighborhood our method increases the neighborhood size (number of nodes $\pm$ std) with the following numbers in layers:

0: 5.0 $\pm$ 0.0

1: 28.9 $\pm$ 1.1

2: 164.8 $\pm$ 7.4

3: 818.1 $\pm$ 37.1

4: 2587.7 $\pm$ 94.1

However, whole neighborhood sampling has these number of nodes:

0: 5.0 $\pm$ 0.0

1: 274.1 $\pm$ 45.1

2: 6172.3 $\pm$ 356.0

3: 7650.0 $\pm$ 0.0

4: 7650.0 $\pm$ 0.0

The graph has 7650 nodes in total; in 2 layers it already extends to the whole graph. This expansion can be smaller by reducing the expander degree, but the expander graph approximates full-attention Transformers, and thus if the expansion does not reach the whole graph, obviously the expander graph will fail in approximating the full-Transformer.

This is not the only advantage, and as we have previously posted in the response to Reviewer vT4j, this lets us have regularity in our model and caused over 16% speed-up in the calculations for the Tolokers and Photo calculations over the sparse but not regular calculation scheme. The memory saving over the full neighborhood even without batching (which means we are not using the advantage of the slower expansion of our method in this comparison) has been explored in Figure 1 of our rebuttal PDF.

Pokec Dataset and Polynormer Model Baseline: For the Pokec dataset, we can confirm that we have used SGFormer’s data split, which is different from the one used in the Polynormer paper. Our experiments are valid and we are comparing with numbers from SGFormer’s table, thus all the numbers are using the same split. Our method is not comparable with the GCN, NAGphormer, and Polynormer numbers mentioned by the reviewer due to different data splits, and if you compare our model with GCN with the SGFormer split our method has far superior results.

For comparing with Polynormer in general as is common in the machine learning community, usually, methods with a similar number of parameters should be compared. Our model for the Pokec dataset has only 83,781 parameters, and while Polynormer does not report the number of parameters explicitly they use a hidden dimension of 512, which with only one layer of a linear map would need 262,144 parameters. They have 9 layers in total for this dataset, each having multiple linear mappings. Thus the scale of the number of parameters is not comparable at all.

We have made our best effort to answer all the reviewer's concerns. We would greatly appreciate it if the reviewer could read our rebuttal and let us know if there are still concerns remaining. We would be happy to answer, though the remaining time from the discussion period is very short.

Thanks for your valuable feedback! Here you can find our response; we also encourage you to check our rebuttal pdf for some new experiments.

the two stage process

Even for one-stage training, hyperparameter tuning is a significant part of the training. In our setup, the first stage is relatively easier to tune — we only need one training run that converges from the initial network, and as the number of layers and the hidden dimensions are fixed, usually just tuning the learning rate works. So the “overhead” of training two networks is not so huge, compared to the usual process of hyperparameter tuning. We think this plus the additional extra complexity – mitigated by us releasing documented code for the two-stage pipeline – is not so extreme, and often worth it for the benefits it brings.

Also, the training process of the second network is independent of the initial network given the attention scores. We intend to share the attention scores we have from our initial network’s training publicly for possible future uses.

error propagation

It is true that the error can propagate. However, this is not too drastic because the second network learns its own attention scores. The neighborhood sampling is done with the initial network, and the second network just gets the updated list of the neighbors and learns its own attention scores. As long as the initial network gives reasonably high attention scores to important neighbors, the second network can adjust the attention scores. Even if some of the selected neighbors are not very informative, the second network can ignore them by giving a lower attention score. If an informative link is missed, it is likely that the needed information can propagate through other (longer) paths. Also, different neighborhood sampling per Transformer layer helps to mitigate the possibility of missing important links. Even if some of the important neighbors are underestimated in the initial network, increasing the number of sampled neighbors can also help bring them back to the sampled neighborhood.

Our experimental results in Table 1, compared to Exphormer, show that the total amount of introduced error is not that large, and our approach can give competitive results.

The experiment section shows reduced memory usage ...

In Figure 3, the memory has been reported from the actual running of the methods for smaller datasets. It is not feasible to train the original Exphormer on a 40GB GPU even with only 1 layer and hidden dimension of 4 or 8. Among the tested datasets, the largest feasible for Exphormer is ogbn-arxiv, where we see a ~5x reduced memory usage. The memory savings is considerable even without any batching. Training the model with two layers and hidden dimension 8 requires near 200GB for Amazon2M dataset; training it with the hidden dimension 256, and without advanced techniques that trade memory and time, requires between 32 (256/8) to (32)^2 times more memory, which is not feasible even with most CPU devices.

theory

Indeed, this is a limitation of our theory at time of submission. Since submission, we’ve continued to work on that, and have been able to extend the result to deeper layers for a scheme similar to (but not the same as) the narrow network we use in practice: a narrow attention dimension, but dimensions for the other parts of the network (e.g., the MLP layers) the same as the target network. (This “in-between” scheme saves in the most expensive part of the network.) The result is as follows:

Theorem: ​​Assume we have a Transformer network $\mathcal{T}$ with arbitrary large hidden dimension $d_l$, $L=O(1)$ layers, and in this network, in all layers, we have $\lVert h_\cdot\rVert_2 \leq \sqrt{\alpha}$, and $\lVert W_\cdot \rVert_{op} \leq \beta$. There exists a Transformer $\widehat{\mathcal{T}}$, that for any layer $\mathbf{W}_Q$ and $\mathbf{W}_K$ are in $\mathbb{R}^{d_s \times d_l}$ for a $d_s=\mathcal{O}(\frac{\log n}{\epsilon^2})$, with a sufficiently small $\epsilon$, and for all $i \in [n]$, $\lVert\mathcal{T}(X)_i - \widehat{\mathcal{T}}(X)_i\rVert_2 = \mathcal{O(\epsilon)}$.

The proof idea is to use Johnson-Lindenstrauss for the attention scores, bound them by relative error, and then bound the Euclidean distance of the other vectors based on the errors happening in the attention scores. The difference here with the previous proof is that the error can propagate from the Euclidean distance from the input of the layer too. But, we bound all these errors and show that the total error is still $\mathcal{O}(\epsilon)$.

This proof only works for narrow attention score estimation, and assumes the other parts of the network have the same dimensions. This is, however, the most memory/time-intensive part of a Transformer architecture. The rest of the sections are node/token-wise and linear with respect to the number of nodes. The attention score estimation part of a full-Transformer layer requires $\mathcal{O}(n^2d)$ operations and $\mathcal{O}(md)$ operators are required for a sparse Transformer with $m$ attention edges. Thus this change effectively changes the computational complexity of the model.

Also, another interesting thing about this proof is that it not only shows that there is a network with similar attention scores but narrow attention calculation hidden dimension; but also by bounding the error in the output, it shows that this network is nearly as good as the large one, and thus if the large network is optimal, this would be at least near-optimal network in the lower dimension.

This expanded version of the theorem helps extend our guarantees about the possibility of approximation to deeper layers. Thank you for pointing out this important issue.

Writing format

We very much agree with this comment and we will definitely move that figure for the next revision.