Unitary Convolutions for Learning on Graphs and Groups

We thank the reviewer for their feedback and for their insightful questions. We have conducted additional experiments to answer the reviewers' main questions and will also amend the manuscript accordingly. We respond to individual points below.

 

---

While I generally appreciated the paper (hence my slightly positive score), I’m not fully convinced by the Lie UniConv approach. The authors propose the two methods I highlighted above as a solution to avoid overfitting. As overfitting is caused by a contraction of all frequencies associated to all but the lowest frequency of the Laplace operator, converting the adjacency matrix of the provided graph into a unitary operator (as it is done in UniConv) makes sense to avoid the issue (as all eigenvalues of the obtained diffusion operator now have norm equal to 1).

This is a neat perspective on the UniConv operator that we appreciate you highlighting. Indeed, one can view this as replacing standard graph convolution which is a dissipative process with a unitary one.

Let us now address the point about Lie UniConv. To continue, you write:

However, the Lie UniConv approach doesn’t go in that direction, but rather applies an exponential map on operator g_conv to obtain an orthogonal map. In this case, why the exponential of such operator is an orthogonal map is not specified in the paper, and I struggle to understand the rationale behind this approach.

We believe there are two points here that need clarifying.

First, we want to clarify what we think is a simple misunderstanding. The operator $g_{conv}$ is a linear map which lies in the Lie algebra of the orthogonal/unitary group. The matrix exponential maps from the Lie algebra to its Lie group, so the exponential map of this operator will be unitary/orthogonal meaning that $\|g_{conv}(X)\| = \|X\|$ and $g_{conv}$ is invertible. One can easily check that enforcing $W + W^\dagger=0$ enforces that the map $g_{conv}$ is in the Lie algebra.

Second, there is a question about the rationale. In practice, Lie UniConv and UniConv are both unitary, but the UniConv operator acts as a tensor product $\exp(iA)\otimes W$ whereas $\exp(g_{conv})$ cannot be written in this tensor product form. The reason we separate these two is because in practice one may only want to change the message passing to be unitary leaving the feature transformation acting as before. Also, the tensor product structure makes implementation faster as one can apply feature transformation and message passing separately in sequence. Remark 2 aims to highlight these points, which we will make clearer in the revised manuscript. To return to your intuition from earlier, both of these will act as diffusion operators that are unitary. The sole difference is in how one wants to transform the feature space of the nodes. For example, one may use Lie UniConv in situations where the 'magnitude' of message passing depends on the features.

We also refer the reviewer to the overall response to all reviewers for additional details about this. Since this was also a point of confusion for another reviewer, we will clarify this further in the paper (see overall response for specifics).

 

---

I didn’t see any result using Lie Uniconv besides the experiment of the synthetic long rage benchmark I highlighted above. On top to this, on such benchmark, UniConv is not applied in the paper. If this was achieving comparable, or even superior, results to Lie Uniconv, then there would not be any result in the paper showing the benefit of using the Lie Uniconv method.

We thank the reviewer for sharing this feedback which was also shared by other reviewers. There are now a number of updates for this (see overall response with attached pdf). We have added a table comparing Lie UniConv to UniConv on the Peptides and TU datasets including a split for whether one parameterizes in the orthogonal or unitary group. UniConv performs slightly better though the differences are minor. Furthermore, we include a figure showing that Lie UniConv and UniConv are both equally capable of learning the toy model task.

In practice, the differences between the two methods can be hard to discern. In real-world experiments, we often find that UniConv often performs slightly better since it addresses oversmoothing issues in the message passing leaving feature transformations unconstrained. Please see the overall response for more details.

I would kindly ask the authors if they could clarify my doubts above, and present additional results with Lie Uniconv in the rebuttal (unless I missed them somehow in the paper).

Please see prior response. We would be happy to answer any additional questions or perform additional experiments that could help clarify the contributions of our paper.

 

---

In addition to this, I would like to point out that the results presented in Table 2 are incomplete, and while the proposed Uniconv outperforms the reported methods, it is possible to see better performing solutions on Roman E in “A critical look at the evaluation of gnns under heterophily: Are we really making progress?”.

