Equivariant Matrix Function Neural Networks

Thank you very much for your review. We appreciate that you found that our method "performs well on numerical experiments" and works "nicely with geometric data". In our revision, we made a significant effort to make the method more understandable. In the following, we respond to your questions and suggestions to further improve the paper. We respectfully hope that our responses will be satisfactory and increase your score. ### Relationship to spectral convolutions > The proposed method is an extension of spectral convolutional approaches which have enjoyed great success. The particular strategy proposed by the authors, via resolvent expansions, seems like a decent approach. We want to emphasize here that there are key differences between Matrix Function Networks and previous spectral convolution approaches. - In all spectral convolution approaches known to us, the convolution operators were used as filters. We believe that the step of parameterizing features in a neural network architecture to be learnable graph operators and updating the node features based on matrix functions of these operators is new. To be precise, previous spectral convolution approaches have an update rule of \$X^{t+1} = f(L)X^{t}\$ where \$L\$ is fixed (the Laplacian), while MFN acts as \$X^{t+1} = f(H(X^{t}))\$ where both \$H\$ and \$f\$ are learnable. In particular, each entry in the learnable matrices is a many-body function of the local environment of each edge. In subsequent layers, new operators will be learned based on these extract features. - Furthermore, no spectral convolution methods known to us included equivariance to Lie groups, and therefore these methods were not competitive for geometric point clouds. Our method to parametrize the operators naturally expend to geometric point clouds. ### Clarification on the matrix construction step > The method is laid out in section 4 and to be honest, I had quite a bit of trouble following along. We understand that we attempted to make Section 4 the most general possible. However, we now realize that this leads to a very abstract and abstruse definition. We have rewritten the matrix construction step to simplify the language, and we added a new section that gives examples of concrete implementation of the matrix construction in the case of the \$O(3)\$-group. We also have largely expanded Figure 1 to provide a better illustration of the method, including a visual representation of each step of the matrix function network (see https://ibb.co/2NhsmRb for a preview). We hope that these edits will enable a better understanding of the key contributions of our method. Here is a more condensed version of the key part of our approach, simplifying a bit the notation: - First, the method takes a graph with node labels and edge labels as entries and constructs **local** many-body node features using a graph neural network layer (for example, a convolution graph neural network for pure graphs, or an equivariant graph neural network for geometric graphs), that we call here \$V_i^{(t)}\$, \$V_j^{(t)}\$. - Second, a graph matrix \$H_{ij}\$ is learned by parameterizing its entries with a learnable function of the node features \$H_{ij} = \phi(V_i^{(t)}, V_j^{(t)})\$. The notion of the size of the basis is relevant only for geometric graphs, where a Lie group is acting, because this function \$\phi\$ needs to be equivariant. - The node features are updated using the output of a learnable matrix function of these matrices. We want to stress that this is a **key** novelty in the context of GNNs. The output features are now **non-local** features. ### On the experiments > First, there is no code shared [...] For example, it appears complex numbers are being counted as a single parameter per number whereas it should be two per number, see question below asking to clarify this. [...] I am not sure how certain features are being extracted. For example, the authors say that they take descriptors from the PDF model in the appendix for the ZINC data. 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.

### On the Combes Thomas Theorems > I simply could not follow section 4.3. What is the Combe-Thomas theorem? Also, doesn’t going from the bound stated in the paragraph before Eq (16) to the bound in Eq. (16) require some dependence on the norm of the weights? 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. ### Runtime comparison > Runtime comparisons are also not provided, and this can be a concern here, since the method requires spectral decompositions (or approximations thereof). 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 no other architectures can model, including Transformers. Due to the unstructured non-locality of global attention, Transformers cannot 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 want to stress that this solves an important open problem in the community of machine learning force fields (see [3]). We made changes to the contributions to better highlight this key point. ### On relationship with infinite layer MPNN > Later on, the geometric expressiveness statements are stated without proof and missing too many details for me to follow. First, what exactly are the authors saying is equivalent to the infinite layer MPNN (does the matrix channel need an infinitely sized pole expansion)? Where are the formal proofs of these statements? Does expressiveness here imply approximating a given function and if so, in what norm? What are two-body entries? 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. ### On the parameters count > Are these complex numbers counted as two parameters per number or one? 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.

