Paper: Equivariant Matrix Function Neural Networks (ICLR-2024)

Reviewer comment (Theory):
There is confusion regarding the indexing in equation (7). While $ij$ seems to refer to node indexing, the matrix is $bn \times bn$, making $H_{cij, m_1, m_2}$ inconsistent with how matrix entries work.

Author response:
We thank the reviewer for iterating on the feedback and helping us improve the quality of the submission.

Based on this feedback, we have further revised the manuscript to use improved notation for the matrix indices. Additionally, we have added a section in the appendix on runtime comparison, which includes timings for the MFN models and the MACE model on the cumulene example. We have also added the following reference in the 'Linear Scaling' paragraph [1], providing a more formal exposition of the statement presented in the text.

[1] T. A. Davis, 'Direct Methods for Sparse Linear Systems'

Reviewer comment (Experiments):
There is uncertainty regarding how certain features are extracted. For example, the authors state they take descriptors from the PDF model in the appendix for the ZINC data, which would be unfair if those features come from a trained model on the same dataset.

Author response:
We have attached a copy of the two codes used to run the experiments, both for the Cumulene case and the pure graph case. The parameter count is correct; see the response below for more details, but only the poles are complex parameters. The PDF descriptors are **not** extracted from a trained model; we just use the Laplacian and fixed powers of the Laplacian of the graph to extract initial edge labels to enrich the information in pure graphs as initial embedding. We want to stress that this does not affect the number of parameters, and this extraction was not trained in any way. They are not related to **anything** learnable, just a fixed precomputable initial embedding. As our MFN learns custom operators beyond the Laplacian, we wanted to include this operator as a baseline. Sorry for the confusion this has generated.

Reviewer comment (Theory):
Section 4.3 is not understandable. The Combe-Thomas theorem is mentioned without citation, and the step from the bound stated in the paragraph before equation (16) to the bound in equation (16) may require dependence on the norm of the weights.

Author response:
We understand that this section was very technical and we moved most of this to the appendix and used the available space to clarify Section 4.1 on the core of the architecture. We now cite the Combes-Thomas theorem paper [1]. 

Thank you for highlighting the dependence on the norm of the weights, this is an important point that we forgot to mention. We now added a brief explanation that the bounds can be taken uniform for suitable bounds on the operator which are guaranteed by our normalization procedure.

Reviewer comment (Theory):
The geometric expressiveness statements are stated without proof and missing too many details to follow. Questions arise such as what exactly is claimed to be equivalent to the infinite layer MPNN, whether the matrix channel needs an infinitely sized pole expansion, where the formal proofs are, in what norm expressiveness is defined, and what two-body entries are.

Author response:
Although we agree that a full proof of this statement would be interesting, we believe that it is beyond the scope of this paper due to the very limited length. That is why we did not claim any proof but intended to present that as a remark for future investigation. We moved this part of the paper to the appendix and used the gain space to clarify other sections. We also have toned down the statement to show that it is a well-motivated remark rather than a theorem. Here, we give an intuition of the relationship between the two approaches:
Any analytic matrix function \$f\$ admits a formal power series expansion, valid in its radius of analyticity,
\begin{equation}
    f(H) = \sum_{k=0}^{\infty} c_{k} H^{k}
\end{equation}
where \$c_{k}\$ are the complex (or real) coefficients of the expansion. Lets look at the diagonal elements of each of these powers,

$(H)^{2}_ {ii, c00} = \sum_ {j \in \mathcal{N}(i)} H_ {ij,c0m}  H_ {ji,cm0}, \quad (H)^{3}_ {ii, c00} = \sum_ {j \in \mathcal{N}(i)}  \sum_ {k \in \mathcal{N}(j)} H_ {ij,c0m} H_ {jk,cmm'} H_ {ji,cm'0}$

The sum over \mathcal{N}(i) is forced due to the sparsity of the constructed matrix \$H\$ in our method. Therefore, we observe a convolutional structure similar to message-passing methods. The power \$2\$, corresponds to a single convolution, the power \$3\$ to two convolutions, all the way to infinity. Each coefficient of the matrix function that is **learnable** can act like weights in the linear update of the analogous layer. Because the power series expansion is infinite, one sees a clear analogy between the two.

Reviewer comment (Evaluation):
When reporting parameter counts, are complex numbers counted as two parameters (real and imaginary parts) or one? A fair comparison should count each complex number as two parameters.

Author response:
Thank you for your inquiry; all parameters were appropriately counted. We want to stress that only the poles of the resolvent expansion are complex parameters, the rest being real. The poles usually represent at most 256 parameters (times two for real and complex), which is negligible in parameter counts.