The reviewer is correct in that there are models in “A critical look at the evaluation of gnns under heterophily: Are we really making progress?” that perform better, for example on the Roman Empire dataset. However, these methodologies use different embeddings in their pre-processing. As noted in Appendix F under 'Heterophilous Graph Datasets', "we do not separate ego- and neighbor-embeddings, and hence also do not report accuracies for models from the original paper that used this pre-processing (e.g. GAT-sep and GT-sep)".

| Number of Nodes | Vanilla GCN | Spectral GCN | Gated GCN |   GPS  | Unitary GCN |
|-----------------|:-----------:|:------------:|:---------:|:------:|:-----------:|
| 200             |    0.0024   |    0.0054    |   0.0028  | 0.0047 |    0.0061   |
| 1000            |    0.0036   |    0.0106    |   0.0064  | 0.0124 |    0.0122   |
| 5000            |    0.0084   |    0.0383    |   0.0331  | 0.3009 |    0.0479   |
| 25000           |    0.0733   |    0.2640    |   0.1541  | 7.1890 |    0.1894   |
| 100000          |    0.2383   |    1.1645    |   1.1916  |   oom  |    1.3550   |

We thank the reviewer for their feedback and questions. As an overall point, the reviewer made many comments about normalized convolution which we were unsure of the specific message. We kindly ask the reviewer to clarify.

---

The originality of this approach is a bit limited, since normalizing message passing is not a very novel idea. Doing it with unitary constraints has not been explored before for GNNs.

Unitary layers normalize the 2-norm, but we politely disagree that they are simply a type of normalization as they also guarantee invertibility, lack of oversmoothing, stable Jacobian, etc.

What's the relative benefit compared to more naive normalizations?

As stated above, we are not sure what a "naive normalization" is referring to. We believe this may refer to other interventions listed below. We stress that that unitary layers are not competing with these, but complement them: any architecture can be implemented in tandem with such techniques.

1. batch norm: batch norm helps normalize layer outputs, but does not have useful properties of unitary layers listed above, e.g. vanishing/exploding gradients, over-smoothing, and stability.
2. regularizers and feature adjustments: There are methods to constrain, regularize or adjust features or layer outputs to make neighboring feature vectors more dissimilar and mitigate oversmoothing. These include Dirichlet energy constraints (e.g. Zhou et al. Dirichlet energy constrained learning...), Pairnorm (Yang et al. Pairnorm...), and others. These may help over vanilla methods in addressing oversmoothing though their benefit is arguably limited -- see survey of [RBM23] and comments below.
3. normalizing adjacency matrix: One typically normalizes adjacency matrices by multiplying it by $D^{-1/2}AD^{-1/2}$ where $D$ is a diagonal matrix with $i$-th entry as the degree of node $i$. This is by default used in message passing architectures, and we use it too. This to us is a separate topic from how to perform message passing on an adjacency matrix.
4. residual connections: Though not a type of normalization, residual connections have been proposed to help avoid oversmoothing. We experimentally compare to this intervention in our experiments and also give a theoretical argument for why residual connections alone can only alleviate the problem (see Proposition 4).

For more detail, consider a concurrent paper [SWJS24] using feature normalizations and residual connections. They have weaker empirical results on datasets such as Mutag and Proteins with their normalization, GraphNorm, and residual connections. We can include more details if this is what the reviewer means by "naive normalization".

---

Other questions/concerns:

if [instabilities with depth] is widely observed, please indicate a few references where this was observed?

Instabilities come in various forms and have inspired various architectectural changes, including for CNNs and GNNs (see related work). We have also included some examples. Figure 4 in the appendix shows an example of instabilities with increasing depth for several message-passing GNNs. The attached one page response also shows this issue with vanishing/exploding gradients.

"linear convolution" -- as opposed to what other kind? isn't convolution always linear?

Linear here is linearity in the input $x$. Some convolutions are not of this form, e.g. steerable convolution acts over a set of potentially nonlinear filters.

"When W contains only real-valued entries, the above returns an orthogonal map" -- I don't follow this, why is that? this statement would require a proof, and in general it is false... $\exp(g_{conv})(X)\neq e^{A} X e^{W^*}$ in this case.

