Paper: On the Stability of Expressive Positional Encodings for Graphs (ICLR-2024)

Reviewer comment (Theory):
The complexity of the proposed SPE is much larger than previous positional encoding methods because it needs to reconstruct graph structures, i.e., $\boldsymbol{V} \operatorname{diag}\left(\phi(\boldsymbol{\lambda})\right) \boldsymbol{V}^{\top}$, whose complexity is $\mathcal{O}(KN^{2})$. In contrast, the Transformer-based methods, e.g., BasisNet, only have the complexity of $\mathcal{O}(NK^{2})$, where $K \ll N$.

Author response:
Note that the complexity of BasisNet is also $N\times N\times K$, because it also needs to compute the matrix $V_kV^{\top}_k\in\mathbb{R}^{N\times N\times d_k}$ in $k$-th eigensubspace for $k=1, 2,...$. Also in general a transformer-based method should have complexity $O(N^2)$ because it requires to compute pair-wise correlation. 

Empirically, we also find SPE's running time is much faster than BasisNet. Kindly check our unified response ``*running time evaluation*'' above.

Reviewer comment (Experiments):
In the molecular property prediction task, SPE has more parameters than baselines. It would be better if the authors could align the number of parameters across different methods. Additionally, in the OOD tasks, the improvement of SPE over baselines is marginal.

Author response:
Thank you for the suggestion. We will align the number of parameters in the revised manuscript. 

For OOD task, the core idea is not to achieve state-of-the-art performance but to show how previous unstable methods suffer from the **risk of instability** on OOD data. In this sense, standard GNN model is actually a strong baseline because it is very stable.

Reviewer comment (Writing):
In equation (1), what does $\mathbb{R}^{n \times d} \times \mathbb{R}^d \rightarrow \mathbb{R}^{n \times p}$ mean? I understand that $\mathbb{R}^{n \times d} \rightarrow \mathbb{R}^{n \times p}$ represents the function applied to the position features of each node. What does $\mathbb{R}^d$ indicate? Operation on eigenvalues?

Author response:
Yes, it means a general positional encoding method can take both eigenvalues and eigenvectors as input.

Reviewer comment (Theory):
Is it possible to replace the element-wise MLPs of $\phi$ with polynomial functions? In this situation, I think the complexity can be significantly reduced and the expressiveness can be preserved since polynomials are also non-linear.

Author response:
We think MLPs are not be the bottleneck of computation complexity. This is because the input size is simply the number of eigenvectors to use, which is usually fixed and does not scale with graph size. On the other hand, MLPs are usually believed to be more expressive than polynomials. So overall MLPs are preferable for us.

Reviewer comment (Theory):
I feel that the paper lacks some explanations. It is not clear which modules of the proposed method contribute to it being stable. If no $\phi_\ell$ layers are added, wouldn't $K_\ell$ be equal to 1? In my understanding, this wouldn't hurt stability. Also, it seems to me that as $m$ increases the bound becomes looser and looser. If that's the case, why do we need multiple such permutation equivariant layers?

Author response:
> **W1**: I feel that the paper lacks some explanations. It is not clear which modules of the proposed method contribute to it being stable. If no $\phi_i$ layers are added, wouldn't  $K_i$ be equal to 1? In my understanding, this wouldn't hurt stability.

The stability stems from the operation $V\text{diag}(\phi(\lambda)) V^{\top}$. The key intuition is to combine multiple eigensubspaces with smooth composition, avoiding the splitting of individual eigensubspaces (which is what previous methods did).

If you simply let $\phi(\lambda)$ be $\lambda$, it is true that this model is also stable. However, its **expressive power** will be damaged. For instance, if we let $\phi(\lambda)=\lambda$ and use full eigenvectors, then $V\text{diag}(\phi(\lambda))V^{\top}$ simply becomes Laplacian $L$ (and equivalently adjacency matrix $A$). In this case SPE is just acting like a normal graph neural networks on $A$, bringing no aditional expressive power. Similarly, larger $m$ (number of $\phi$) may damage stability but could bring more expressivity. So in practice we let $m$ be a hyperparameter for the stability-expressivity trade-off.

Reviewer comment (Experiments):
One of the main weaknesses of this paper is the proposed model's complexity. Function $\rho$ takes a tensor of dimension $n \times n \times m$ as input. This might not be problematic in case the model is trained on molecules since molecules are small graphs. But in case of other types of graphs such as those extracted from social networks which are significantly larger, this can lead to memory issues.