Reviewer comment (Theory):
What are the $\phi_m$ basis functions for pure graph inputs?

Author response:
We have added a more detailed explanation of the matrix construction in pure graph that was previously in the appendix. Here is a short summary, using the new notation introduced in Section 4 to clarify our architecture. We use a convolutional graph neural network to learn the node features $V_{i}^{(t)}$ for each node $i$. We then use a simple multilayer perceptron $H_{ij,c} = \text{MLP}(V_{i}^{(t)}, V_{j}^{(t)}, e_{ij}^{(t)})$, where $e_{ij}^{(t)}$ are a concatenation of fixed edge attributes (see the PDF explanation) and learnable edge features extracted from the off-diagonal blocks of the matrix function. This was counted in the parameter count. We hope that the code will help clarify these points.

Reviewer comment (Reproducibility):
If $\phi_m$ functions are learnable, were they included in the parameter count?

Author response:
We have added a more detailed explanation of the matrix construction in pure graph that was previously in the appendix. Here is a short summary, using the new notation introduced in Section 4 to clarify our architecture. We use a convolutional graph neural network to learn the node features $V_{i}^{(t)}$ for each node $i$. We then use a simple multilayer perceptron $H_{ij,c} = \text{MLP}(V_{i}^{(t)}, V_{j}^{(t)}, e_{ij}^{(t)})$, where $e_{ij}^{(t)}$ are a concatenation of fixed edge attributes (see the PDF explanation) and learnable edge features extracted from the off-diagonal blocks of the matrix function. This was counted in the parameter count. We hope that the code will help clarify these points.

Reviewer comment (Experiments):
Regarding the statement “The initial edge features ... are obtained by concatenating PDF [Yang et al. (2023)] descriptors used for ZINC”, are these descriptors the outputs of the trained model? If so, this would be unfair for comparison.

Author response:
We have attached a copy of the two codes used to run the experiments, both for the Cumulene case and the pure graph case. The parameter count is correct; see the response below for more details, but only the poles are complex parameters. The PDF descriptors are **not** extracted from a trained model; we just use the Laplacian and fixed powers of the Laplacian of the graph to extract initial edge labels to enrich the information in pure graphs as initial embedding. We want to stress that this does not affect the number of parameters, and this extraction was not trained in any way. They are not related to **anything** learnable, just a fixed precomputable initial embedding. As our MFN learns custom operators beyond the Laplacian, we wanted to include this operator as a baseline. Sorry for the confusion this has generated.

Reviewer comment (Theory):
One area of the paper that is not completely understood is the introduction of a choice of basis (Section 4.1). It is stated that this choice should be made with consideration of the application domain. This seems to align with spectral GNNs that lean on hand-designed features for effective learning, but the authors point out that this can be approximated well in some cases. The limitations of this are unclear — when are we unable to automatically learn these functions, and are there practical considerations for learning them?

Author response:
Thank you for an excellent question. In the case of pure graphs, enumerating all possible graph operators is NP-hard. The approach we take to do the parameterization in a data-driven way takes the assumption that the model would converge to learn the optimal set of operators for the specific task. There are no limitations in the parameterization except that the matrix needs to be permutation invariant. However, for a graph with N nodes, the entry of the matrix might be an N-body function of all the other nodes. In our experiments, we simply used a multilayer perceptron to learn these entries, but it is likely that other solutions perform better. There are many theoretical open questions related to learning these operators, linked to the cospectrality to a set of spectrums of operators. In the case of geometric graphs (which are embedded in a vector space), the answer can be made more complete. In this case, the entries are \$G\$-equivariant functions of the graph. The proof of completeness of equivariant architectures like ACE or MACE for general Lie groups shows that for an infinitely large basis, any of these functions can be uniformly approximated with arbitrary accuracy. \[1,2\]

Reviewer comment (Theory):
It is unclear whether we need to ensure that the learnable weights and poles satisfy any constraints for the resolvent calculus definition (Section 4.2). From the explanation in the paper, the choice of weights and poles determine the curve implicitly via the quadrature rule. It is questioned whether it is possible to learn weights and poles that lead to invalid curves or otherwise undesirable behaviour.

Author response:
Yes, poles that are too close to being purely real lead to divergences. Normalization of the matrix was fundamental to ensure that the optimizer did not explore these points. In practice, finding a good parameterization of the poles was important to making the method work. We parametrize the poles as \$z_{s} = w_{1}\exp(w_{2}) + w_{3}\$, with \$w_{1}, w_{2}, w_{3}\$ learnable weights initialized from a standard Gaussian. For the weights, conditions on the norm of the weights would mean targeting a specific space of functions. We are just initializing these weights with a Gaussian centered on zero and a variance one, and let the optimizer move them freely. However, we believe that more clever initialization based on some inductive bias might improve the convergence in applications where one knows the kind of matrix function that gives rise to non-locality (e.g. Fermi Dirac in electronic systems).

