Equivariant Matrix Function Neural Networks
Authors: Ilyes Batatia
Lars Leon Schaaf
Gabor Csanyi
Christoph Ortner
Felix Andreas Faber
Venue: ICLR 2024 spotlight (2024)
Keywords: equivariance
graph neural networks
long range
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)

Abstract: Graph Neural Networks (GNNs), especially message-passing neural networks (MPNNs), have emerged as powerful architectures for learning on graphs in diverse applications. However, MPNNs face challenges when modeling non-local interactions in systems such as large conjugated molecules, metals, or amorphous materials.
Although Spectral GNNs and traditional neural networks such as recurrent neural networks and transformers mitigate these challenges, they often lack extensivity, adaptability, generalizability, computational efficiency, or fail to capture detailed structural relationships or symmetries in the data. To address these concerns, we introduce Matrix Function Neural Networks (MFNs), a novel architecture that parameterizes non-local interactions through analytic matrix equivariant functions. Employing resolvent expansions offers a straightforward implementation and the potential for linear scaling with system size.
The MFN architecture achieves state-of-the-art performance in standard graph benchmarks, such as the ZINC and TU datasets, and is able to capture intricate non-local interactions in quantum systems. The code and the datasets will be made public.

Published as a conference paper at ICLR 2024
EQUIVARIANT MATRIX FUNCTION NEURAL NETWORKS
Ilyes Batatia Lars L. Schaaf Gábor Csányi' Christoph Ortner* Felix A. Faber'
University of Cambridge, UK Uniiersity of British Columbia, Canada
ABSTRACT
Graph Neural Networks (GNNs), especially message-passing neural networks (MPNNs), have emerged as powerful architectures for learning on graphs in di- verse applications. However, MPNNs face challenges when modeling non-local interactions graphs such as large conjugated molecules, and social networks due to oversmoothing and oversquashing. Although Spectral GNNs and tradi- tional neural networks such as recurrent neural networks and transformers miti- gate these challenges, they often lack generalizability. or fail to capture detailed structural relationships or symmetries in the data. To address these concerns, we introduce Matrix Function Neural Networks (MFNs), a novel architecture that pa- rameterizes non-local interactions through analytic matrix equivariant functions. Employing resolvent expansions offers a straightforward implementation and the potential for linear scaling with system size. The MFN architecture achieves state- of-the-art performance
standard graph benchmarks, such as the ZINC and TU datasets, and is able to capture intricate non-local interactions in quantum systems, paving the way new state-of-the-an of-the-art art force fields.
INTRODUCTION
Graph Neural Networks (GNNs) have proven be powerful architectures for learning on graphs on a wide range of applications. Various GNN architectures have been proposed, including message
passing neural networks (MPNN) (Gimer et al. 2017 Battaglia et al., 2018. & Welling 2017a
Kipf
Veličković et al., 2018: Wu et al. Semior et al. 20188] Hu u et al. 2019 Batzner et al. 2022 and
higher-order equivariant MPNNs (Batatia et al. 20226).
MPNNs struggle to model non-local interactions both due to computational constraints and over- smoothing (Di Giovanni et al. 2023). Spectral Graph Neural Networks attempt address the lim- itation of this kind by encoding the global structure of a graph using eigenvectors and eigenvalues of a suitable operator. These approaches predominantly focus on Laplacian matrices, exploiting the graph's 's inherent spectral features. Many spectral GNNs apply polynomial or rational filters (Bianchi et al. 2021b Gasteiger et al. 2018 Wang & Zhang 2022 He et al., 2021 Defferrard et al. 6b) Zhu et al. 2021 Kreuzer al.) 2021) eigenvalues of graph structures, reaching state-of-s oiee--fftthe the-hhhrr art accuracy on pure graphs tasks. However, these methods often exhibit rigid architectures that re- quire extensive feature engineering, potentially limiting their adaptability to various types of graphs. Moreover, they have been restricted to non men-geomttric graphs, making them unsuited for molecules
and materials.
Traditional neural network architectures, such as recurrent neural networks (Elman 1990; Hochre- iter & Schmidhuber 1997] Cho al., 2014: Graves 2013) and transformers (Vaswani et al. 2017) also face challenges when modeling non-local interactions. While transformers can capture some non-local dependencies through their self-attention mechanisms, they come significant com- putational cost due to their quadratic complexity with respect to the input sequence length, even existing methods can mitigate their cost Jaegle et al.) (2022). Furthermore, transformers lack in- herent structural relationships positional information within input data, necessitating the use of additional techniques, such as positional encodings (Vaswani et al. 2017: Shaw et al. 2018).
In chemistry and material science tasks, some models incorporate classical long-range interactions
through electrostatics (Grisafi & Ceriotti 2019 Behler [2021] Unke & Meuwly 2019), disper-
al. 2023 Kosmala et al.) 2023). However, no existing architecture effectively address non-local interactions, where quantum
sion, reciprocal space (Gao & Remsing 2022 Buguenin-Dumiteee
effects can propagate over extensive distances through electronic delocalization, spin coupling, or

---

Published as a conference paper at ICLR 2024
other many-body non-local mechanisms. This is particularly problematic in systems such as large ente conjugated molecules, amorphous materials, or metals.
Consequently, there is need for new neural network architectures that can efficiently and accu- rately model complex non-local many-body interactions, while addressing the limitations of current approaches. We propose Matrix Function Networks (MFN) possible solution to this challenge. Concretely, we make the following contributions.
We introduce Matrix Function Networks (MFNs), new graph neural network architecture able model non-local interactions in a structured, systematic way.
We introduce the resolvent expansion as convenient and efficient mechanism to learn a general matrix function. This expansion can in principle be implemented in linear scaling cost with respect to the size of the input.
We demonstrate the ability of our architecture to learn non-local interactions on a dataset of challenging non-local quantum systems, where standard GNNs architectures, including those incorporating global attention, fail to give even qualitatively accurate extrapolative predictions.
We show that MFNs achieve state-of-the-ar perfooomnn on ZINC and TU graph datasets.
RELATED WORK
Overlap matrix fingerprints (Zhu et al., 2016) introduced overlap matrix fingerprints (OMFPs), vector of spectral features of atomic environment more generally point cloud. Given point cloud, an overlap operator (identity projected on an atomic orbital basis) is constructed, and its taken as the features of that point cloud. Although a the- oretical understanding of OMFPs is still lacking, computational experiments have shown excellent properties distance measure (Zhu et al., 2016; Parsaeifard et al. 2021).
ordered eigenvalues (or other invariants)
Spectral Graph Neural Networks Spectral GNNs (Wu et al. 2021) are GNNs that use spectral filters operating on the Fourier decomposition of the Laplacian operator of the graph. Spectral GNNs are categorized by the type of filters they apply to the spectrum of the Laplacian: ChebyNet (Def-) ferrard et al.] 601 approximates the polynomial function of the Laplacian using Chebychev expansions, GPRGNN (Chien et al.) (022)) directly fits coefficients of a fixed polynomial, while ARMA (Bianchi et al.) 2021a) uses rational filters.
Equivariant Neural Networks Equivariant neural networks are the general class of neural net- works that respect certain group symmetries (Bronstein et al. (021). Noteably, convolutional neural networks (CNNs) (LeCun et al.] (1989) are equivariant to trans laanslatoon, while G-convolutions (Cohen & Welling 2016 Cohen et 20.8 Kondor & Trivedi 2018) generalized CNNs to equivariance of
compact groups Lately, equivariant message passing neural networks (Anderson et al. 2019 Sator-
ras et al. 2021
Brandstetter et al., 2022 Batzner al. 2022 Batatia al. [002bb a) have emerged as a powerful architecture for learning on geometric point clouds. Most of these architectures have been shown to lie in a common design space Batatia et al. (202aa 2023).
Hamiltonian Learning A natural application of equivariant neural network architectures is ma- chine learning of (coarse-grained) Hamiltonian operators arising in electronic structure theory. This task of parameterizing the mapping from atomic structures to Hamiltonian operators is currently re- ceiving increasing interest because of the potential extension accessible observables over purely mechanistic models. The recent works Nigam et al. (2022) Zhang et al. (2022) introduce such parameterizations in terms of a modified equivariant Atomic Cluster Expansion (ACE) (Drautz) 2020), precursor to the architecture we employ in the present work. Alternative approaches in-
clude (Hegde & Bowen 2017 Schütt et al. 2019 Unke et al. 2021a Gu et al. 2023).
3 BACKGROUND
3.1 SPECTRAL GRAPH NEURAL NETWORKS
We briefly review spectral graph neural networks and explain their limitations that our work will overcome. Consider a graph g = (X,E) with node set X and edge set E. A graph defined purely by its topology (connectivity) is called a pure graph. . Let denote the number of nodes in the graph, and let A E R=x x be its adjacency matrix. Define a vector of ones as 1, E R.. The degree matrix
of the graph is D diag( A n,, and the Laplacian matrix is L D A. The Laplacian is
2

---

Published as a conference paper at ICLR 2024
symmetric positive semidefinite matrix and admits a spectral decomposition, L = UAU", where U is an orthogonal matrix of eigenvectors and A is a diagonal matrix of eigenvalues.
A popular approach to learning functions graphs is to convolutional neural networks on the graph. Spectral graph convolutional networks (SGCNs) are a class of graph convolutional networks that use the spectral decomposition of the Laplacian matrix to define convolutional filters. Let s E IR" be a function of a graph G and be a convolutional filter. SGCNs take advantage of the spectral decomposition of the Laplacian matrix of graph to compute the convolution of and t:
= kg 8 = t((UTt) o s)) - [o()ss
(I)
where fo is a matrix function of the Laplacian matrix L of the graph & and @ the Hadamard product. Various works have proposed different matrix functions fa. such as polynomial functions (Defferrard et al., 2016a), rational functions (Bianchi et al. 2021a). or neural networks (Wu et al.) 2021). Two non-isomorphic graphs
share the same spectrum of their Laplacian operators, while they have different spectra for other graph operators (Johnson & Newman 1980). Therefore, the use of other graph operators as basis for matrix functions can be beneficial for learning functions on graphs that are strictly more expressive than Laplacian SGCNs. However, this approach has two main limitations.
Expressiveness: Performing only convolutions with fixed graph operator, (usually the Laplacian), limits the expressivity of the model. Choosing the most expressive graph oper- ator requires problem-dependent feature engineering.
Lie-group symmetries: The approaches proposed so far have been restricted pure graphs and in particular do not use additional symmetries for graphs embedded in a vector space. For example, graphs embedded in R33 often lead to (3))--vqivari learning tasks.
To overcome these limitations, we propose MFNs in Section 4 which allow parameterization of both the graph operator H and the matrix function fe. and can be formulated to preserve the equiv- ariance under all known group actions. Since learnable operator H prevents the precomputation of diagonalization of H during training, we also introduce method that avoids diagonalization and allows (in principle) linear scaling with the number of nodes.
3.2 EQUIVARIANT MESSAGE PASSING NEURAL NETWORK
Equivariant Message Passing Neural Networks (MPNNs) Batzner et al. 2022); Batatia et al. (2022b) graph neural networks that operate on graphs G (X, E) embedded in a vector space V. The nodes F; E X are no longer only a list of indices, but belong to a configuration space Q that extends the vector space V. For example, in the atomistic point cloud, v; (i, (i, 0,) Ω N R x Z describing the positions and chemical species of each atom through which the graph is embedded into V R The case of the pure graph can be recovered by setting Ω N. We are interested in learning graph maps of the form
Φ msets(()) + Z
(2)
where is an abstract target space, usually a vector space and msets( the multi-set of states. As the input is graph, we impose the mapping to be permutation invariant (invariant under relabeling of the nodes). In many applications, the target properties satisfy additional symmetries: When group G acts on both (and, therefore, on 9) and Z, we say that + is G-equivariant if,
o g p
Ag G,
where p is representation of the group on the vector space Z. typical strategy is then to embed the nodes x; E X into feature space, where a suitable representation of the group is available.
We represent the state of each node o, in the layer of the MPNN by a tuple.
(4) where x; defines the collection of node attributes of the graph as defined previously and h are its learnable features. A forward pass of the network consists of multiple message construction, update, and readout steps. During message construction, message m is created for each node
(vi h (1)).
3

---

Published as a conference paper at ICLR 2024
MFN Layer
[x:]
embed
Local Graph
[
Matrix Construction
Matrix Function
fd (wnt
Update
[
MFN Layer
leamaale
Create learnable emberddingg from node strfaunts/hhatures log MPMN MACE, CoNN...
MFN Layer
v(O v/o
(V v
(N
Evs (2st H
Update using nun-local features
m
MFN Layer
O-O-O-O
Construct sparse couivar mEDI
Readout
target
Oieam non-localty from Isarmable matr fanction
Figure 1: Matrix function network architecture. Illustrating matrix construction and non-locality of matrix functions molecular graph.
by pooling over its neighbors.
m; (1)
+ JEN(()
( (0) ,o,"), n((t)== _o (,()))mm()))
( dour ({(R(c01)))1
where the individual operations have the following meaning:
Mj is a learnable message function and 4 )ev(()) is a learnable, permutation invariant pooling operation over the neighbors of atom (c.g., a sum);
Ue is a learnable update function, transforming the message m, (t) into new features h ((+1);
R₂ is learnable node readout, mapping node states of to the per-node outputs;
wont is global readout map, typically øout ([(R))) }:},) ):)))) Σ., R(".
(t)
Equivariant MPNNs are widely used for learning properties of 3D point clouds, such as molecules and matcrials. However, there are several limitations to their expressiveness,
Non-local MPNNs are restricted to a finite range, dictated by the number iterations T. A large number of properties in physics are long-range, which requires 7 - co. Adding large number of layers leads to high computational overhead and poor expressivity due to oversmoothness (Di Giovanni al. 2023).
Correlation order: Most MPNNs use two body messages, which means that the pooling in Equation 5 is just a sum. The MACE architecture (Batatia et al., (022b)) extended the message construction to arbitrary order, but increasing the order is still computationally demanding, especially for large T.
MATRIX FUNCTION NEURAL NETWORKS
4.1
THE DESIGN SPACE OF MATRIX FUNCTION NEURAL NETWORKS
Our MFN models act in the space group equivariant matrix operators of the graph, that we denote H( where most commonly G O(3) the rotations group (including in examples), but framework applies to the more general class of reductive Lie groups (e.g., the Lorentz group in high energy physics SU(3) in quantum chromodynamics). In the same setting as in section 3.2 let G be an undirected graph with n nodes, and let the states a, be embedded in a configuration space Ω. The architecture operates in four stages each layer, a local graph layer that learns node features from the states, the matrix construction, the matrix function, and the update. These are visually outlined in figure [
Local graph layer For each iteration i, the first step in an MFN layer is to form equivariant node features using a local graph layer (abstractly noted L) as a function the local environment node i,
(0) JJEN()) ((1) (a)
(())
tem
4

---

Published as a conference paper at ICLR 2024
where v:(0) the learned equívariant node features and m indexes the representation of G acting on V. i.e. Vm g = Emm Pmmme (g) vm) We denote the dimension of the representation by M and the channel dimension by c. In the case where G = O(3), the m indices correspond to the usual angular momentum. The specific choice of the graph layer L depends on the application. One could use any equivariant architecture to learn these node features, such as a layer of equivariant MPNNs. In the case of the rotation group, any such function can be approximated with arbitrary accuracy using an ACE or MACE layer (Zhang al. 2022] Batatia et al.] 20226). This result has been generalized to any reductive Lie group, using the G-MACE architecture (Batatia et al., 2023). is important to note that the use of more expressive approximators will result in better coverage of the operator space H()) and therefore better general expressivity.
Matrix construction The second step involves constructing a set of graph matrix operators from the learned node features. The space of graph matrix operators, H())) corresponds to the space of operators that are (1) self-adjoint, (2) permutation equivariant, and (3) G-equivariant. The entries of the matrix operator correspond to learnable equivariant edge features expanded in the product of two representations my and m2. Hence, for each channel c, the equivariant operators square matrices H. E M consistent with x n blocks of size M x M where the number of nodes. Each nonzero block Hei corresponds to learnable edge features between the nodes and j. such that
(7) where i.j denote the indices of the matrix blocks, while m1, my represent the indices within each block. Additionally, (1) @omjma is a learnable equivariant function mapping a tuple of node features to an edge feature. For a concrete example of matrix construction in the case of isometry group O(3),
H exj.ming - mg
5 if E (i) and é E E(),
see section 4.3 and see Figure [ for an illustration of the block structure.
The full matrix inherits the equivariance of node features,
Hog g ((g) (g), E
where p is a block diagonal unitary matrix that denotes the group action of G on the basis element gb In the case where G is the rotation group, the blocks of p correspond to the usual D-Wigner matrices.
Matrix function The central operation of the MFN architecture is the matrix function, which introduces long-range many-body effects. Any continuous function fe : RR + R, with parameters 0. can be used to define a matrix function that maps square matrix to another square matrix of equal size. Formally, matrix function on self-adjoint matrices H can be defined by its spectral decomposition. Using the fact that H is diagonalizable such that H UAUT with U orthogonal and A diagonal,
(H) = Uf)
where f(A) is a diagonal matrix obtained by applying f to the each diagonal element of A. An essential observation is that any continuous matrix function fo preserves equivariance (see proof in Appendix A.6).
(H g) The matrix function can be related to a high- many-body equivariant function via the Cayley- Hamilton theorem. The cigendecomposition in Equation ils responsible for the non-locality of our approach. In practice, computing matrix functions is expensive as requires diagonalization, scaling as n with the number of nodes. Many approaches are available to approximate matrix functions, such as Chebyshev polynomials or rational approximation, which can leverage potentially cheaper evaluation schemes. Furthermore, the matrix H is sparse in many applications, which can be further exploited to reduce computational cost. To further optimize this, we propose to employ a resolvent expansion to parameterize fe. detailed in Section 4.2 Similar approaches have been successfully applied to large matrices in other fields such as electronic structure calculations (Lin et al.) 2009 2013).
(g), E.
(10)
Update The diagonal update updates the state of each node with non-local features extracted from the diagonal blocks of the matrix function,
h ((+1)
im = v(t)
tem +
Σ
w e(t)
ca m0
(11)
6
5

---

Published as a conference paper at ICLR 2024
This method is the most computationally efficient since selected inversion techniques (Lin et al. 2009) can be employed to efficiently evaluate the diagonal blocks of a matrix function; cf. Sec- tion Note that the diagonal blocks are symmetric and therefore extracting m0 Om is equiva- lent. Alternative updates can be defined from the matrix function that we detail in the appendix AA55 The optimal kind of update is trade-off between expressivity and computational cost. All of the updates differ fundamentally from the standard spectral GNN update, which is a filtering operation.
The node states are then updated using these new nonlocal node features 0; form the next states. The steps are repeated for T iterations, starting from the local graph layer.
(+++1 (z,, h (b(())) to
Readout
The readout phase is the same as the usual MPNN readout in Equation 5
4.2 RESOLVENT PARAMETERIZATION OF MATRIX FUNCTIONS
The evaluation of the matrix function in Equation is the practical bottleneck of our method. The cost of the evaluation depends on the choice of parameterization of the univariate function fo. For general analytic f C + C, resolvent calculus allows us represent
f(H) = [(=)(=I- H)-1 de
(12)
2
where C is curve encircling the eigenvalues of H and excluding any poles of f. Approximat- ing the contour integration with quadrature rule with nodes 2, and weights w. yields f(H) x E. (.f(z)) H)-1, ,, and merging Wa (2z) we arrive at the parameterization
(o = Σ w. W2,I H) H)
(13)
Pole expansions for evaluating matrix functions have a long and successful history, especially when the arguments are sparse matrices (Higham 2008). The novelty in equation [3] over standard usage of pole expansions in computing matrix functions (Higham 2008) is that both the weights w, and the poles 2, are now learnable parameters.
The derivation shows that in the limit of infinitely many pole-weight pairs (2,, W.) any analytic matrix function can be represented. Since analytic functions are dense in the space of continuous functions, this means that all continuous matrix functions can be represented in that limit well (at the same time letting the poles approach the spectrum). In practice, the poles should be chosen with non-zero Imaginary parts in order avoid the spectrum of H, which is real since H is assumed to be self-adjoint. Therefore, we choose adjoint pole weight pairs (w,, z.) and (w > to ensure that fe is real when restricted real arguments.
Linear scaling cost The pole expansion framework is the first key ingredient in linear scaling electronic structure methods Goedecker (1999) such as PEXSI[Ln et al. (2009) 2013). The sec- ond ingredient is the selected inversion of the resolvent. Instead of computing the full inverse, (=I-H)-1 one first computes sparse L.DL" factorization and then selectively computes only the diagonal entries of the resolvent. The bottleneck in this approach is the LDL" factorization. For full matrix, it scales as O(r²²) operations and O(n²) memory. The complexity improves consider- ably for sparse matrices. Suppose that the sparsity pattern is d-dimensional, corresponding topo- logically d-dimensional graph; e.g. the cumulenes in Section 5.] are topologically one-dimensional despite being embedded in 3R3 Using nested dissection ordering minimize the fill-in, the cost of the LDL factorization reduces to ((n) for d = 1 (e.g., natural language processing and quasi- (22222 operations and O(n log n) memory for = 2 (e.g., Image recognition); and Onn?) operations and (n44/3) memory for d = 3. A more for- mal exposition of this material is given in Davis (2006). The final step reduce the computational cost of our architecture to linear scaling is to replace the L.DI. factorization with an incomplete factorization, as proposed by Ettee (2020) in the context of electronic structure calculations. This would lead to approximations in non-licality that need to be investigated in future work.
1D molecular systems such as
4.3 EXAMPLE OF MEN FOR THE 0(3) GROUUP
Here we detail a concrete example for the MFN layer in the case -rroup. For concreteness, will consider point clouds of atoms input. We will reuse this architecture in the Experiment

---

Published as a conference paper at ICLR 2024
section. At each layer i, the states of each node described by tuple o, = (re, Oc, 1e (. where ri E Rd denotes the position of atoms in the 3D space, 0; E is the nuclear charge and n! the learnable node features that the MFN layer will update that we initialize just as learnable embedding of the species. The node features are expanded on the spherical basis and use the usual convention of the Im index (with the correspondence with the single in Section 4.1 being (m =
+ 0, 0) and 1 (I 1, m -1) etc.). For the local graph layer. we use a MACE layer (see Batatia et al. (2022b) for details) that constructs the local node features Vi, in Equation 6 We then construct the matrix using tensor product of the features of the nodes on nodes and J.
aijd Re(ray) @ jdgg))) j (i))) and NNO).
(14)
where R coorrssponss to a function of the distance riy between and j. As the graph is undirected, the matrix is symmetric with respect j. The matrix is sparse, as only the elements corresponding to
Φ x
H
PY
17/22
Figure 2: Block structure of Euclidean MFN operator, H. Each entry in H corresponds different product of representations of the group of rotations (ss, sp, pp, ...). Example for L=1.
an edge of the graph are nonzero. This matrix exhibits block structure illustrated in Fig. similar to Hamiltonian matrices in quantum mechanics with and sp orbitials. We refer to L as the maximal value of each spherical index i and la of the matrix (0 S l1,2 S L), and we refer to this architecture as (MACE) in the Results section. I.
4.4 EXPRESSIVITY OF MATRIX FUNCTION NETWORKS
Purely formally, one may think of matrix function /(H) as an infinite power series. This suggests that MFNs are inherently non-local and exhibit a convolution structure, similar to message- passing methods with infinite layers and linear update (see Appendix A.2 for details). This of interest for modeling causal relationships or non-local interactions by proxy, such as in chemistry or natural language processing. In these cases, the propagation of local effects over long distances results in multiscale effects that are effectively captured by method. The degree of non-locality of the interaction MFN can be precisely quantified using the Combes-Thomas theorem (Combes) & Thomas 1973). We provide detail analysis in the Appendix A.I A.1
RESULTS
5.1 CUMULENES: NON-LOCAL 3D GRAPHS
We compare the non-locality of MFNs local MPNNs and global attention MPNNs using linear carbon chains, called cumulenes. This example is notorious challenge for fate-of-the-an architec- tures machine learning force fields Unke al. (2021c), and our main motivation for introducing the MFN architecture.
Cumulenes are made up of double-bonded carbon atoms terminated with two hydrogen atoms at each end. Cumulenes exhibit pronounced non-local behavior as result of strong electron delocal- ization. Small changes in chain length and relative angle between the terminating hydrogen atoms can result in large changes in the energy of the system, visually represented in Figure 3 These structures are known to illustrate the limited expressivity of local models (Unke et al., 2021c) and similar to the k-chains introduced by Joshi et al. (2023) in the context of the geometric WL test. Frank et al. (2022) showed that global attention is capable of capturing the angular trends of cumulenes with fixed length. We go beyond and demonstrate that MFNs capable of accurately
7

---

Published as a conference paper at ICLR 2024
extrapolating longer chains, simultancously capturing length and angle trends. In contrast, global attention models such as Spookynet (Unke et al. 2021b), are unable to extrapolate to longer chains, highlighting the benefit of the matrix function formalism. For all MFN models in this section, use MACE layers (Batatia et al. 2022b) to form the matrix at each layer. We refer to the model as MFN (MACE). Details of the training and the specific choice of parameterization of the matrix entries are included in the Appendix A.9
DFT
MFN free
MFN (re:
(r.
1.5
receptice
Field
0.2
rel. energy [eV]
1.0
0.0
0.5
0.0
-0.2
12
-0.5
5
10
15
20
25
30
20
40
60
80
Figure 3: Visualizing MFN expressivity on cumulene chains. The left panel depicts energy trends with respect to cumulene chain length at fixed angle 5%. The right panel shows the DFT (ground truth) and the predicted energy function of the dihedral angle & between the hydrogen atoms for a cumulene chain containing 12 carbon (MACE) are only able capture average trends, atoms. Local many-body equivariant models though test configurations the training set. Invariant MFNs (L = 0) capture only the trends with respect to length, while equivariant MFNs (L capture both non-local trends. All models have cutoff' distance re of 3A, corresponding to the nearest neighbors, with two message-passinn layers. The cutoff distance well MACE's receptive field for the first carbon atom is annotated in the left panel.
included in
Trends with chain length and rotation The lowest energy structure of cumulenes alternates be- tween 90- and O-degree angles for odd and even carbon atom counts, respectively. Consequently, varying the number of carbon atoms at fixed angle unveils distinctive zigzag pattern in the ground truth energy (Fig.Bleet panel). Although the local model, with two message passing layers, is trained on exactly these configurations, this system is beyond its expressivity,
can be seen by the linear trend for ne > 7 (Figure BiggureBleff panel). In contrast, the invariant and equivariant MFN models per- fectly reproduce density functional theory (DFT) energies, thanks to their inherent non non-locality. To demonstrate the necessity of cquivariance, we train models on the energy of a fixed size cumulene function of the dihedral angle between the hydrogen atoms. Figure demonstrates that only the equivariant MFN (L=1) captures the angular trend.
Guaranteed Non-Local dataset Datasets designed to test non-local effects often yield unexpectedly high accuracy when evaluated with local models (oovacs et al., (2023), complicating the assessment of model non-locality. The dataset introduced here is based cumulenes, whose strong electronic delocalization results in a directly observable non-locality. The training set contains geometry- optimized cumulenes with 3-10 and 13, 14 carbon atoms, which are then rattled and rotated at various angles. The test set contains cumulenes created in similar fashion with the same num- ber of carbons (in-domain) and cumulenes of unseen length, not present in the dataset (out-domain 11,12 and 15,16). Table shows that the MFN architecture significantly outperforms both local and attention- based models (Spookynet). Attention captures some non-locality, resulting in marginally lower errors the train and in-domain test set. However, the learned non-locality does not gener- alize to larger molecules, obtaining energy and forces worse than those obtained with a simple local model. The structured non-locality of MFNs enables generalization to larger unseen system sizes.
5.2 PERFORMANCE ON PURE GRAPHS
In this section, evaluate the performance of our MFNs models in graph-level prediction tasks using GCN layers for the matrix construction. Detailed of the datasets and hyperparameters of the models can be found in the Appendix A.9
8

---

Published as a conference paper at ICLR 2024
(me Spookyyei (Global ostention))
PAc:
MACE (Local)
MEN(MOCK) (ours)
MACE (Local)
Spacks (Giobal uttention)
MFN
(ours)
Train
3-10,13,14
43.1
11.4
2,0
179.6
114.1
31.7
Test Demain)
3-10,13,14
41.8
205.6
162.3
34,0
(Out Domain)
11.12
1
(Out Domain)
15.16
Table 1: Guaranteed non-local cumulene dataset containing rattled cumulene chains, with various chain lengths (nc) and hydrogen dihedral angles (#). The table compares energy (E) and forces (F) RMSEs between local two-layer MPNNS (MACE), global attention MPNNs (SpookyNet), and equivariant MFNs. Train and in-domain test sets contain cumulenes of lengths 3-10 and 13.14. Two out-domain test sets compare different levels of extrapolation to unseen cumulene lengths, containing cumulenes with 11, 12 and 15, 16 carbon atoms, respectively. Bold is best and underline second best.
ZINC. We use the default dataset splits for ZINC, aligned with the leaderboard baseline settings, with approximately 500K parameters set. Table shows that MFN surpasses previous architectures, demonstrating the utility of learning various operators,
pure graphs
Table Results on ZINC with the MAE and number of parameters used, where the best results are
in bold. Baselines are taken from Yang et al. 2023) and model citations are in A.9.2
Method
GCN
GAT
MPNN
GT
SAN
Graphommer
PDF
MFN (GCN)
MAE
0.367=0.011
0.38410.007
0.14520:007
0.22610.014
0.13910.006
0.122:0.006
0.066.0.004
0.06310.002
Sparn
505k
481k
NA
50%
48%
472k
512k
TU Datasests. We test our model on five TUDataset datasets involving both bioinformatics datasets (MUTAG, ENZYMES, PTC MR, and PROTEINS) and social network dataset (IMDB-B). To ensure fair comparison with baselines, follow the standard 10-fold cross-validation and dataset split in table
Table 3: Results on TUDataset (Higher is better). Bold is best, and underlined second best within +0.5%. Baselines are taken from (Yang et al. 2023) and model citations are in A.9.3
Method
MUTAG
ENZYMES
PTC MR
PROTEINS
IMDB-B
GK
81.5242.111
32.70
55.6520.5
71.39:0.33
79.1142.11
2..16+1.64
55.91-0.3
59.57-00.1
PK
76.0::2.77
59.5:12.4
73.68
AWE
87.87449.76
35.774-5.93
74.45-+5.80
PSCN
88.95-14.4
62.29
75+2..5
71
ECC
76.11
45.67
DGK
87.44:12.72
$3.43:0091
60.08:2..
75.68:10.5
66.96:100.6
GraphSAGE
85.117.66
58.216.0
72.345..
CapsGNN
88.67-16.88
54.6715..7
76.2413.6
73.1
GIN
89.4++-.66
64.6++7.0
76,2+-2.8
75.1+5..1
J-GNN
86.1
60.9
75.5
74.2
IGN
83.89=12.95
58.53416.86
76.58:155.49
52.0+5.54
PPGNN
90.55 8.7
6641--4
77.30-14.73
73.0+5.77
GCN²
89.39:E1.60
66. 1.784111.79
71.111104
74.80:12.01
PDF
89.91:14.35
63.50
68.36:18.38
76.2815.11
75.60:12.69
MEN (GCN)
91.5=7.35
72.9+:7.55
68.9::8.09
76.18:14.07
7..1
CONCLUSION
We have introduced Matrix Function Networks (MFNs), an architecture designed to address the limitations of existing GNNs and MPNNs in modeling non-local many-body interactions. Utilizing a resolvent expansion, MFNs achieve potentially linear scaling with respect to system size, offering computationally efficient solution. Our evaluations indicate state-of-the-art performance ZINC and TU graph datasets without human-designed features to capture the global graph topology. We also demonstrate that our architecture is capable of modeling the complex non-local interactions of cumulene quantum systems. Future work could focus on extending MFNs to other complex systems, further validating its adaptability and efficiency, and exploring its interpretability.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
CO's work was supported by the NSERC Discovery Grant IDGR019381 and the NFRF Exploration Grant GR022937. LLS acknowledges support from the EPSRC Syntech CDT with grant refer- ence EP/S024220/1. Computational resources were provided by the Cambridge Service for Data Driven Discovery (CSD3), which was accessed through the University of Cambridge EPSRC Core Equipment Award EP/X034712/1. The authors would also like thank Rokas Elijošius for useful discussions.
ETHICS STATEMENT
FAF is employed by AstraZeneca at time of publication; however, none of the work presented in this manuscript was conducted at or influenced by this affiliation.
REPRODUCIBILITY STATEMENT
To ensure reproducibility and completeness, we include detailed descriptions the model used, hyperparameters. and data sets in the Appendix. The ZINC and TU datasets are pub-
licly available. The cumulene dataset is available at: https://glthur The MFN code will be available at: https:
Noom/LarsSchaff
Guaranteed hooo--ooolloolcllars-. //gathub.ccmm1Iyes551 wtt.ecoeoomyyyyymcc
REFERENCES
Brandon Anderson, Truong Son Hy, and Risi Kondor. Cormorant: Covariant molecular ral networks. H. Wallach, H. Larochelle, Beygelzimer, AlcheBue, E. Fox. and
R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL. hhtpp:: htttts::/prcceedinng neuri ps. ct. wwtpcp/ppppppppeeee 12aper/2011 /201 9/ffle/ 2117383322274666666..
Paper.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv:1607.066. 607.006500000...0 2016.
arXiv preprint
Ilyes Batatia, Simon Batzner, Dávid Péter Kováes, Albert Musaclian, Gregor N. C. Simm, Ralf Drautz, Christoph Ortner, Boris Kozinsky, and Gábor Csányi. The design space of eq3ivequi atom-centen centereeeerd interatomic potentials, 2022a. /tt.:tttpss//rxxi..rr////////// /205 06643
Ilyes Batatia, Dávid Péter Kováes, Gregor N. C. Simm, Christoph Ortner, and Gáhor Csányi. Mace: Higher order equivariant message passing neural networks for fast and accurate force fields, 2022b. URL[sttps:
// //arxiv.crrt abspaab...2222 0797
Ilyes Batatia, Mario Geiger, Jose Munoz, Tess Smidt, Lior Silberman, and Christoph Ortner. A general frame- work for equivariant neural networks reductive lie groups, 2023.
Peter W Battaglia, Jessica B Hamrick, Victor Bapst, Alvaro Sanchez-Gonzalez, Vinicius Zambaldi, Mateusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, et al. Relational inductive biases, deep learning, and graph networks. arXiv preprint arXiv: (006:v//8060/21. 2018.
Simon Batzner, Albert Musaelian, Lixin Sun, Mario Geiger, Jonathan P. Mailoa, Mordechai Kornbluth, Nicola Molinari, Tess E. Smidt, and Boris Kozinsky. E())-qqiivaiiat graph neural networks for data-efficient and accurate interatomic potentials. Nature Communications, 13(1):2453, 2022.
Jörg Behler. Four generations of high-dimensional neural network potentials. Chemical Reviews, 121(16): 10037-10072, 2021.
Filippo Maria Bianchi, Daniele Grattarola, Lorenzo Livi, and Cesare Alippi. Graph neural networks with convolutional ARMA filters. IEEE Transactions Pattern Analysis and Machine Intelligence, 2021a.
Filippo Maria Bianchi, Daniele Grattarola, Lorenzo Livi, and Cesare Alippi. Graph neural networks with convolutional arma filters, IEEE transactions on pattern analysis and machine intelligence, 44(7):3496 3507, b021b.
Johannes Brandstetter, Rob Hesselink, Elise der Pol, Erik Bekkers, and Max Welling, Geometric and physical quantities improve e(3) equivariant message passing, 2022.
10

---

Published as a conference paper at ICLR 2024
Michael M. Bronstein, Joan Bruna, Taco Cohen, and Petar Veličković. Geometric Deep Leamnnng Grids, Groups, Graphs, Geodesies, and Gauges, 2021.
Eli Chien, Jianhao Peng, Pan Li, and Olgica Milenkovic. Adaptive universal generalized pagerank graph neural network. In International Conference on Learning Representations, 2021.
Kyunghyun Cho, Bart Merriënboer, Caglar Gulcehre, Dzmitry Bahdanau, Fethi Bougares, Holger Schwenk, and Yoshua Bengio. On the properties of neural machine translation: Encoder-decoler ap- proaches, In Proceedings of the 2014 Conference on Empirical Methods in Natural Language Processing (EMNLP), pp. 1724 1734, 2014.
Taco S. Cohen and Max Welling. Steerable cnns. ICLR 2017, 2016. doi: 1048550/RXXX..61120089999 URL 1612.1222 08498
Taco S. Cohen, Mario Geiger, Jonas Köhler, and Max Welling. Spherical CNNs. In International Conference
on Learning Representations, 2018. URL /ttpo/////.///////////o////////r http://www.....er//oooooooalaaaa
J. M. Combes and L. Thomas. Asymptotic behaviour of eigenfunctions for multiparticle schrödinger operators. Communications in Mathematical Physics, 34(4):251-270, December 1973. doi: 10. 10.0000//00164477, URL
https: //cc///c.i.... .rg/ 10. heee///111110//111 10w./160111111 646473
James P. Darby, Dávid P. Kovács, llyes Batatia, Miguel A. Caro, Gus L. W. Hart, Christoph Ortner, and Gábor Csányi. Tensor-reduced atomic density representations. Phys. Rev. Lett, 131-028001, Jul URL atttps://iink aps org/doi/10
2023. doi: 10.1 R0..ios 13 131.028001.
dwi/l0 1103/
PhysRevLett 131. 028001
Timothy Davis. Direct Methods for Sparse Linear Systems. Society for Industrial and Applied Mathematics, siam.
2006. doi: 011137717780887718888
URL
97 978089871888
Pim de Haan, Taco S Cohen, and Max Welling. Natural graph networks. Advances in Neural Information Processing Systems, 33:3636-3644 2020.
Michaël Defferrard, Xavier Bresson, and Pierre Vandergheynst. Convolutional neural networks on graphs with fast localized spectral filtering. In Advances in Neural Information Processing Systems, pp. 2016a.
Michael Defferrard, Xavier Bresson, and Pierre Vandergheynst. Convolutional neural networks graphs with fast localized spectral filtering. Advances neural information processing systems, 29, 2016b.
Francesco Giovanni, Lorenzo Giusti, Federico Barbero, Giulia Luise, Pietro Lio, and Michael Bronstein.
On over- squsssuaahnn in message passing neural networks: The impact of width, depth, and topology. arXiv preprint arXiv:2302.0294 022441, 2023.
Ralf Drautz. Atomic cluster expansion of scalar, vectorial, and tensorial properties including magnetism and charge transfer. Phys. Rev. B, 102:024104, Jul 2020. doi: 10. 10. 1ww.110/////..B.B... 102.000240 URL https:
//link. wrlink.aps.cre www.pppppp..a.aa.aaa.aaaapaae..p /wwd/o.o11111/// http:///..113//hyeeeeeeee
102. 024104
Vijay Prakash Dwivedi, Chaitanya K Joshi, Thomas Laurent, Yoshua Bengio, and Xavier Bresson. Benchmark- ing graph neural networks. arXiv preprint arXiv:2003.009%: 2020.
Jeffrey L Elman. Finding structure time. Cognitive Science, 14(2):19--111, 1990.
Simon Etter. Incomplete selected inversion for linear-scaling electronic structure calculations, 2020.
Thorben Frank, Oliver Unke, and Klaus-Robert Müller. So3krates: Equivariant attention for interactions bitrary length- in molecular systems. Advances in Neural Information Processing Systems, 35:29400- 29413, 2022.
Ang Gao and Richard Remsing. Self-consistent determination of long-range electrostatics in neural network potentials. Nature communications, 13512, 2022.
Johannes Gasteiger, Aleksandar Bojchevski, and Stephan Günnemann. Predict then propagate: Graph neural networks meet personalized pagerank. arXiv preprint arXiv:1810.05000 2018
Justin Gilmer, Samuel S Schoenholz, Patrick F Riley, Oriol Vinyals, and George E Dahl. Neural message passing for quantum chemistry. In Proceedings of the 34th International Conference on Machine Learning- Volume 70, pp. 1263-1272. JMLR. org, 2017.
11

---

Published as a conference paper at ICLR 2024
Stefan Goedecker. Linear scaling electronic structure methods. Rev. Mod. Phys. 71:1085-11113 Jul 1999. doi: 10. 100//e//oo/hyyy.7.1055. URL
77. 085
Alex Graves. Generating sequences with recurrent neural networks. In arXiv preprint a.X8v:/308..000 2013.
Andrea Grisafi and Michele Ceriotti. Incorporating long-range physics in atomic-scall machine learning. The Journal of Chemical Physics, 151(20). 2019. ISSN 0021-9606. doi: 10.1063/1.1128375 URLhttps:
7/doi .ogg10.
/0. 1063/1 5128375 204105.
Qiangqiang Gu, Zhanghao Zhouyin, Shishir Kumar Pandey, Peng Zhang, Linfeng Zhang, and E Weinan, DecPTB: A deep learning-bses based hight-binding approach with ab initio accuracy. arXiv, July 2023.
Will Hamilton, Zhitao Ying, and Jure Leskavec. Inductive representation learning large graphs. In Advances in Neural Information Processing Systems, pp. 1024-1034, 2017.
Mingguo He, Zhewei Wei, Hongteng Xu, Bernnet: Learning arbitrary graph spectral filters via bernstein approximation. Advances in Neural Information Processing Systems, 34:1 34:14239-14251, 2021.
Ganesh Hegde and Chris Bowen. Machine-leamed approximations density functional theory hamiltonians. Sci. Rep., 7:42669, February 2017.
Nicholas J Higham. Functions of matrices: theory and computation. STAM, 2008.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. In Neural computation, volume 9,8, pp. 1735-1780. MIT Press, 1997.
Zhixuan Hu, Li Jin, Yi Yang, and Ronghang Zhang. Deep graph pose: semi-supervised deep graphical model for improved animal pose tracking. In Proceedings of the 27th ACM International Conference on Multimedia. pp. 1365-1374. ACM. 2019.
Kevin K. Singuenin-Dumittan Philip Loche, Ni Haoran, and Michele Ceriotti. Physics-inspired equivariant descriptors non-bonded interactions, 2023.
Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. Internutionaa conference on machine learning, pp. 448-456, pmils. 2015.
John Irwin and Brian K. Shoichet. ZINC free database of commercially available compounds for virtual screening, Journal of Chemical Information and Modeling, 1882)1177-182, December 2004, doi: 10.1021/ ci049714+. URL httt:s///aa...ort//...........
Sergey Ivanov Evgeny Buraev. Anonymous walk embeddings. Jennifer Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, 1111122001 Stockholmsmässan, Stockholm Sweden, 10-15 Jul 2018. PM... URI.http: http::/procsed:ii ngs. ml 280ss/v80/ ivanovl 8a. html
Andrew Jaegle, Sebastian Borgeaud, Jean- Saan-Bapiiste Alayrac, Carl Doersch, Catalin Ionescu, David Ding, Skanda Koppula, Danicl Zoran, Andrew Brock, Evan Shelhamer, Olivier Hénaff, Matthew M. Botvinick, Andrew Zisserman, Oriol Vinyals, and João Carreira. Perceiver io: A general architecture for structured inputs and outputs, 2022.
Charles R Johnson and Morris Newman. note on cospectral graphs Journal of Combinatorial Theory, Series B, 28(1):96-103, 1980. ISSN 0095-8956. doi: http:://boi.ogg/00.1000000 00t.:/00000000000995550000000000 8956(80)00088-11 URL https: / www. www.sciiiieeeee...c..mm. com/ science
Chaitanya Joshi, Cristian Bodnar, Simon Mathis, Taco Cohen, and Pietro Liò. the expressive power of geometric graph neural networks. In International Conference on Machine Learning, 2023.
Thomas N Kipf and Max Welling. Semi-supervised classification with graph convolutional networks. In Proceedings of the International Conference on Learning Representations (ICLR), 2017a.
Thomas N. Kipf and Max Welling. Semi-supervised classification with graph convolutional networks. In International Conference on Learning Representations (ICLR), 2017b.
Risi Kondor and Shubhendu Trivedi. On the generalization equivariance and convolution neural net- works to the action of compact groups. In Jennifer Dy and Andreas Krause (eds.). Proceedings of the 35th International Conference on Machine Learning. volume 80 of Proceedings of Machine Learning Re- search, pp. 2747-2755. PMLR, (0-15 Jul 2018. URL. kondor Ra. htm1
nga ml prrsss/v8C/
12

---

Published as a conference paper at ICLR 2024
Arthur Kosmala, Johannes Gasteiger, Nicholas Gao, and Stephan Günnemann. Ewald-based long-range mes- sage passing for molecular graphs, 2023.
David Peter Kovacs. llyes Batatia, Eszter Sara Arany. and Gabor Csanyi. Evaluation of the MACE Force Field Architecture: From Medicinal Chemistry to Materials Science. The Journal Chemical Physics, 159(4): 044118, July 2023. ISSN 1021-9606, 1689-7690. doi: 10. 015.53351111332...
Devin Kreuzer, Dominique Beaini, Will Hamilton, Vincent Létorreau, and Prudencio Tossou. Rethinking graph transformers with spectral attention, Advances Neural Information Processing Systems, 34:21618- 21629, 2021.
Y. LeCun, B. Boser, J. S. Denker, D. Henderson, R. E. Howard, W. Hubbard, and L. D. Jackel. Backpropagation applicd to handwritten zip code recognition. Neural Computation, 144):541-551, 1989. doi: 10.1 162/neco. 19.9.
L. Lin, J. Lu, L. Ying. R. Car, and W. E. Fast algorithm for extracting the diagonal of the inverse matrix with application the electronic structure analysis of metallic systems. Comm. Math. Sci., 7:755, 2009.
L. Lin, M. Chen, C. Yang. and L. He. Accelerating atomic orbital-based electronic structure calculation via pole expansion and selected inversion. J. Phys. Condens. Matter, 25:295501, 2013.
Haggai Maron, Heli Ben-Hamu Nadav Shamir, and Yaron Lipman. Invariant and equivariant graph networks.
arXiv preprint arXiv: www.:812..990 09902. 2018.
Haggai Maron, Heli Ben-Hamu, Hadar Serviansky, and Yaron Lipman. Provably powerful graph networks. In Proceedings of the 33rd International Conference on Neural Information Processing Systems, pp. 2156- 2167, 2019.
Christopher Morris, Martin Ritzert, Matthias Fey, William L Hamilton, Jan Eric Lenssen, Gaurav Rattan, and Martin Grohe. Weisfeiler and leman go neural: Higher-order graph neural networks. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pp. 4602-4609, 2019.
Marion Neumann, Roman Gareett, Christian Bauckhage, and Kristian Kersting. Propagation kernels: efficient graph kemelss from propagated information. Machine Learning, 102(():209-245, 2016.
Mathias Niepert, Mohamed Ahmed, and Konstantin Kutzkov. Learning convolutional neural networks for graphs. In International Conference on Machine Learning, pp. 2014-2023 2016.
Jigyasa Nigam, Michael J. Willatt, and Michele Ceriotti. Equivariant representations for molecular hamiltoni- ans and n-center atomic-scale properties. J. Chem. Phys., 156. 2022.
Behnam Parsacifard, Deb Sankar De, Anders S Christensen, Felix A Faber, Emir Kocer, Sandip De, Jörg Behler, O Anatole Lilienfeld, and Stefan Goedecker. An assessment of the structural resolution of various fingerprints commonly used in machine learning. Machine Learning: Science and Technology, 2(1):015018, apr 2021. doi: 10. 10.1088/2632-21133/bbb [pp212]
URL https:://////. / /x. doi, 10t.:11....1/1122222221112
Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E(n) equivariant graph neural networks, 2021.
T Schütt, M Gastegger, Tkatchenko, K-R Müller, and R Maurer. Unifying machine lcarning and quantum chemistry with deep neural network for molecular wavefunctions. Nat. Commun., 10(1):5024, November 2019.
Andrew W Senior. Richard Evans, John Jumper. Jeff Kirkpatrick, Laurent Sifre, Green, Chongli Qin, Augustin 7idek, Alexander W Nelson, Alex Bridgland, et Protein folding using distance maps and deep learning. bioRxiv, 2018.
Peter Shaw, Jakob Uszkoreit, and Ashish Vaswani. Self celf-attention with relative position representations, 2018.
Nino Shervashidze, SVN Vishwanathan, Tobias Petri, Kurt Mehlhom, and Karsten Borgwardt. Efficient graphlet kernels for large graph comparison. In Artificial Intelligence and Statistics, pp. 488-495 2009.
Martin Simonovsky and Nikos Komodakis. Dynamic edge-conditioned filters in convolutional neural networks on graphs. In Proceedings of the IEFE conference computer vision and pattern recognition, pp. 3693-
3702. 2017.
O. T. Unke, M. Bogojeski, M. Gastegger, M. Geiger, T. Smidt, and K.-.. Müller. SE (3) (3u-vvvrvr predic- tion of molecular wavefunctions and electronic densities, 2021a. 35th Conference on Neural Information Processing Systems (NeurIPS 2021).
13

---

Published as a conference paper at ICLR 2024
Oliver T. Unke and Markus Meuwly. PhysNet: A Neural Network for Predicting Energies, Forces, Dipole Moments, and Partial Charges. Journal of Chemical Theory and Computation, 15(6):3678 3693, 2019, ISSN 15499626. doi: 10. 10.102//ass.ott..b000..
Oliver I. Unke. Stefan Chmiela, Michael Gastegger, Kristof Schütt, Huziel E. Sauceda, and Klaus-Rober Müller. SpookyNet: Learning force fields with electromic degrees freedom and nonlocal effects. Nature Communications, 12(1), December 2021b.
doi: 10. 0000000/41447-00222755000 URL https://doi.
Oliver T Unke, Stefan Chmiela. Huziel E Sauceda, Michael Gastegger, Igor Poltavsky, Kristof T Schütt, Alexandre Tkatchenko, and Klaus-Robert Müller. Machine learning force fields. Chemical Reviews, 121
(16):10142-10186, 2021c.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances neural information processing systems, 30: 5998-6008, 2017.
Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pictro Lió, and Yoshua Ben- gio. Graph Attention Networks. International Conference on Learning Representations, 2018. URL
https:// net httee//ommmmmmdddppppe
Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Lio, and Yoshua Bengio.
Graph attention networks. In Proceedings of the International Conference on Learning Representations (ICLR), 2018.
S Vichy N Vishwanathan, Nicol N Schraudolph, Risi Kondor, and Karsten M Borgwardt. Graph kernels. Journal of Machine Learning Research, 101p)))12011114 2010.
Xiyuan Wang and Muhan Zhang. How powerful are spectral graph neural networks. In International Confer- ence on Machine Learning, pp. 23341-23366 PMLR, 2022.
Zonghan Wu, Shirui Pan, Fengwen Chen, Guodong Long, Chengqi Zhang, and Philip S Yu. A comprehensive survey graph neural networks. IEEE Transactions on Neural Networks and Learning Systems, 32(1): 4-24, 2020.
Zonghan Wu, Shirui Pan, Fengwen Chen, Guodong Long, Chengqi Zhang, and Philip S. Yu. A comprehensive survey graph neural networks. IEEE Transactions on Neural Networks and Learning Systems, 32(1): 4-24, 2021. doi: 10.1 10.10NNSS...22000022099
Zhang Xinyi and Lihui Chen. Capsule graph neural network. In International Conference on Learning Repre- sentations, 2019. URL BRLL/ttppe///ppezeviie net forte//oummmmm/mm wtwp:e/wwwwweeeeeeotemmmmmaaa
Keyulu Xu, Weihua Hu, Jure Leskovec, and Stefanie Jegelka. How powerful graph neural networks?
In International Conference on Learning Representations, 2019. URL htttpss:///peerevviee net/ Eorumpid=ryCsolA53
Pinar Yanardag and SVN Vishwanathan. Deep graph kernels. Proceedings of 21th ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, pp. 1365-1374. ACM, 2015.
Mingqi Yang, Wenjie Feng, Yanming Shen, and Bryan Hooi. Towards better graph representation learning with parameterized decomposition and filtering, 2023.
Chengxuan Ying, Tianle Cai, Shengjie Luo, Shuxin Zheng, Guolin Ke, Di He, Yanming Shen, and Tie- Yan Liu. Do transformers really perform bad for graph representation? arXiv preprint arXiv:2/06.05234, 2021.
Liwci Zhang, Berk Onat, Genevieve Dusson, Gautam Anand, Reinhard Maurer, Christoph Ortner, and James R Kermode. Equivariant analytical mapping of first principles hamiltonians to accurate and transferable materials models. upj Computational Materials, 8, 2022. doi: https://doi.rr//1 10. 1038/ $41524-022-00843- URLhttts:/ http://aaxiv.v.. crg/ wttp://////./..////e1
Li Zhu, Maximilian Amsler. Tobias Fuhrer, Bastian Schaefer, Somayeh Faraji, Samare Rostami, S. Alireza Chasemi. Ali Sadeghi. Migle Grauzinyte, Chris Wolverton, and Stefan Goedecker. A fingerprint based metric measuring similarities of crystalline structures. Journal of Chemical Physics, 144(3), I 2016. doi: 10. 49400000
Meiqi Zhu, Xiao Wang. Chuan Shi, Houye and Peng Cui. Interpreting and unifying graph neural networks with an optimization framework. In Proceedings of the Web Conference 2021, pp. 1215 1226,,26,2122 2021.
14

---

Published as a conference paper at ICLR 2024
A APPENDIX
A.1 NON-LOCALITY OF MENS AND COMBES-THOMAS THROREM
The degree of non-locality of the interaction can he precisely quantificd. If H. has finite interaction range, then the Combes- Thomass theorem Combes & Thomas (1773) implies that [==I H S Ce where
T = dist(2, z, a (H)). da the length of the shortest path from node node j. and C, c may depend on the norm of but can be chosen uniformly provided [H] (operator max-norm) has an upper bound. Our normalisation of weights in Appendix [A.4 is an even stronger requirement and hence guarantees this.
Since we have taken H self-adjoint with real spectrum, an estimate for 7s is Imag(z). As a result, if we constrain the poles in the parameterization of fo to be Imag(za) 7, then the resulting matrix function will satisfy
(15) Therefore, the degree of locality can be controlled by constraining Imag(ze) zx) to be some fixed value. While equation [5] only gives an upper bound. it is not difficult to see that it can in fact be attained for suitable choice of matrix function. In practice, however. fo can be seen as approximating an unknown target with learnable weights u2g- If f is analytic in neighborhood of a (H)) then [[) S Ce- where measures how smooth f is (distance of poles of f to (H)). For our parameterization fe we therefore expect the same behavior 7 = 7 in the pre-asymptotic regime of moderate diy but transition to the asymptotic rate equation [5 for large dij. An in-depth analysis of such effects go beyond the scope of this work.
[H)) Ce
A.2 GEOMETRIC EXPRESSIVENESS AND RELATIONSHIP TO INFINITE LAYERS OF MPNNS
The WL-test quantifies a network's 's ability to distinguish non-isomorphic graphs. Its extension to graphs embed- ded in vector spaces is known the geometric WL-test (Joshi et al. 2023). The expressivity in the following discussion adheres these definitions. relationship exists between equivariant MPNNs with linear updates and infinite layers and a one-layer MFN due to the matrix function's 's product structure. The features constructed by a one-layer MFN with a single ma- trix channel and two-body matrix entries closely resemble the features of a two-body equivariant MPNN with infinite layers and lincar updates. When matrix entries incorporate features beyond two-body, the one-layer MFN becomes more expressive than its body feature counterpart.
the context of graphs in practical relevance is high. A clear
Here, we give an intuition of the relationship between the two approaches. Any analytic matrix function f admits formal power series expansion, valid in its radius of analyticity,
fII)
Σ k=
Ci
H
(16)
where e& the complex (or real) coefficients of the expansion. Let us look at the diagonal elements of each of these powers that we will extract the next doe features.
(H)
Σ
Hoj.ctmHyy..cm.
(H)
Σ Σ
(17)
(EN(()
(EN( REN())
The sum over (i)) is forced due to the sparisity 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 two convolutions, all the way infinity. Each coefficient of the matrix function that learnable can act like weights the linear update the analogous layer. Because the power series expansion is infinite, one sees a clear analogy between the two.
A.3 INTERPRETABILITY OF THE MFN OPERATORS
In the case of the Euclidean group. the matrices learned in MFNs have the same symmetries as Euclidean operators. Euclidean operators play a crucial role in quantum mechanics. They are defined as self-adjoint operators (in the space of square-integrable functions) that are equivariant the action of the group of rotations, translations, reflections, and permutations. When these operators are expended on the basis the irreducible representations of the rotation group, they exhibit block structure (see Appendix A.6) in which cach entry is labeled with representations (as, ps, pp, dd, ...).
The most central of these operators is the Hamiltonian operator. The energy of quantum system is related to the trace a specific matrix function of the Hamiltonian,
(18) The Hamiltonian is usually computed as a fixed point a self colf-cooniitent loop that minimizes the energy. This loop introduces many-body non-local effects. This motivates us to use many-body functions to parameterize our matrices and to learn the fixed point directly via learning a matrix function of many matrices. This is in opposition tight-binding methods that usually construct a single low-body Hamiltonian with physics-inspired functional forms but require several self-consistent iterations converge.
E = TH)(H)
15

---

Published as a conference paper at ICLR 2024
Φ
H
Pi
m1
P
sp
PPP
7/22
Figure 4: Block structure of a Euclidean Operator, H, learnt in the MFNs. Each entry in corre- sponds to a different product of representations of the group of rotations (ss, sp, PP,....
A.4 MATRIX NORMALIZATION
Batch normalization (loffe & Szegedy 2013) and layer normalization al. 2016 are two techniques widely used in deep learning to stabilize trainin. improve convergence, and provide some form of regulariza- tion.
In the context of MFNs, similar normalization strategy is needed to ensure stable training and improved convergence. Instead of normalizing the features directly conventional normalization layers do, normalize matrices with batch dimension batch, channels, n, rd]. where n is the size of
the eigenvalues of set the matrix. The aim to adjust their mean and variance to 0 and 1, respectively, much like standardization in traditional batch or layer normalization. For the batch matrix norm, the mean is taken across the batch dimension, and for the layer matrix norm, the normalization is taken across the channel dimension. The mean and the variance of the eigenvalues are computed using the following formulas,
B(A) tr(H) Var(A)
tr(H?)
tr(H)²
(19)
f1 1)
This normalization of eigenvalues that the spectral properties of the graph used for representation in MFN are effectively standardized, contributing to better training stability and convergence.
A.5 ALTERNATIVE UPDATES FOR MFN
Many updates can defined in order to extract non-local features from the matrix function.
1. The dense update utilizes the entire matrix, including all off-diagonal terms, to update the next matrix function,
A
(++1)
v- icm
Σ w
(H + f6 (t-1)) (H H).m
(20)
scm
This update can not be performed using selected inversion techniques, it can offer additional ex- pressiveness guarantees.
2. The sparse update: uses only the parts of the matrix corresponding to the connected nodes to update the nodes and edges of the graph obtain the matrix function in the next layer,
(1+1)
(i() Σ
(())
(H -US--"(( (eA())
(21)
A.6 EQUIVARIANCE
Proof of equivariance of continuous matrix functions We will first prove the invariance of the spec- trum of operator. Let H C= be self-adjoint matrix. H admits eigendecomposstion, = T The cigenvalues A the roots of the characteristic polynomial of H,
det(H >A.) = 0, A
(22)
Let G reductive Lie group. Then any finite- dimen-dammniional representation G. can viewed represen- tation of the maximal compact subgroup of the complexification of G. by the Weyl unitary trick. Therefore, they can analytically continued to unitary representations. Let G + Caxn he a representation of G taken unitary, p(g))(g)" Tn, Vgg G. Let H1 and H2 be two matrices related by action of G, then Hi p(g) Haplg)" for some g Let As and A2 be their respective spectrums. For all À E A
det(Hi >Ia) = det (plg) Haplg) - Ala) det(p(g) - Xxx)p(g)) (eet(p())p(g)) (det(H2 Adn) det( Al.)
(23)
(24)
16

---

Published as a conference paper at ICLR 2024
Therefore the eigenvalues of H2 also eigenvalues of H2 and A1 A2-
We will now prove that if E U1 is cigenvector of H₂ with eigenvalue À then p(g)u eigenvector of H2. For all u E U1
Hz(p(g)u) = p(g)H (p(g))) p(9)
(25)
= >(p(g)u)
(26)
Now let's prove equivariance of matrix function. Let f - continuous function, then matrix function is defined as,
=(H)= = Uf
(27)
Therefore for any g E G.
S(H g) = H((()) = P(S)US(AJUU" p(g)" (()
(28)
(29)
Equivariance of the resolvent
The pole expansion yields straightforward proof of (a(H) equivariance:
[a(H o g)
Σ w. (z.I pHp') Σ p.,I - H) ">''p"-- p.
(30)
For general continuous function f. the analogous result follows by density.
A.7 MULTIVARIATE MATRIX FUNCTION
As MFNs extract features from multiple operators, it is useful to extend the univariate matrix function in
Equation to multivariate matrix function. The space of multivariate functions F(M(GG) isomorphic the closure of the span of the tensor product of functions single operators F(H(G)) 88 We the number of variables n, the correlation order of the matrix function. The resolvent expansion can be generalized to the multivariate matrix function, using matrix products of resolvents,
H(G)) is
f(H,..... H.,) (2mr)" /// -fanyx--
f(21,..., 2n)
da dz.
(31)
x...x. (1) (2, (2,I-H.)
Multivariate matrix functions create higher order features. one-layer MFN with a matrix function corre- lation order n is as expressive as a (n + 1)-body equivariant MPNN with infinite layers and linear updates. The space F(H(G)) is of very high dimension and it is possible to use standard compression techniques to find more efficient approximations such tensor decomposition (Darby et al., 2023). The use of linear combinations of matrices Equation approximates some of this space. We leave further exploration of these kinds network future work.
A.8 RUNTIME COMPARISON
The current MFN implementation a prototype implementation and has cubic scaling. However, fast inver- sion methods are well established, and we explain in Section[4.2]how this scaling can 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 known, they require significant effort to implement and integrate them into machine learning librarics like Pytorch due the need for specialized CUDA kernels. We intend make this effort and release open-source code 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.
For the sake of completeness, we show run-time comparison on the cumulene examples between MACE and the MFN. We time the models energy and forces evaluation of the MFN model is MACE layer with the same hyperparameters as MACE model. We use one layer MFN model with poles and 16 matrix channels. The MACE model has channels, maximal angular resolution of liaa = 3 and message passing equivariance L = 1.
A100 GPU, The graph layer construction
A.9 DETAILS OF NUMERICAL EXPERIMENTS
A.9.1 CUMULENES
Dataset The cumulene dataset is designed to test the expressivity of graph neural networks and their non- local capabilities. The task is to regress energy and forces from 3D molecular configurations. The ground truth energy is obtained using the ORCA quantum chemistry code at the density functional theory level of accuracy using the wB97X-D3 functional, the def2-TZVP basis set, and very tight SCF convergence criteria.
17

---

Published as a conference paper at ICLR 2024
Table 4: Runtime comparison of a energy and force call of a MACE model and MFN model on 20 atoms cumulene. The MFN models constructs equivariant matrices up to L = 1.
Method
MACE
MFNMACE
Time (ms/atom)
0.7
2.8
ccc
Hydragen
Ahch
Tau
Carbon
1gg
Figure 5: Angles and distances that define cumulene graph used test expressivity in Figure The carbon-hydrogen (rCH), first carbon-carbon (1))) and remaining carbon-carbon distances (r())) set to 1.086 A, 1.315 Å and 1.279 Å respectively. The angle between the hydrogen- carbon-hydrogen (0нCH) is fixed at 118.71 degrees and the dihedral angle e depends on the exper- iment as detailed in the main text.
Length and angle scans The configurations for the length and angle scans consist of chain chain-like graphs as Figure[5] The length scans range from to 30 atoms fixed at angle of degrees. Additionally, scan the angle hetween the two terminating hydrogen groups for cumulene with twelve carbon At increments of degrees, we sean from 0 to 90 degrees. By symmetry, this covers the entire range unique configurations from to 360 degrees. For the expressivity scan, all internal coordinates (see Figure [5) that uniquely define the cumulene are kept constant and set to the geometry-optimise 30 cumulenes. The carbon-hydroged 1.086 A, 1.315 Å and 1.279 Á respectively (see Figure 5 For the length 12 cumulene, the distance between the most distant carbons 14.1A 14. Thus, it becomes impossible for purely local equivariant MPNN, with cutoff 3 and two message- passing layers, to differentiate between two graphs at different angles.
(ne 12).
configuration of the length
first carbon-carbon, and remaining carbon-carbon distances are
Note that the relative energies offer obtained by subtracting individual atom energies. These are -6.3eV and -1036.1eV for each hydrogen and carbon atoms respectively.
Guaranteed Non-local dataset The GNL dataset tests how well a model can extrapolate an unseen number atoms. Furthermore, configurations are rattled to see how well the model can capture both local and non-local effects. Configurations are generated starting from the relaxed geometries of cumulenes with length 0-20. Relaxations are carried out with the smaller 6-3/G basis set. The cumulenes rotated by equally spaced scan with increment degrees, starting from random angle, and the positions are subsequently randomly perturbed by Gaussian noise with standard deviation 0.01 Å. The training and validation set contain cumulenes of length 0-10 and 13-14, with 10 and 3 samples, respectively. The test set contains in ni-domain configurations with two samples for each of the lengths present in the training set, Furthermore, contains configurations with 11-12 and 15-16 carbon atoms, labeled oft-on-dommin. In total, the train, validation, and test contain 200, 50 and 170 configurations.
subsequently
MACE and MFN Model Both the local MACE MFN are trained on same graphs, which means that the cutoff distance is fixed at 3A, including information from the nearest neighbor. For the MFN model, we use the architecture described in Section 4.3 Both models are trained with two layers, allowing the local model to distinguish changes that are separated by 12A. The MFN is trained using 16 matrix channels and 16 poles. For the matrix construction step, use MACE Batatia al. (2022h) layer with correlation order lmx = 3, = and 128 channels to generate node features hç. We use the diagonal update [ to update the node features of MACE and reiterate. The readout at the first layer is linear and at the second it is one layer MLP with 16 hidden dimensions.
Spookynet Model The Spookynet architecture was trained on the same dataset as the MFN model. We use one-layer model with 5.5 A cutoff reproduce the receptive field of MACE. We the global attention and electrostatic interactions. The model has 128 channels.
Training All models underwent initial training on the GNL dataset before transfer learning was applied to the relevant length and angle scan data sets. The saved model corresponds to the epoch that exhibits the minimum loss values. Details settings such learning rate and epoch count are disclosed in Table 5 Training incorporated both energy and forces, with adjustable weights for each observable. In particular, for
18

---

Published as a conference paper at ICLR 2024
length and angle scans, additional 100 epochs with force weight were used after initial training, ensuring the depiction of the minimal possible energy error Figure [3
the lowest-loss model saved.
Table 5: Model training parameters. For the matrix functions the number of poles (np) and matrix
channels (c) are indicated.
dataset
model
epochs
Ir
Eweight
Fwvight
Playets
Timxs
other
Length Scan
MFN (L=0)
1240
le-2
1000
100
2
3
11=16,c=16
MFN (I==1)
5628
le-5
1000
100
2
3
n,=16,c==16
MACE
5288
0.005
1000
100
2
3
Angle Scan
MFN (L=0)
1240
le-2
1000
100
2
3
n==16,c=16
MFN (L=1)
656
lc-S
1000
100
2
3
11=16,,=16
MACE
954
0.005
1000
100
2
3
GNL
MFN (L=1)
4958
le-2
100
3
ny=16,c=16
MACE
1260
1e-2
100
2
3
Spookynet
3500
le-4
0.10
0.90
5.5
attention
A.9.2 ZINC
Dataset The ZINC dataset [Irwin & Shoichet 2004) contains 12,000 small 2D molecular graphs with an average of 23 nodes with information on the node attributes and the edge attributes. The task is to regress the constrained solubility log P- -A #cycle where log P is the octanol-water partition coefficients, SA is the synthetic accessibility score, and #tyccle is the number of long cycles. For each molecular graph, the node features are the species of heavy elements.
Model The model made up of layers. Only the first two layers are MFN layers (in order save on the number of parameters). The initial edge features ciy obtained by concatenating PDF Yang al. 2023 descriptors used for ZINC the input edge features the data set. each layer f, the initial nodes features, (( are computed using a convolutional GNN with 140 channels. For the first two layers, matrix then formed using multilayer pereeptron (MLP) of size [560, 16) with the GELU activation function and batch norm,
(32) We compute the matrix function, using 16 poles and 16 channels. We extract the diagonal elements to update the node features, and use the off diagonal elements to update the edge features by passing them into a residual 2 layers MIP of size [512, 140, 140.. We used average pooling followed by linear layer for the readout.
Hy = MLP(M!). (())
Training Models were trained with AdamW, with default parameters of 31 0.9, 8g 0.999, and 10- and weight decay of 5e 5. We used learning rate of 0.001 and batch size of 64. The learning rate was reduced using an on-plateau scheduler.
The baseline models used for the ZINC dataset comparisons include: GCN (Kipf & Welling 20176), GAT (Velčkkvvć et al., 2018), MPNN (Gilmer et al. 2017), GT (Dwivedi et al. 2020). SAN (Kreuzer
Baselines
et al. 2021) Graphormer (Ying et al. 202T). PDF (Ynng et al. 2023
A.9.3 TUDATASETS
Dataset We train subset of the TUDatasets including the MUTAG. ENZYMES, PTC-MR, PROTEINS and IMDB-B subsets. The MUTAG, ENZYMES, PTC-MR, and PROTEINS datasets are molecular datasets that contain node-level information on the molccules. The IMDB-B a social network datasets.
The number layers for each subset is in Table 6 (in order to save on the number of parameters). The number of hidden channels, nigpprs for each model is given in 6 The initial edge features ef are obtained by concatenating PDF Yang et al. (023) descriptors used for each subset and the input edge features of the data set. In each layer the initial node features, h,!), computed using convolutional GNN. For the first two layers, the matrix is then formed using a multilayer perceptron (MLP) size [4 x Vireamees, 16) with the activation function GELU a batch norm,
Model
We use MFN layers only for the first 3 layers
Hay = MLP(M 16
(33)
We compute the matrix function, using 16 poles and 16 channels. We extract the diagonal elements to update the node features, and the off diagonal elements update the edge features by passing them into residual
19

---

Published as a conference paper at ICLR 2024
2 layers MLP of size [512, Meaures,, Tefeature]. We used average pooling followed by linear layer for the readout.
Training Models were trained with AdamW, with default parameters of 31 = 0.9, 82 0,999, and = 10 and a weight decay of 5e 5. We used learning rate of 0.001 and batch size of 64. The learning rate reduced using an on-plateau scheduler.
Table 6: Model training parameters. For the matrix functions the number of poles (n,) and matrix
channels (c) are indicated
dataset
epochs
Ir
Players
Tefamures
other
MUTAG
301
le-3
6
256
n2=16,c=16
ENZYMES
le-3
256
n,=16,0=16
PTC-MR
151
le-3
6
128
n,=16,c=16
PROTEINS
le-3
128
1p=16.c=16
IMDB-B
301
le-3
3
256
7p=16,c=16
Baseline
The baseline models used for the TU datasets comparisons include: GK (Shervashidze al.
2009
RW /Vishwanathan al. 2010), PK (Neumann al.) (2016), AWE (Ivanov Burnaev 2018)) PSCN (Nieper)
al. 2016), ECC (Simonovsky
Komodakis 2017), DGK (Yanardag Vishwanathan 2015).
Caps-
GNN (invi & Chen 2019), GIN (Xu et al. 2219). k-GNN (Morris et al. [2019) IGN Maron ct al
2018
PPGNN (Maron et al. 2019), GCN (de Haan et al. 2020 GraphSage (Hamilton et al. 2017
PDF
Yang
2023
20


Reviews:
[{"review_id": "8vurddXaP7", "reviewer": "Reviewer_qFqk", "summary": "This paper proposes MFN, a graph neural network where the intermediate representation takes a matrix form. The matrix form makes the neural network equivariant for a group action where the representation is unitary. An efficient computation based on the power expansion of the matrix is also proposed. The experiments on the molecule datasets show the competitive performance of MFN.", "questions": "In experiments:\n- What is L in the caption of Fig 2?\n- It seems H is constructed with Kronecker products as in Eq. (21). What is the actual size of H?", "limitations": "", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 1, "contribution": 3, "strengths": "Strong empirical results.\n\nMatrix variate formulation and its efficient approximation by the resolvent expansion are novel.", "weaknesses": "There is room for improvement in the presentation. I found several incomplete descriptions. For example,\n- Given a matrix $A$, $A^*$ is not defined (I guess it's adjoint of A)\n- $f_\\theta(H)$ is defined as \"scalar to scalar,\" but it seems \"matrix to matrix\". \n- In Eq. (10), since $H$ is a matrix, it's not obvious how $g$ acts on $H$. I guess H is the function of $\\sigma$ and $g$ acts on it, but it's not clearly described. \n- I couldn't find the proof of Eq. (10).\nAlso, see the \"Questions\" area below. \n\nIt 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. A clear comparison is needed. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "y9ziJLJwjM", "reviewer": "Reviewer_2GgR", "summary": "This paper introduces a novel architecture for processing graph data. The goal of the proposed method is to model non-local interactions between nodes in graphs without the shortcomings of alternate approaches. Namely, over-smoothing and excessive compute requirements for message-passing GNNs and the need for hand-designed features with spectral GNNs.\n\nMatrix function neural networks learn equivariant representations over graphs. The approach consists of three stages: 1) building a matrix that corresponds to a group-equivariant matrix operator 2) applying a learned matrix function to this matrix 3) updating the matrix. To my understanding, the major contributions appear in the second part. The authors make use of resolvent calculus to determine an efficient parameterization of the matrix function. Given that the underlying matrix is sparse, this parameterization admits an efficient (potentially linear) algorithm for computing the function.\n\nThe method is evaluated on three different graph learning datasets, where it generally outperforms the baseline methods (or performs at least as well).", "questions": "One area of the paper that I feel I couldn't completely understand was on 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. On the surface, this seems to align with spectral GNNs that lean on hand-designed features for effective learning. However, the authors point out that this can be approximated well in some cases. 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?\n\n\nDo we need to ensure that the learnable weights and poles satisfy any constraints for the resolvent calculus definition (Section 4.2)? From reading the explanation in the paper, I was under the impression that the choice of weights and poles determine the curve implicitly via the quadrature rule. Is it possible that we could learn weights and poles that lead to invalid curves, or otherwise undesirable behaviour?\n\nAs I mentioned above, the MFN models used in experiments always contain only a small number of layers. Was this a computational consideration?\n\nMinor comments:\n\n- There are a couple of symbols that are not defined in the main text. These can be mostly be inferred from context, but should ideally be defined.\n    - Above equation (1), the domain and codomain of $s$ should be defined.\n    - I could not find a definition of $\\mathfrak{Z} z$ and related symbols used in Section 4.3.\n- You cite transformers as having a problematic quadratic scaling. While this is true in the original formulation, it is quite common nowadays to use methods with far more efficient scaling. For example, PerceiverIO can reduce the computational cost to linear.\n- Shaw et al. 2018 are cited for positional encodings, but these are also used in the original Vaswani et al. transformer paper.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "To my knowledge, this is an original and novel approach to designing GNNs. The authors describe theoretical and practical differences between their approach and existing methods. The proposed approach pulls on learnings from several other areas of study (resolvent calculus, electronic structure methods, and more) to design an expressive and efficient class of graph neural networks. The resulting method is elegant and performs well in the empirical evaluation.\n\nOverall, I found the paper to be clear and well-written. I am not familiar with resolvent calculus or some of the other technical background introduced in section 4, but I was able to understand and follow the details. Furthermore, full details of the experiments are provided in the supplementary.\n\nI feel that the contributions in this work are significant. GNNs are a busy area of research and it is no small feat to propose a novel and practical architecture. Empirically, the proposed method performs well. While the margin of improvement is often narrow, MFNs consistently achieve performance at least on par with leading methods in comparison.", "weaknesses": "It is unclear whether MFNs can be used as a drop-in replacement for GNNs. In two of the experiments, the models used are not end-to-end matrix function networks, due to model size considerations. Instead, a small number of MFN layers are used with more standard layers following. In all experiments, initial edge features are also extracted from other architectures --- I haven't quite understood the implication of this in terms of comparison to baselines.\n\nThe method is fairly complex with several design decisions. For example, the choice of basis, choice of update method, number of learnable weights and poles, etc. Moreover, making the method computationally tractable requires some techniques that are unfamiliar to the learning community at large. Therefore, I expect that the method may be adopted quite slowly by the community. This could be alleviated with an open-source implementation.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "87wgrmi0xN", "reviewer": "Reviewer_Qi1r", "summary": "GNNs need to balance the contributions from local (nearby nodes) data and more non-local data. The authors here present an approach based on spectral graph convolutional methods which can capture local and non-local features through a parameterization via resolvent expansions. The authors propose various means of implementing their method to both improve performance and speed up the technique. Their results are backed by experiments on a few datasets which show relatively good performance of the method.", "questions": "Other questions/concerns:\n- When reporting parameter counts, the authors have some parameters which are complex numbers. Are these complex numbers counted as two parameters per number or one? A fair comparison should count each complex number as two parameters since there is a trainable real and complex part.\n- 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? If so, what are the practical limitations that arise from this implementation?\n- What are the $\\phi_m$ basis functions for pure graph inputs? \n- If the $\\phi_m$ are being learned, did you include this in the parameter count?\n- When you say \"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? That would be unfair for comparison if so.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 1, "contribution": 3, "strengths": "The most important strength is that the paper presents a method which performs well on numerical experiments. Though I have some concerns about the experiments, if the results are confirmed and these concerns are appropriately addressed, I think this method should be documented and tested further.\n\nSeparate from the experiments, the proposed method is an extension of spectral convolutional approaches which have enjoyed great success. The particular strategy that the authors propose, via resolvent expansions, seems like a decent approach to take when including long range interactions. It is tied together relatively nicely with geometric data as well. Implementation can be tricky in certain parts (e.g. implementing complex numbers) but the architecture and overall idea is relatively sound.", "weaknesses": "The method is laid out in section 4 and to be honest, I had quite a bit of trouble following along. A number of notational concerns are pointed out below. But let me just start at the beginning and make sure I have the right picture. As I understand, the authors are choosing $b$ basis operators which are functions taking inputs consisting of a set of $\\mathcal{X} \\times V$ pairs where $V$ is some space for the hidden state. Then, they apply linear maps on top of these outputs from the bases which respect the graph permutation equivariance, hence a $bn \\times bn$ matrix. I still am not sure this is the right interpretation, but the authors need to state this more formally and cleanly. They should specify what the functions $\\phi_m$ do. I.e. specify input and output space. Indexing of eq. (7) is confusing. I believe $ij$ refers to the node indexing, but the matrix is $bn \\times bn$ so $H_{cij, m_1, m_2}$ is not consistent with how the entries of the matrix work. Also, eq. (7) seems to indicate the entries are outputs of fixed basis functions, but these can be learnable functions as well? The authors should also give examples of $\\phi_m$ as promised (is (21) an example of such?). I also believe that eq. (7) should be indexed over $k$ in the intersection of the neighbors, not the union. This is just a sampling of the confusions that arise and eventually required me to make decisions as to what is actually happening.\n\n\n\n\nMoving on, I have concerns about the experiments. First, there is no code shared so I cannot verify numbers. What is shared in the zip file appears to be a dataset and some details on how to train a model for that dataset. Sharing code is obviously not a requirement, but it means I have to judge the result solely on what is written. Here, the methods of comparison to other results are potentially in question. For example, it appears complex numbers are being counted as a single parameter per number whereas it should be two per number (real and complex parts are both variables); see question below asking to clarify this. Runtime comparisons are also not provided, and this can be a concern here, since the method requires spectral decompositions (or approximations thereof). Finally, see questions below, but I am not sure how certain features are being extracted. E.g., the authors say that they take descriptors from the PDF model in the appendix for the ZINC data, but this would be unfair if they take features from a trained model on the same dataset. Nonetheless, if these issues are resolved, results look good as stated and if confirmed, this would add support to the method.\n\nI simply could not follow section 4.3. What is the Combe-Thomas theorem? There is no citation for this. 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? Later on, the geometric expressiveness statements are stated without proof and missing too many details for me to follow. I have many questions. 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?\n\nNotation comments:\n- Eq (3): $\\mathcal{G}$ in the notation was defined as a specific graph; however, here it is used to denote the space of graphs.\n- Eq (4): $\\Phi \\odot g$ is confusing. Why is there a representation for the output space but not the input space? \n- Below Eq (8): $\\rho$ is bold in the equation but not bold there. Also, $\\rho$ is unitary in general, not orthogonal as far as I understand. If it is orthogonal, then why is there a complex conjugate transpose instead of just a transpose?\n- ... and many more. I would ask the authors to go through the paper and significantly revise and clean things up. \n\n\nAltogether, I cannot recommend acceptance for this paper given the writing quality and concerns laid out above. Nonetheless, I welcome the authors' responses to these points.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (spotlight)