Author response:
> **W2**: One of the main weaknesses of this paper is the proposed model's complexity.

We agree with the reviewer that complexity of current implementation of SPE could be a problem for large graphs. This problem also appears in BasisNet, whose complexity is also $O(N^2)$. Empirically, our unified response ``*running time evaluation*'' above shows the running time of SPE and other methods. We can see that SPE's running time is comparable to SignNet and is much faster than BasisNet. 

One alternative way for reducing complexity could be to sparse the $V\text{diag}(\phi) V^{\top}$ by only computing and storing its $(i,j)$-entry where $(i,j)$ is an edge. Then these positional encodings are treated as edge features for downstream tasks. The complexity will reduce from $O(n^2)$ to $O(|E|)$ where $|E|$ is the number of edges. In practice, the sparsing of $V\text{diag}(\phi) V^{\top}$ can be efficiently done by firstly assigning $V_{i, :}$ to $i$-th node features and using an one-layer message passing neural network to update and store edge features $[V\text{diag}(\phi) V^{\top}]\_{i,j}=\sum_{k}[\phi(\lambda)]\_k V\_{i,k}V\_{j,k}$.

Reviewer comment (Evaluation):
The proposed approach is much more complex than a standard GNN model, but in most cases it provides minor improvements over a model that does not use positional encodings. For instance, the improvement on Alchemy is minor, and also on DrugOOD, SPE provides minor improvements in the Assay and Scaffold domains and no improvements in the Size domain.