Reviewer comment (Experiments):
The MFN models used in experiments always contain only a small number of layers, and it is unclear whether this was purely a computational consideration.

Author response:
For the pure graph case, it was mostly for restricting the number of parameters and for computation consideration. We intend to implement specialized CUDA kernels for fast inversion as outlined is Section 4.2, that would enable us to scale to larger number of matrices. For the molecular case, we do not see any improvement when increasing the number of layers past 2. This is expected from previous work with many-body methods \[3\].

Reviewer comment (Writing):
There are a couple of symbols that are not defined in the main text. These can mostly be inferred from context, but should ideally be defined. Above equation (1), the domain and codomain of \( s \) should be defined. The definition of \( \mathfrak{Z} z \) and related symbols used in Section 4.3 could not be found.

Author response:
Thank you for pointing out these small notation issues. We changed to the more conventional $\text{Imag}$ to denote the imaginary part in Section 4.3. We also have added a reference to PerceiverIO, mentioning that better scaling can be achieved with newer implementations.

Reviewer comment (Theory):
The paper cites transformers as having a problematic quadratic scaling. While this is true in the original formulation, it is now common to use methods with far more efficient scaling — for example, PerceiverIO can reduce the computational cost to linear.

Author response:
Thank you for pointing out these small notation issues. We changed to the more conventional $\text{Imag}$ to denote the imaginary part in Section 4.3. We also have added a reference to PerceiverIO, mentioning that better scaling can be achieved with newer implementations.

Reviewer comment (Presentation):
There is room for improvement in the presentation, as several incomplete descriptions were found.

Author response:
We tried to make the architecture as general as possible. However, we now see that this abstraction made the method quite hard to grasp. We have now added several sections, including examples of the MFN architecture for specific groups, to exemplify the architecture in a concrete setting. We also have rewritten part of Section 4 to clarify the matrix construction step. Finally, we have improved Figure 1 to give a better illustration of the method. We hope that has significantly improved the overall presentation of the method.

Reviewer comment (Writing):
Given a matrix \(A\), \(A^*\) is not defined (reviewer guesses it is the adjoint of \(A\)).