### On the expressiveness and WL test > It seems this parameterization is a nonlocal form of message passing and the resulting layers manipulate the eigenvalues of the Laplacian or adjacency matrix. We know from expressiveness results (WL or otherwise) that this is limited because the eigenvalues are insufficient to encode all the information in a graph. Is my understanding here correct? Thank you for this interesting question. First, we emphasize that MFN layers do not manipulate the eigenvalues of the Laplacian or adjacency matrix. We learn operators that are arbitrary operators. In particular, when we apply the MFN for geometric graphs, the learned operators are group equivariant, while the Laplacian or adjacency would just be invariant. For expressiveness, there are two separate interesting cases: - The first case is about expressiveness on pure graphs (not embedded in a vector space). In this case, it is true that many graphs share the same spectrum for the Laplacian or for the adjacency. However, the MFN learns **multiple** graph operators. Therefore, a well-converged MFN layer can learn a set of operators whose spectrum maximally discriminates graphs relevant to the learning task. Moreover, the fact that multiple operators are learned, even if some graphs are cospectral for one of the operators, means that they need not be cospectral for the union of the spectrums of all the operators. Therefore, an MFN layer should be more expressive than simply using the Laplacian spectrum. - The case for geometric graphs is different. The question is related to the geometric WL-test [1]. We know that if the matrix entries are learnt using a $G$-MACE [2] layer with correlation order equal to the number of points of the graph, then the spectrum of this operator is complete, in the sense that it will be able to discriminate any graph that are not related by permutations or symmetries of $G$. ### Basis functions > What are the basis functions for pure graph inputs? If they are being learned, did you include this in the parameter count? 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. [1] On the Expressive Power of Geometric Graph Neural Networks, Joshi, et al [2] A General Framework for Equivariant Neural Networks on Reductive Lie Groups, Batatia, et al [3] Machine learning force fields, Oliver T Unke, et al

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'

Thank you very much for your review. We appreciate that you found our work "novel" and "strong empirical evidence". In the following, we respond to your questions and suggestions to further improve the paper. ### On the basis > It is stated that this choice should be made with consideration of the application domain. I'm curious to know what the limitations here are. When are we unable to automatically learn these functions? And are there any practical considerations for learning these? 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\] ### Conditions on the poles > Is it possible that we could learn weights and poles that lead to invalid curves, or otherwise undesirable behaviour? 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). ### Number of layers > As I mentioned above, the MFN models used in experiments always contain only a small number of layers. Was this a computational consideration? 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\]. ### Minor comments 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. \[1\] Atomic Cluster Expansion: Completeness, Efficiency and Stability, Dusson, et al \[2\] A General Framework for Equivariant Neural Networks on Reductive Lie Groups, Batatia, et al \[3\] MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields, Batatia, et al

Thank you very much for your review. We appreciate that you found our work "novel" and with "strong empirical evidence". In the following, we respond to your questions and suggestions to further improve the paper. We respectfully hope that our responses will be satisfactory and will increase your mark. ### General remark > There is room for improvement in the presentation. 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. ### Some details > Given a matrix \$A\$, \$A^{*}\$ is not defined (I guess it's adjoint of A). > \$f_{\theta}(H)\$ is defined as "scalar-to-scalar," but it seems "matrix-to-matrix". \$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]()). ### Action of \$G\$ on the matrix \$\textbf{H}\$ > In Eq. (10), since \$\textbf{H}\$ is a matrix, it's not obvious how \$g\$ acts on. I guess H is the function of and acts on it, but it's not clearly described. 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}\$. ### Proof of equivariance of matrix function \$H\$ > I couldn't find the proof of Eq. (10). 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. ### Computation cost of MFN > It is not easy to grasp the computational cost of MFN since it depends on the graph dimension [...] For example, you can evaluate the actual wall clock time and memory usage, which would be clear evidence that MFN is better than other methods, including Transformer. 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.

### Construction of \$H\$ in the \$O(3)\$ group > What is L in the caption of Fig 2? > It seems H is constructed with Kronecker products as in Eq. (21). What is the actual size of H? 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.

Dear Reviewer qFqk, Thank you again for your valuable feedback and comments! As the discussion period is ending soon, we would greatly appreciate it if you could let us know whether you are satisfied with our response. We will be happy to address any remaining concerns. Sincerely, Authors