Author response:
> **W3**: `The proposed approach is much more complex that a standard GNN model, but in most cases it provides minor improvements over a model that does not use positional encodings.'

It is true that on Alchemy all PE methods provide minor improvements. But still SPE is the best. We conjecture that Laplacian encodings may not be a very good inductive bias on this dataset. For DrugOOD, the core idea is not to achieve state-of-the-artperformance but to show how previous unstable methods suffer from the **risk of instability** on OOD data. Standard GNN model is actually a strong baseline because it is very stable.

Reviewer comment (Evaluation):
No running times of the different models are reported in the paper.

Author response:
> **W4**: No running times of the different models are reported in the paper.

Thank you for raising this point. We add running time of all methods on ZINC and DrugOOD in unified response ``*running time evaluation*'' above. Overall we can see that for our current implementation SPE has comparable running time to SignNet and much faster than BasisNet.

Reviewer comment (Experiments):
The proposed model seems to advance the state of the art in the field of positional encodings for graphs, however, it is not clear whether it also advances the state of the art in the graph learning community. I would suggest the authors compare the proposed approach against some recently proposed GNN models, and not only against methods that produce positional encodings.

Author response:
> **W5**: The proposed model seems to advance the state of the art in the field of positional encodings for graphs, however, it is not clear whether it also advances the state of the art in the graph learning community. I would suggest the authors compare the proposed approach against some recently proposed GNN models, and not only against methods that produce positional encodings.

Thank you for the valuable sugeestion. We plan to add a few recently proposed GNN models for comparison during the rebuttal period. Again, graph positional encodings are important for graph learning community since (1) it plays a fundamental role for constructing powerful graph transformers [1, 2]; (2) it can be plugged into any GNN models as feature augmentation [3] and improve their expressive power.


[1] Vijay Prakash Dwivedi and Xavier Bresson. A generalization of transformer networks to graphs. AAAI Workshop on Deep Learning on Graphs: Methods and Applications, 2021.

[2] Devin Kreuzer, Dominique Beaini, Will Hamilton, Vincent Lé etourneau, and Prudencio Tossou. Rethinking graph transformers with spectral attention. Advances in Neural Information Processing Systems, 34, 2021.

[3] Vijay Prakash Dwivedi, Anh Tuan Luu, Thomas Laurent, Yoshua Bengio, and Xavier Bresson. Graph neural networks with learnable structural and positional representations. In International Conference on Learning Representations, 2022.

Reviewer comment (Theory):
In Figure 2, how did you compute the Lipschitz constant of MLPs? We can compute the Lipschitz constant for models that consist of a single layer, but exact computation of the Lipschitz constant of MLPs is NP-hard [1].\
[1] Virmaux, Aladin; Scaman, Kevin. Lipschitz regularity of deep neural networks: analysis and efficient estimation. Advances in Neural Information Processing Systems, 2018.

Author response:
> **Q1**: In Figure 2, how did you compute the Lipschitz constant of MLPs? We can compute the Lipschitz constant for models that consist of a single layer, but exact computation of the Lipschitz constant of MLPs is NP-hard

Thank you for pointing this out. Technically, the paper mentioned refers to the best (or minimal) Lipschitz constant. Here by Lipschitz constant we actually mean an upper bound for best Lipschitz constant. Note that we can upper-bound the best Lipschitz constant of MLPs by multiplying the operator norm of weight matrices. We clarify this in the Appendix Section B.4 in the revised manuscript.

Reviewer comment (Theory):
The instability of prior method (i.e., the so-called hard partition method) is not proved. The authors point out under Eq.2 that **hard partition** is induced when $\[\phi_{\ell}(\boldsymbol{\lambda})\]_j=\mathbb{1}$ (other places in $\phi(\cdot)$ are zeros), and then $\boldsymbol{V}\text{diag}(\phi_{\ell}(\boldsymbol{\lambda}))\boldsymbol{V}^{T}$ is the $\ell$-th subspace. The problem is that, if we set $\{ \phi_i \}_{i=1}^{m}$ that induce the hard partitions, then they are **constant functions** and meet the $K_{\ell}$-Lipschitz continuous assumption in Assumption 3.1, which is then used to prove the stability of SPE. Therefore, the question is, **is the prior method (i.e., the counterpart that uses hard partitions) really unstable?** It seems that hard-partitioned SPE can be proved to be stable via exactly the same proof of Theorem 3.1.

Author response:
Thank you for pointing this out. Kindly check our unified response ``*Why previous methods are unstable?*'' above for detailed proof. 

Regarding your argument that hard paritition is stable by Theorem 3.1, it is not valid because $\phi$ is actually **discontinuous**. To see this, let us consider basis-invariant methods using hard partition of eigensubspaces (e.g., BasisNet). In this case, the $k$-th entry of the hard partition $\phi_{\ell}(\lambda)$ is $$[\phi\_{\ell}(\lambda)]\_k=1,  \text{if $\lambda_k$ is $\ell$-th smallest eigenvalue},\quad [\phi\_{\ell}(\lambda)]\_k =0, \text{otherwise}$$ 
This $\phi$ function  is more like a **step function** instead of a constant function. It is **discontinuous**, and thus is not Lipschitz continuous. So Assumption 3.1 does not hold and Theorem 3.1 cannot be applied for such hard partition functions. You may kindly check Appendix C in the revised manuscript for details and examples.

Reviewer comment (Theory):
Equivalence for $\{\phi_i\}_{i=1}^{m}$. The authors restrict $\{\phi_i\}_{i=1}^{m}$ to be permutation equivariant, whose input is the Laplacian spectrum. Here, the authors are asking for equivalence under the reordering of eigenmaps/eigenvalues, instead of the reordering of graph nodes.

Author response:
The overall model is indeed permutation **equivariant to node reordering**. This is because the permutation equivariance of $\phi$ (w.r.t eigenvalues) gurantees the permutation equivariance of the $V\phi(\lambda)V^{\top}$ to node indices reordering. To formally see this, suppose we have eigendecomposition for $L=V\text{diag}(\lambda)V^{\top}$. Applying permutation matrix $P$ to $L$ may produce an new decomposition in form of $PLP^{\top}=PVQ\text{diag}(\lambda)Q^{\top}V^{\top}P^{\top}$ with certain $Q$, where $Q\in O(\lambda)$ is an block-diagonal matrix where each block is a rotation matrix (i.e., $O(\lambda)=\\{\oplus_iQ_i: Q_i\in O(d_i), d_i\text{ is the dim of $i$-th eigensubspace}\\}$). That means the new eigenvectors after permutation become $PVQ$. If $\phi$ is permutation equivaraint, then it means $Q\text{diag}(\phi(\lambda))Q^{\top}=QQ^{\top}\text{diag}(\phi(\lambda))=\text{diag}(\phi(\lambda))$. Here the first equality holds since permutation equivariant $\phi$ produces same values for entries of the same eigensubspace and thus we can switch the order of multiplication with block-diagonal $Q$.  Therefore the new positional encoding $PVQ\text{diag}(\phi(\lambda))Q^{\top}V^{\top}P^{\top}=PV\text{diag}(\phi(\lambda))V^{\top}P^{\top}$, that is, the permutation version of the old positional encoding.

Reviewer comment (Experiments):
On the universal expressiveness. The proof of this SPE's universality relies on being reduced to BasisNet. Therefore, two problems arise: (1) The experiment regarding expressiveness, i.e., the graph substructure counting, does not include BasisNet. (2) According to Lim et al. (2023), the instance of BasisNet, Unconstrained-BasisNet, universally approximates any continuous basis invariant function. In Unconstrained-BasisNet, IGN-2 (Maron et al., 2018) is the core part to achieve such expressiveness. However, in implementation, the authors set $\rho$ to be one identical GIN (Xu et al., 2019), which would surely limit the expressiveness.

Author response:
1. Please kindly check the new counting results of BasisNet in Figure 3 in the refined manuscript. 
2. We agree that $\rho$ has to be sufficiently powerful to achieve universality. But in practice we find GIN can already get promising results. This may be because the proposed architecture $V\text{diag}(\phi(\lambda))V^{\top}$ is already expressive and thus a complicated $\rho$ may not be that necessary for the current tasks.

Reviewer comment (Theory):
Would the learned $\{\phi_i\}_{i=1}^{m}$ be close to each other? This would lead to similar position encodings.

Author response:
We don't observe such phenomenon in practice, but if it is the case it just means one channel for such $\phi$ is sufficient for positional encoding on this specific task.

Reviewer comment (Novelty):
The novelty of the method itself is partially limited as the idea to use a weighted correlation over the eigenvectors closely resembles the correlation used in BasisNet.

Author response:
Note that the key difference between our method and BasisNet is not just adding a weight: BasisNet computes the correlation $V_kV_k^{\top}$ between eigenvectors $V_k$ **within** the **$k$-th eigensubspace** for different $k$; in contrast, SPE computes the correlation $V\phi V^{\top}$ between **all the eigenvectors $V$ across different eigensubspaces**. This allows us to combine multiple eigensubspaces with smooth composition. The main argument is that this operation equipped with a continuous and equivariant $\phi$ is the key to achieve **stability**.

Reviewer comment (Experiments):
The experiments are limited. The performance of SPE in Table 1 is sub-par, and details of the experimental results in Figure 2 are unclear and the hyperparameters seem not to be reported, which makes it hard to reproduce the experiments.

Author response:
The SPE's result on ZINC is 0.0693, which greatly improves the baseline such as SignNet (0.0853) and BasisNet (1.5555). In fact, this result also outperforms many other recently proposed positional encodings methods [1, 2, 3] and even highly expressive GNN models [4, 5, 6]. 

For Alchemy dataset, SPE is still the best, though all positional encodings bring a little improvement. We conjecture that Laplacian encodings may not be a good inductive bias for this dataset.

Also thanks for pointing the hyperparameter things out. Some of the hyperameters were reported in the Appendix, Section B and we add more details of them in the revised manuscript.

Reviewer comment (Theory):
The point regarding the trade-off between expressivity and generalisation is unclear. Is there a formal explanation which we can quantify?

Author response:
A classic interpretation is the **variance bias trade-off**. Let $\hat{\epsilon}_n$ be the training loss, $\epsilon$ be the test loss, $\mathcal{H}$ be the hypothesis class and $\hat{h}_n=\arg\min\_{h\in \mathcal{H}}\hat{\epsilon}\_n(h)$ be the trained model. If we increase the size of the hypothesis class (i.e., make the model more complicated and expressive, which corresponds to enlarge the Lipschitz constants of $\phi$ in Figure 2), intuitively
- training loss $\hat{\epsilon}_n(\hat{h}_n)$ is going to decrease, because we are optimizing over a larger hypothesis class. Empirically this corresponds to the overall decreasing trend of training loss as shown in the left-most plots in Figure 2.
- generalization gap $\epsilon(\hat{h}_n)-\hat{\epsilon}_n(\hat{h}_n)$ is going to increase, because from Proposition 3.1 this term is upper bounded by the complexity of the model and the divergence between training and test data distribution. Empirically this corresponds to the overall increasing trend of generalization gap shown in the right-most plots in Figure 2.

Reviewer comment (Experiments):
Additional experiments could be useful, such as an experiment comparing the generalisation gap of LapPE/BasisNet/SPE.

Author response:
Thank you for the good suggestions. 

1. Now we add the training loss and the generalization gap (test loss - training loss) on ZINC dataset as shown in Table 3, Appendix B.5 in the refined manuscript. We can see that though SignNet and BasisNet achieve a pretty low training MAE (high expressive power), their generalization gap is larger than other baselines (poor stability) and thus the final test MAE is not the best. For baseline GNN and PEG, they are pretty stable with small generalization gap, but the poor expressive power make them hard to fit the dataset well (training loss is high). In contrast, SPE has not only a lowest training MAE (high expressive power) but also a small generalization gap (good stability). That is why it can obtain the best test performance among all the models.

2. We are working on the TUDatasets you suggest and the results will be released before the end of rebuttal. LRGB consists of more and larger molecular graphs and is hard to be covered due to limited time of rebuttal period.

Reviewer comment (Experiments):
Additional experiments could be useful, such as evaluating the performance of LapPE/BasisNet/SPE on LRGB or TUDatasets.

Author response:
We are working on the TUDatasets you suggest and the results will be released before the end of rebuttal. LRGB consists of more and larger molecular graphs and is hard to be covered due to limited time of rebuttal period. 

We just get the new results on TUDataset as you suggested. Please kindly refer to Appendix B.8, Table 6 for results. Overall we can see that SPE achieves the best results on 3 out of 4 different datasets.

Reviewer comment (Theory):
The reason for Hölder continuity ($c\neq1$) is not fully clear. For example, does PEG / standard basis-net already satisfy your stability criterion Definition 3.1 and/or Assumption 3.1? If yes, what is the conceptual/theoretical benefit of your proposed architecture? If not, could you please provide an argument why the assumptions fail for PEG / basis-net?

Author response:
BasisNet does not have stability guarantee Def 3.1 because it does not satisfy Assumption 3.1. In short that is because we can see it as using discontinuous $\phi$ functions, which violates Assumption 3.1. Kindly check our unified response  ``*Why previous methods are unstable?* '' above. 

PEG is indeed stable, but its expressive power is extremely limited. This is becasue PEG cannot distinguish eigenvectors $V\in\mathbb{R}^{n\times d}$ and eigenvectors $VQ\in\mathbb{R}^{n\times d}$, where $Q\in O(d)$ is an arbitrary $d$-dim rotation matrix. This point is also summarized in our Introduction section `` But, to achieve stability, it (PEG) completely ignores the distinctness of each eigensubspaces
and processes the merged eigenspaces homogeneously. Consequently, it loses expressive power and
has, e.g., a subpar performance on molecular graph regression tasks''. And this also got verified by our experiments on ZINC, Table 1, which shows that the performance of PEG is limited. Additionally, we show that it is due to the high training loss of PEG in Table 3, Appendix B.5 in the revised manuscript. This verifies the poor expressive power of PEG.

Reviewer comment (Theory):
Is $c\neq 1$ crucial for any proof, guarantee, or assumption?

Author response:
$c\neq 1$ serves for describing the stability property of SPE, which is the general Holder continuity instead of Lipschitz continuity. For the second question, we think it is an open question. At this point we are not sure whether SPE can have a Lipschitz bound instead of a Holder bound. But fortunately even $c\neq 1$ is sufficient to provide a OOD generalization guarantee.

Reviewer comment (Theory):
Are there cases where SPE satisfies the stability assumption of Wang et al (ICLR 2022), i.e., with $c=1$?

Author response:
$c\neq 1$ serves for describing the stability property of SPE, which is the general Holder continuity instead of Lipschitz continuity. For the second question, we think it is an open question. At this point we are not sure whether SPE can have a Lipschitz bound instead of a Holder bound. But fortunately even $c\neq 1$ is sufficient to provide a OOD generalization guarantee.

Reviewer comment (Evaluation):
Please provide runtimes for your proposed method, preferably for pre-processing and overall runtime. This would help to put the achieved results into context with basis-net, etc.

Author response:
Thank you for pointing this out. We add the average running time of training and inference, which can be found in our unified response ``*running time evaluation*'' above. The pre-processing is the standard Laplacian decomposition same as others.

Reviewer comment (Theory):
If I am not mistaken, SPE and basis-net should have the same expressivity and thus at least theoretically be both equally capable of counting cycles, etc. Can you provide some intuition why SPE performs significantly better than basis-net in this task (Figure 3)?

Author response:
This is because BasisNet is not a stable method and thus generalizes poorly. Actually both SPE and BasisNet can fit pretty well on training dataset (training loss < 0.05). However, BasisNet suffers from overfitting with a large generalization gap (test error - training error). In contrast, the generalization gap of SPE is almost zero. This phenomenon is also observed on ZINC, as shown in Table 3, Appendix B.5 in the refined manuscript.