Author response:
\$A^{*}\$ is indeed the adjoint of \$A\$ in the paper. Thank you for spotting these; we have added clearer introductions to all the notation in the paper, including the adjoint. The ambiguous notation of the function \$f\$ is to highlight that any real value function can be "overloaded" to a symmetric matrix value function (matrix to matrix) by simply applying it point-wise to its eigenvalues. An alternative view if \$f\$ is analytic is to consider the formal power series expansion of the function \$f\$, and formally replace the scalar powers to matrix powers. In this context, it is widely accepted to use the same symbol for the real value map and the matrix value map (see, for example, the Wikipedia page on matrix functions [https://en.wikipedia.org/wiki/Analytic_function_of_a_matrix]()).

Reviewer comment (Theory):
\(f_\theta(H)\) is defined as "scalar to scalar," but it seems to be "matrix to matrix".

Author response:
\$A^{*}\$ is indeed the adjoint of \$A\$ in the paper. Thank you for spotting these; we have added clearer introductions to all the notation in the paper, including the adjoint. The ambiguous notation of the function \$f\$ is to highlight that any real value function can be "overloaded" to a symmetric matrix value function (matrix to matrix) by simply applying it point-wise to its eigenvalues. An alternative view if \$f\$ is analytic is to consider the formal power series expansion of the function \$f\$, and formally replace the scalar powers to matrix powers. In this context, it is widely accepted to use the same symbol for the real value map and the matrix value map (see, for example, the Wikipedia page on matrix functions [https://en.wikipedia.org/wiki/Analytic_function_of_a_matrix]()).

Reviewer comment (Writing):
In Eq. (10), since \(H\) is a matrix, it is not obvious how \(g\) acts on \(H\). The reviewer guesses \(H\) is a function of \(\sigma\) and \(g\) acts on it, but this is not clearly described.

Author response:
Yes, \$g\$ is a group action that acts on the state of each of the nodes. At the matrix level, the action is given by Eq. 9, since the functions \$\phi\$ are assumed to be equivariant. The action of \$G\$ propagates from the actions in the states to the action in the matrix. A good example is for 3D graphs. In the case \$G=O(3)\$ the group of 3D rotations and the element \$g\$ are rotation matrices \$R\$. The state of each node would be a 3D vector of positions \$(x, y, z)\$. Rotations can be applied to the graph by rotating each position \$(x, y, z) \to R (x, y, z)\$. The function \$\phi\$ being equivariant will result in matrix entries that are also equivariant. Therefore, if a rotation is applied to the whole graph, transforming by a rotation \$R\$, then there exists an orthogonal matrix \$D(R)\$ such that the graph operator \$H_{1}\$ of \$G_{1}\$ is transformed into \$H_{2} = D(R) H_{1} D(R)^{T}\$.

Reviewer comment (Theory):
The proof of Eq. (10) could not be found.

Author response:
We have written a detailed proof of the equivariance of generic continuous matrix functions in the appendix and reference it in the main text. Thank you for noticing that.

Reviewer comment (Evaluation):
It is not easy to grasp the computational cost of MFN since it depends on the graph dimension \(d\), the choice of the update equations (11)–(13), etc.

Author response:
We agree that computational cost is a key parameter in any method. As explained in Section 4.2 in the Linear Scaling Paragraph, the prototype implementation has a cubic scaling and is therefore quite slow compared to other approaches (a factor of x2 to x5 compared to local models). However, fast inversion methods are well established, and we explain how this scaling can be reduced significantly to less than quadratic scaling, all the way to linear scaling, by exploiting the sparsity of the matrices we construct and selected inversion methods. 

Although these methods are known, they require significant effort to implement and integrate them into machine learning libraries like Pytorch due to the need for specialized CUDA kernels. We intend to make this effort and release open-source code in future work, but we believe that this is beyond the scope of this paper, where we focus on the novelty of the architecture and its expressivity. 

Finally, we stress that one of the key results of this paper is that MFN can model systems that Transformers cannot model. Due to the unstructed non-locality of global attention, Transformers can not extrapolate to large systems in the Cumulene example. This means that regardless of speed, there are applications in which MFNs are clearly the only solution to date. We made changes to the contributions to better highlight this key point.

Reviewer comment (Evaluation):
A clear comparison of computational cost is needed. For example, actual wall clock time and memory usage could be evaluated to provide clear evidence that MFN is better than other methods, including Transformer.

Author response:
We agree that computational cost is a key parameter in any method. As explained in Section 4.2 in the Linear Scaling Paragraph, the prototype implementation has a cubic scaling and is therefore quite slow compared to other approaches (a factor of x2 to x5 compared to local models). However, fast inversion methods are well established, and we explain how this scaling can be reduced significantly to less than quadratic scaling, all the way to linear scaling, by exploiting the sparsity of the matrices we construct and selected inversion methods. 

Although these methods are known, they require significant effort to implement and integrate them into machine learning libraries like Pytorch due to the need for specialized CUDA kernels. We intend to make this effort and release open-source code in future work, but we believe that this is beyond the scope of this paper, where we focus on the novelty of the architecture and its expressivity. 

Finally, we stress that one of the key results of this paper is that MFN can model systems that Transformers cannot model. Due to the unstructed non-locality of global attention, Transformers can not extrapolate to large systems in the Cumulene example. This means that regardless of speed, there are applications in which MFNs are clearly the only solution to date. We made changes to the contributions to better highlight this key point.

Reviewer comment (Presentation):
In experiments — What is \(L\) in the caption of Fig. 2?

Author response:
The \$L\$ index corresponds to the maximal spherical order in the basis expansion of the matrix for the spherical case. In order to make this notation clearer and improve the presentation of the paper, we have added a section to detail the implementation of the matrix construction case in the \$O(3)\$ group. We hope that this will significantly help the readers understand the method. \$H\$ is a tensor of size \[\$N_{\text{nodes}} \times (L+1)^{2}, N_{\text{nodes}}\times(L+1)^{2}, N_{\text{channels}}\$], where \$N_{\text{nodes}}\$ is the number of nodes in the graph, and \$N_{\text{channels}}\$ is a hyper-parameter corresponding to the number of operators we are learning.

Reviewer comment (Reproducibility):
In experiments — It seems \(H\) is constructed with Kronecker products as in Eq. (21). What is the actual size of \(H\)?

Author response:
The \$L\$ index corresponds to the maximal spherical order in the basis expansion of the matrix for the spherical case. In order to make this notation clearer and improve the presentation of the paper, we have added a section to detail the implementation of the matrix construction case in the \$O(3)\$ group. We hope that this will significantly help the readers understand the method. \$H\$ is a tensor of size \[\$N_{\text{nodes}} \times (L+1)^{2}, N_{\text{nodes}}\times(L+1)^{2}, N_{\text{channels}}\$], where \$N_{\text{nodes}}\$ is the number of nodes in the graph, and \$N_{\text{channels}}\$ is a hyper-parameter corresponding to the number of operators we are learning.