We think there is a simple misunderstanding. This fact is true for the simple reason that $\exp(W)$ is real-valued whenever $W$ is real-valued (the exponential map cannot return complex numbers when its input is real-valued). So the output is always orthogonal provided $W$ meets the Lie algebra criteria of eq. (7), namely $W+W^T=0$. We also do not claim that $\exp(g_{conv})(X)=e^{A} X e^{W^*}$. In general, it will not take this tensor product form. Perhaps the reviewer can clarify further their concern?

I don't see what's the importance / relevance of examples 1 and 2 from Section 3. I think these are standard examples, they can be skipped or moved to the appendix.

Please see general response. We will shorten section 3.2 by moving the discussion of the Fourier/spectral based methods to the appendix.

Saying that Prop 6 "gives a proof that unitary convolution avoids oversmoothing" is unjustified without further proofs/details...

We will be more precise stating that UniConv provably avoids oversmoothing as measured by the Rayleigh Quotient or Dirichlet Energy, but we push back that this is simply "vague intuition". The Rayleigh quotient is a widespread measure of oversmoothing (see references below). Beyond this, unitary message passing is invertible and isometric, two properties that ensure signals always propogate through a network without getting "lost".

[vanishing/exploding gradients and dynamical isometry]: what is the connection of this definition with the definitions from section 3.1?

Due to lack of space available here, please see response to reviewer vJVB asking the same question.

 

References:
[AL24] Roth, Andreas, and Thomas Liebig. Rank Collapse Causes Over-Smoothing and Over-Correlation in Graph Neural Networks. PMLR, 2024.
[RBM23] T Konstantin Rusch et al. A survey on oversmoothing in graph neural networks. arXiv:2303.10993, 2023
[RCR+22] T. Konstantin Rusch et al. Graph-coupled oscillator networks. PMLR, 2022.
[SWJS24] Scholkemper, Michael, et al. "Residual Connections and Normalization Can Provably Prevent Oversmoothing in GNNs." arXiv preprint arXiv:2406.02997 (2024).
[WAWJ24] Wu, Xinyi et al. "Demystifying oversmoothing in attention-based graph neural networks." NeurIPS (2024).

No problem; we are happy to clarify. Let us answer your questions in order.

1. Formally, unitary transformations must be invertible so the input and output dimension must be equal. Thus the row and column dimensions of $W \in \mathbb{C}^{b\times b}$ must be equal and similarly for $A \in \mathbb{C}^{a \times a}$. However, this does not mean that $X$ has equal row and column dimension. The map $f_{Uconv}$ maps matrices of size $a \times b$ to matrices also of size $a \times b$ so dimensionality is preserved in this sense. As noted in the draft though, this does not limit implementation: for $W$ with arbitrary row and column dimension, one can also work within the Stiefel manifold or even more simply pad inputs and outputs so that dimensionalities are equal to actually implement things in practice (see appendix D).
2. It does not need to hold that $AXW =- (AXW)^\dagger$. In fact, this will not hold true in general since the row and column dimensions of $X$ are not equal so this equation formally may not even make sense. What does hold true however is that the linear operator $g_{conv}$ is in the Lie algebra. As a reminder, exponential map is equal to $\exp(g_{conv})(X) = \sum_{k=0}^\infty \frac{g_{conv}^{(k)}(X)}{k!} = \sum_{k=0}^\infty \frac{A^kXW^k}{k!}.$ Viewed in the vectorized form from our previous comment this is equivalent to: $vec[\exp(g_{conv})(X)] = \sum_{k=0}^\infty \frac{(A \otimes W^T)^k}{k!} vec[X],$ so as stated earlier, one can check that $A \otimes W^T$ is in the Lie algebra to ensure unitarity.

Please let us know if we can clarify anything further.

Ok, thanks for the detailed reply. I hope it won't bother you if I continue a bit with the clarifications on the same topic (sorry if it does):

1. (minor point) So if $A=A^T$ and $W=-W^\dagger$ then these are square matrices, and then $X$ must be a square matrix, right? Or how can I compute $\exp(AXW)$ if $AXW$ is not a square matrix? (I'm confused by the notation $X\in \mathbb C^{a\times b}$, is it that necessarily $a=b$?)

2. Assume now that $A=Id$, and $X,W$ are real-valued square matrices of equal dimension $n$, of which $W=-W^T$, to simplify things. Say $A=Id$, and $W=\left[\begin{array}{cc}0&-1\\\\ 1&0\end{array}\right]$ and $X=\left[\begin{array}{cc}0&1\\\\ 0&0\end{array}\right]$. Then $AXW=\left[\begin{array}{cc}1&0\\\\ 0&0\end{array}\right]\neq -(AXW)^T$ right? Then $\exp(AWX)$ is not unitary or orthogonal. I am confused, so maybe clarifying what I misunderstood with this example can help.

We thank the reviewer for praising the novelty and theoretical results in our work. We have read their concerns about the format and writing of the paper, and hope to address them below.

---

In Section 3.1 (Lie) UniConv layers are introduced... it was not clear that this is the central definition.

Can the reviewer expand on what they mean by central definition? From our understanding, the definitions of UniConv operations are central to the paper.

The paper generalizes the setting of unitary convolutions... it is unclear what this section intends to communicate...

Our goal was two-fold: to show that unitary convolution on graphs is a simple method applicable to general symmetric domains, and to connect it to previous CNN implementations, demonstrating how to generalize beyond graph convolution, such as applying unitary convolution to hypergraphs or other message-passing types. That being said, we understand the reviewer's concerns. We will shorten this section in the paper focusing specifically on the points above. Please see our general response for details.

Dynamical isometry is defined as a measure as a property to analyze exploding gradients and an example is given that a combination UniConv together with a yet completely undefined activation function (GroupSort) is “perfectly dynamically isometric”... Why is an example important enough to warrant the definition of Dynamical isometry? Are vanishing / exploding gradients a problem with GNNs?

Vanishing/exploding gradients are a general problem with any deep or recurrent network and occur because norms of the states at intermediate layers grow or decay exponentially with depth (see e.g. [LMTG19], [LMGK21]). GNNs are no different from deep networks in that they also have weights initialized randomly causing norms to grow or decay as a product of the number of layers. We found that standard GNNs cannot train at very large depths (see figure 4 showing standard message passing schemes do not train well at around 40 layers for example). We also show an example of a standard GCN featuring vanishing/exploding gradients in the one page response pdf (UniConv does not have this issue in contrast).

For scalability to many layers, stable information and gradient propagation is essential. Dynamical isometry offers a rigorous framework to study this stability, supported by experiments and theory across various architectures. We aim to show that GNNs also exhibit the stability properties described by dynamical isometry.

What is GroupSort and is it actually used in practice?

Groupsort is a nonlinearity that splits the input vector into pairs of numbers and sorts the pairs (see eq. 58). It is a nonlinearity that preserves isometry and used in the adversarial learning literature to get provable adversarial bounds via unitary architectures. E.g., see its use in Trockman & Kolter Orthogonalizing convolutional layers with the cayley transform.

For the graph distance dataset the model is not mentioned in the text. In this case, the model is GCN with the novel Lie UniConv layer (Figure 2). However, it is not explained why the Lie UniConv layers are used and not the UniConv Layers. Similarly, for the experiments on LRGB UniConv layers are used but it is not explained why Lie UniConv layers are not tried.

Please see global response and additional figures/tables for a full answer to this. In summary, this choice did not make a significant difference. For experiments where networks were not very deep, we found that enforcing unitarity in the message passing was helpful and unitarity was not needed in the feature transformation. For example, UniConv performs slightly better in LRGB though little difference is observed.

To me it seems like this could be a good paper with interesting ideas and solid experiments...I think more time is required to allow the authors to rewrite their paper... there is a lot of useful information hidden in the appendix and it might make sense to focus the paper on the main idea (Section 3.1) and move other parts to the appendix to make more space.

If we understand the reviewer correctly, they have requested that we focus more on section 3.1, defer parts of section 3.2 to appendix, and make more space for clarifying some points such as when to use Lie UniConv vs. UniConv. We will make these changes in the revised manuscript. Are there other parts of the paper that the reviewer feels should be moved from the appendix to the main text or vice versa?

We are hesitant to make significant changes to the writing and format given this concern was not shared by all the reviewers, but are of course open to more feedback.

(Lie) UniConv are implemented as a k-th taylor approximation. What is the value for k?... knowing this value seems important.

We found in all our experiments that $k=10$ suffices. The error in the Taylor approximation exponentially decreases with $k$. This value is in line with that in other papers (e.g. see [SF21]). We will state this in the main text.

---

 

Dynamical Isometry References:

[SMG13] Saxe et al. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. arXiv:1312.6120, 2013
[XBSD+18] Xiao et al. Dynamical isometry and a mean field theory of cnns: How to train 10,000-layer vanilla convolutional neural networks. ICML, 2018
[PSG18] Pennington et al. The emergence of spectral universality in deep networks. PMLR, 2018

 

Vanishing/ Exploding Gradient Problems with GNNs: In addition to above, see also:

[LMTG19] Li et al. Deepgcns: Can gcns go as deep as cnns?. Proceedings of the IEEE/CVF international conference on computer vision, 2019.
[LMGK21] Li et al. Training graph neural networks with 1000 layers. PMLR, 2021.

I thank the authors for their thorough rebuttal. The additional experiments do strengthen the paper and including the results for both architectures does in my opinion simplify the experiment section.

Can the reviewer expand on what they mean by central definition? From our understanding, the definitions of UniConv operations are central to the paper.

Yes, they are. However, for some reason this was not immediately clear for me when reading this section. Maybe a bit more guiding of the reader is required.

If we understand the reviewer correctly, they have requested that we focus more on section 3.1, defer parts of section 3.2 to appendix, and make more space for clarifying some points such as when to use Lie UniConv vs. UniConv. We will make these changes in the revised manuscript. Are there other parts of the paper that the reviewer feels should be moved from the appendix to the main text or vice versa?

You understand me correctly, I think that moving 3.2 to the appendix is a good idea. This should allow you to focus more on your essential results (key definitions + theorems/propositions) and intuitively explain what they mean and why they matter.

We are hesitant to make significant changes to the writing and format given this concern was not shared by all the reviewers, but are of course open to more feedback.

I agree with this sentiment. While I do think that my issues with the presentation / clarity persist (and this is somewhat mirrored by feJe), I cannot ignore that this does not seem to be an issue for the other reviewers. Furthermore, upon a further look I have to concede that this paper does contain many non-trivial theoretical insights. I will slightly increase my score but adjust my confidence downward.

We thank the reviewer for their positive comments and feedback. We respond to their questions and comments below.

 

---

Though the general form of convolution is provided. It is not evaluated in experiment. I will suggest putting these content in the appendix in this case. Otherwise, it is better to provide the corresponding experimental results.

We thank the reviewer for pointing this out. Our primary goal was to show a general recipe for constructing equivariant unitary/orthogonal maps, which allows for introducing more general convolution on finite groups beyond just graph convolution. Our main application are GNNs, hence, much of the experiments are focused on the graph setting. We provide experimental results for another instance of the general case, namely group convolution over the dihedral group, in Appendix E.3. We note that previous literature has studied unitary/orthogonal CNNs which also fit within this framework. We discuss and reference these works in the extended related works (see Appendix A.1). We will add more explicit pointers to these sections in the main text.

 

---

The computational cost is higher compared to other methods, but there are no experimental results to show the important details.

The unitary and orthogonal layers have a constant factor overhead over standard message passing. As stated in the paper, the runtime is $O(KnD)$ where $K$ is the approximation in the exponential map, $n$ is the number of nodes, and $D$ is the maximum degree of the graph. Since the approximation error is exponentially small in $K$ (see Appendix C), we found that setting $K$ to be a constant (around 10 for example) sufficed in all experiments. We will specify this in the main text in the revised manuscript.

Note, that this is more efficient than many graph algorithms such as graph transformers which are quadratic scaling in $n$ and other proposed (near) unitary graph message passing algorithms which also scale poorly in $n$ and $E$. The computational cost is detailed further in Appendices A and C. We should also note that this factor $K$ overhead is in some sense unavoidable if one wants invertibility and isometry as detailed in Proposition 4.

We thank the reviewer for their positive comments and going through the paper in detail. There are many insightful questions and points of feedback that we will incorporate into the draft.

---

There aren't many experiments on Lie UniConv... it would have been interesting to see comparison between the two

We agree with the reviewer on this point and have included many comparisons in the global response and attached one page pdf there. This includes experiments for Peptides and TU datasets. In general, the empirical differences are rather minimal.

The part on unitary maps in Fourier basis feels somehow detached from the general discourse of the paper... the theory to understand it is relatively advanced... a short primer on representation theory in the Appendix... would have been of great help. There are no experiments in this setting.

We aimed to show that Fourier space can often speed up operations and simplify convolution as observed in many works. Since this is tangential to the main points, we'll move it to the appendix with a detailed background. In the main text, we'll briefly note that a Fourier basis methodology exists.

The part on unitary maps in Fourier basis feels somehow detached... if one considers using the eigendecomposition of some graph Laplacian as a spectral basis, the (graph) Fourier transform is not given by a set of matrices varying over irreps but a signal of the same dimensionality of the input function.

We believe there is a small misunderstanding here. One can diagonalize the adjacency matrix of an undirected graph and perform convolution over the spectral basis as in spectral graph convolution. When the input and output is a single channel, this can be treated as block diagonal operations, but the blocks are trivially scalars ($1 \times 1$ matrices). More generally, the block size will depend on the input/output channel dimensions. We will provide a lemma explicitly showing this in the more general case for arbitrary graphs/groups.

Sometimes authors leave facts implicit... I had difficulty understanding why Lie UniConv returns an unitary map before looking at the tensor product version of (Eq. (36), Proposition 6). Also the statements in Propositions 4 and 7 are not really clear from a first reading, some more intuition about the statements should be given.

We agree this was slightly confusing on a first read. We will revise these passages for clarity in the updated manuscript.

 

---

Questions:

I don't see how the graph convolutions defined are particular cases of the group convolutions described. Is it the symmetric group S(n) acting on the graph?

To map graph convolution to this case where $conv_G(X) = \sum_i T_i X W_i$, simply set $T_i= A^i$ or the $i$-th power of the (potentially normalized) adjacency matrix. We will explicitly state this after the equation in the main text.

Line 122: Why UniConv has a "tensor product nature"? From what I understand the tensor product is more related to the Lie UniConv because you can write it as $\exp(A \otimes W^T)$.

The reviewer is right that Lie UniConv is parameterized in a tensor product fashion, but crucially, the output is not a tensor product: there may not exist matrices $M_1, M_2$ such that $\exp(A \otimes W^T) \neq M_1 \otimes M_2$. In contrast, UniConv is always of the form $\exp(iA)\otimes W^T$ and one can apply the feature and message passing transformations separately and in sequence. We tried to capture this in remark 2, but will clarify in the main text.

Line 186: "We set input/output representations to be equal so that the exponential map of an equivariant operator is itself equivariant": Assuming this does not limit the class of groups or transforms we can use? Not all groups have irreps of same dimensions...

This statement is written confusingly and we will clarify. What we meant to say is that the input/output representations need to act on the same vector space. For a unitary transformation to be well defined, it needs to be invertible and isometric so the input and output vector space are must have the same dimension. One can generalize beyond these assumptions, but it seems unnecessary for our purposes.

Line 212: is (I − A) a type of graph Laplacian?

That is correct.

Line 272: Why authors don't compare to new architectures as well on this task, like in Table 1? Also I see in Appendix that GPS it is performing pretty well w.r.t. Unitary GCN but it was not put in this graph (maybe to show better improvements?).

We show GPS for sake of completeness, but did not include it in the main plots because it is not a message passing architecture. The goal of the toy model task was to see which architectures can learn long range dependencies with local operations. For transformers, the notion of long range dependency is not as well defined.

Is there another method that the reviewer feels we should compare to? We are happy to include it.

Line 763: I think authors should say that in case of Algo 2 this follows because Fourier transform is unitary. Line 766: "For unitary convolution in the Fourier domain (Algorithm 2), equivariance follows from convolution theorems...". This is too succinct and should be more expanded formally.

Yes, we will say that and expand on this. Thank you for pointing it out.

Line 821: missing trace on following lines and last line should not be put there ($=R_G(X)$)

We think there is a misunderstanding. The following lines have the notation $vec(X)^\dagger ... vec(X)$ which is an equivalent way of writing the trace as an inner product between vectors since $Tr(A^\dagger B)= vec(A)^\dagger vec(B)$ where $vec$ vectorizes the matrices. We will use inner product angle bracket notation to make this clearer. The last line also needs the normalization by the Frobenius norm which we will include as a separate line. Thank you for pointing this out.