Equivariant Matrix Function Neural Networks
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 graphs such as large conjugated molecules, and social networks
due to oversmoothing and oversquashing. Although Spectral GNNs and traditional
neural networks such as recurrent neural networks and transformers mitigate
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
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 stateof-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, paving the way to new state-of-the-art force
fields.
Ilyes Batatia, Lars L. Schaaf, Huajie Chen, Gábor Csányi, Christoph Ortner, Felix A. Faber
# Equivariant Matrix Function Neural Networks

###### 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, paving the way to new state-of-the-art force fields. The code and the datasets will be made public.

## 1 Introduction

Graph Neural Networks (GNNs) have proven to 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) (Gilmer et al., 2017; Battaglia et al., 2018; Kipf & Welling, 2017a; Velickovic et al., 2018; Wu et al., 2020; Senior et al., 2018; Hu et al., 2019; Batzner et al., 2022) and higher-order equivariant MPNNs (Batatia et al., 2022b).

MPNNs struggle to model **non-local** interactions effectively due to computational constraints and over-smoothing (Di Giovanni et al., 2023). Spectral Graph Neural Networks attempt to address the limitation 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 inherent spectral features. Many spectral GNNs apply polynomial or rational filters (Bianchi et al., 2021; Gasteiger et al., 2018; Wang & Zhang, 2022; He et al., 2021; Defferrard et al., 2016; Zhu et al., 2021; Kreuzer et al., 2021) to eigenvalues of graph structures, reaching state-of-the-art accuracy on pure graphs tasks. However, these methods often exhibit rigid architectures that require extensive feature engineering, potentially limiting their adaptability to various types of graphs. Moreover, they have been restricted to non-geometric graphs, making them unsuitable for molecules and materials.

Traditional neural network architectures, such as recurrent neural networks (Elman, 1990; Hochreiter & Schmidhuber, 1997; Cho et al., 2014; Graves, 2013) and transformers (Vaswani et al., 2017) also face challenges when modeling non-local interactions. While transformers can capture non-local dependencies through their self-attention mechanisms, they come at a significant computational cost due to their quadratic complexity with respect to the input sequence length. Furthermore, transformers lack inherent structural relationships or positional information within input data, necessitating the use of additional techniques, such as positional encodings (Shaw et al., 2018).

In chemistry and material science tasks, some models incorporate **long-range** interactions through electrostatics (Grisafi & Ceriotti, 2019; Behler, 2021; Unke & Meuwly, 2019), dispersion, or reciprocal space (Gao & Remsing, 2022; Huguenin-Dumittan et al., 2023; Kosmala et al., 2023). However, no existing architecture effectively address **non-local** interactions, where effects can propagate
over extensive distances through electronic delocalization, spin coupling, or other many-body non-local mechanisms. This is particularly problematic in systems such as large conjugated molecules, amorphous materials, or metals.

Consequently, there is a need for new neural network architectures that can efficiently and accurately model complex non-local many-body interactions, while addressing the limitations of current approaches. We propose **Matrix Function Networks** (MFN) as a possible solution to this challenge. Concretely, we make the following contributions.

* We introduce Matrix Function Networks (MFNs), a new graph neural network architecture able to model non-local interactions in a structured, systematic way and with potentially linear scaling with the size of the input.
* We introduce the resolvent expansion as a convenient and efficient mechanism to learn a general matrix function.
* We demonstrate the ability of our architecture to learn non-local interactions on a dataset of challenging non-local quantum systems.
* We show that MFNs achieve state-of-the-art performance on ZINC and TU graph datasets.

## 2 Related Work

Overlap matrix fingerprints (Zhu et al., 2016) introduced overlap matrix fingerprints (OMFPs), a vector of spectral features of an atomic environment or more generally a point cloud. Given a point cloud, an overlap operator (identity projected on an atomic orbital basis) is constructed, and its ordered eigenvalues (or other invariants) are taken as the features of that point cloud. Although a theoretical understanding of OMFPs is still lacking, computational experiments have shown excellent properties as a distance measure (Zhu et al., 2016; Parsaeifard et al., 2021).

Spectral Graph Neural NetworksSpectral 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 (Deferrard et al., 2016a) approximates the polynomial function of the Laplacian using Chebychev expansions, GPRGNN (Chien et al., 2021) directly fits coefficients of a fixed polynomial, while ARMA (Bianchi et al., 2021) uses rational filters.

Equivariant Neural NetworksEquivariant neural networks are the general class of neural networks that respect certain group symmetries (Bronstein et al., 2021). Noteably, convolutional neural networks (CNNs) (LeCun et al., 1989) are equivariant to translations, while \(G\)-convolutions (Cohen and Welling, 2016; Cohen et al., 2018; Kondor and Trivedi, 2018) generalized CNNs to equivariance of compact groups. Lately, equivariant message passing neural networks (Anderson et al., 2019; Satorras et al., 2021; Brandstetter et al., 2022; Batzner et al., 2022; Batatia et al., 2022b;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. (2022a, 2023).

Hamiltonian LearningA natural application of equivariant neural network architectures is machine 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 receiving increasing interest because of the potential extension in accessible observables over purely mechanistic models. The recent works of Nigam et al. (2022) and Zhang et al. (2022) introduce such parameterizations in terms of a modified equivariant Atomic Cluster Expansion (ACE) (Drautz, 2020), a precursor to the architecture we employ in the present work. Alternative approaches include (Hegde and Bowen, 2017; Schutt et al., 2019; Unke et al., 2021a; Gu et al., 2023).

## 3 Background

### Spectral Graph Neural Networks

We briefly review spectral graph neural networks and explain their limitations that our work will overcome. Consider a graph \(\mathcal{G}=(X,\mathcal{E})\) with node set \(X\) and edge set \(\mathcal{E}\). A graph defined purely by
its topology (connectivity) is called a pure graph. Let \(n\) denote the number of nodes in the graph, and let \(\mathbf{A}\in\mathbb{R}^{n\times n}\) be its adjacency matrix. Define a vector of ones as \(\mathbf{1}_{n}\in\mathbb{R}^{n}\). The degree matrix of the graph is \(\mathbf{D}=\text{diag}(\mathbf{A}\mathbf{1}_{n})\), and the Laplacian matrix is \(\mathbf{L}=\mathbf{D}-\mathbf{A}\). The Laplacian is a symmetric positive semidefinite matrix and admits a spectral decomposition, \(\mathbf{L}=\mathbf{U}\mathbf{\Lambda}\mathbf{U}^{T}\), where \(\mathbf{U}\) is an orthogonal matrix of eigenvectors and \(\mathbf{\Lambda}\) is a diagonal matrix of eigenvalues.

A popular approach to learning functions on graphs is to use 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\) be a function of a graph \(\mathcal{G}\) and \(t\) be a convolutional filter. SGCNs take advantage of the spectral decomposition of the Laplacian matrix of graph \(\mathcal{G}\) to compute the convolution of \(s\) and \(t\):

\[s^{\prime}=t\star_{\mathcal{G}}s=\mathbf{U}((\mathbf{U}^{T}t)\odot(\mathbf{U} ^{T}s))=f_{\theta}(\mathbf{L})s, \tag{1}\]

where \(f_{\theta}\) is a matrix function of the Laplacian matrix \(\mathbf{L}\) of the graph \(\mathcal{G}\) and \(\odot\) the Hadamard product. Various works have proposed different matrix functions \(f_{\theta}\), such as polynomial functions (Defferrard et al., 2016), rational functions (Bianchi et al., 2021), or neural networks (Wu et al., 2021). In a recent work, Yang et al. (2023) proposed to extend filter functions to matrix functions of other **fixed** matrix representations of the graph. Let \(\mathcal{H}(\mathcal{G})\) denote the space of all graph operators on graph \(\mathcal{G}\), i.e., the space of linear self-adjoint and permutation equivariant operators from nodes to nodes. For any **fixed**\(\mathbf{H}\in\mathcal{H}(\mathcal{G})\), generalized spectral GNNs determine filters based on \(\mathbf{H}\) with a matrix function:

\[s^{\prime}=f_{\theta}(\mathbf{H})s, \tag{2}\]

where \(f_{\theta}\) is a matrix function of \(\mathbf{H}\). Two non-isomorphic graphs can share the same spectrum of their Laplacian operators, while they have different spectra for other graph operators (Johnson and Newman, 1980). Therefore, the use of other graph operators as a 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 a **fixed** graph operator \(\mathbf{H}\) limits the expressivity of the model. Choosing the most expressive graph operator requires problem-dependent feature engineering.
* **Lie-group symmetries:** The approaches proposed so far have been restricted to **pure** graphs and in particular do not use additional symmetries for graphs embedded in a vector space. For example, graphs embedded in \(\mathbb{R}^{3}\) often lead to \(E(3)\)-equivariant learning tasks.

To overcome these limitations, we propose MFNs in Section 4, which allow parameterization of both the graph operator \(\mathbf{H}\) and the matrix function \(f_{\theta}\), and can be formulated to preserve the equivariance under all known group actions. Since a learnable operator \(\mathbf{H}\) prevents the precomputation of diagonalization of \(\mathbf{H}\) during training, we also introduce a method that avoids diagonalization and allows (in principle) linear scaling with the number of nodes.

### Equivariant Message Passing Neural Network

Equivariant Message Passing Neural Networks (MPNNs) Batzner et al. (2022); Batatia et al. (2022) are graph neural networks that operate on graphs \(\mathcal{G}=(X,\mathcal{E})\) embedded in a vector space \(V\). The nodes \(x_{i}\in X\) are no longer only a list of indices, but belong to a configuration space \(\Omega\) that extends the vector space \(V\). For example, in the atomistic point cloud, \(x_{i}=(i,\mathbf{r}_{i},\theta_{i})\in\Omega:=\mathbb{N}\times\mathbb{R}^{3}\times \mathbb{Z}\) describing the positions and chemical species of each atom through which the graph is embedded into \(V:=\mathbb{R}^{3}\). The case of the **pure** graph can be recovered by setting \(\Omega=\mathbb{N}\). We are interested in learning graph maps of the form

\[\Phi\colon\mathcal{G}\to Z \tag{3}\]

where \(Z\) is an abstract target space, usually a vector space. As the input is a 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 a group \(G\) acts on both \(\Omega\) (and, therefore, on \(\mathcal{G}\)) and \(Z\), we say that \(\Phi\) is \(G\)-equivariant if,

\[\Phi\circ g=\rho(g)\Phi\qquad\forall g\in G, \tag{4}\]

where \(\rho\) is a representation of the group on the vector space \(Z\). A typical strategy is then to embed the nodes \(x_{i}\in X\) into a feature space, where a suitable representation of the group is available.
We represent the state of each node \(\sigma_{i}\) in the layer \(t\) of the MPNN by a tuple,

\[\sigma_{i}^{(t)}=(x_{i},\mathbf{h}_{i}^{(t)}), \tag{5}\]

where \(x_{i}\) defines the collection of node attributes of the graph as defined previously and \(\mathbf{h}_{i}^{(t)}\) are its learnable features. A forward pass of the network consists of multiple _message construction_, _update_, and _readout_ steps. During message construction, a message \(\mathbf{m}_{i}^{(t)}\) is created for each node by pooling over its neighbors,

\[\mathbf{m}_{i}^{(t)}=\bigoplus_{j\in\mathcal{N}(i)}M_{t}\big{(}\sigma_{i}^{(t)}, \sigma_{j}^{(t)}\big{)},\quad\mathbf{h}_{i}^{(t+1)}=U_{t}\big{(}\sigma_{i}^{(t)}, \mathbf{m}_{i}^{(t)}\big{)},\quad\Phi(\mathcal{G})=\phi_{\text{out}}\bigg{(}\Big{\{} \big{\{}\mathcal{R}_{t}\big{(}\sigma_{i}^{(t)}\big{)}\big{\}}_{i}\Big{\}}_{t} \bigg{)}. \tag{6}\]

where the individual operations have the following meaning:

* \(M_{t}\) is a learnable message function and \(\bigoplus_{j\in\mathcal{N}(i)}\) is a learnable, permutation invariant pooling operation over the neighbors of atom \(i\) (e.g., a sum);
* \(U_{t}\) is a learnable update function, transforming the message \(\mathbf{m}_{i}^{(t)}\) into new features \(\mathbf{h}_{i}^{(t+1)}\);
* \(\mathcal{R}_{t}\) is a learnable node readout, mapping node states \(\sigma_{i}^{(t)}\) to the per-node outputs;
* \(\phi_{\text{out}}\) is a global readout map, typically \(\phi_{\text{out}}(\{\{\mathcal{R}_{i}^{(t)}\}_{i}\}_{i})_{t})=\sum_{i,t} \mathcal{R}_{i}^{(t)}\).

Equivariant MPNNs are widely used for learning properties of 3D point clouds, such as molecules and materials. However, there are several limitations to their expressiveness,

* **Non-local :** MPNNs are restricted to a finite range, dictated by the number of iterations \(T\). A large number of properties in physics are long-range, which requires \(T\to\infty\). Adding a large number of layers leads to high computational overhead and poor expressivity due to oversmoothness (Di Giovanni et al., 2023).
* **Correlation order:** Most MPNNs use two body messages, which means that the pooling in Equation 6 is just a sum. The MACE architecture (Batatia et al., 2022b) extended the message construction to an arbitrary order, but increasing the order is still computationally demanding, especially for large \(T\).

## 4 Matrix Function Neural Networks

### General Matrix Function Neural Networks

Our MFN models act in the space of group equivariant matrix operators of the graph, \(\mathcal{H}(\mathcal{G})^{G}\), with \(G\) a reductive Lie group. In the same setting as in section 3.2, let \(\mathcal{G}\) be an undirected graph with \(n\) nodes, and the states \(\sigma_{i}\) are embedded in a configuration space \(\Omega\). The architecture operates in three stages at each layer, the **matrix construction**, the **matrix function**, and the **update**.

Matrix constructionThe space of matrix operators on graphs, \(\mathcal{H}(\mathcal{G})^{G}\), corresponds to the space of operators that are (1) **self-adjoint**, (2) **permutation equivariant**, and (3) **G-equivariant**. We consider operators expanded in a basis \(\{\phi_{m}\}_{m}\) with an available group action on each basis element. These operators correspond to square matrices \(\mathbf{H}\in\mathbb{C}^{bn\times bn}\) following basis truncation, where \(b\) is the basis size. Matrix blocks \(\mathbf{H}_{ij}\) have dimensions \(b\times b\), and the basis is indexed by a tuple \((m_{1},m_{2})\). An extra channel dimension, denoted \(c\), is introduced to learn multiple operators. Matrix entries in \(\mathbf{H}\) can generally be \(G\)-equivariant functions of the neighborhoods \(\mathcal{N}(i)\) and \(\mathcal{N}(j)\),

\[H^{(t)}_{cij,m_{1}m_{2}}=\sum_{\hat{c}}W^{(t)}_{\hat{c}c}\phi^{(t)}_{\hat{c}m_ {1}m_{2}}(\{\sigma_{k}^{(t)}\}_{k\in\mathcal{N}(i)\cup\mathcal{N}(j)}). \tag{7}\]

Figure 1: **Matrix function network architecture. Illustrating matrix construction and non-locality of matrix functions on a chain graph.**
The choice of basis \(\phi\) depends on the application (see the SI for detailed examples). In the case of the rotation group, any such function can be approximated with arbitrary accuracy using an ACE or MACE layer (Zhang et al., 2022; Batatia et al., 2022b). This result has been generalized to any reductive Lie group, using the \(G\)-MACE architecture (Batatia et al., 2023). It is important to note that one could use any equivariant architecture to learn these matrix entries. However, the use of more expressive approximators will result in better coverage of the operator space \(\mathcal{H}(\mathcal{G})^{G}\) and therefore better general expressivity. The full matrix inherits this equivariance of the basis,

\[\mathbf{H}_{ij}\circ g=\boldsymbol{\rho}(g)\mathbf{H}_{ij}\boldsymbol{\rho}^{ \ast}(g),\quad\forall g\in G \tag{8}\]

where \(\rho\) is an orthogonal matrix that denotes the group action of \(G\) on the basis element \(\phi\).

Matrix functionThe central operation in MFNs, which introduces long-range many-body effects, is the matrix function. Any continuous function \(f_{\theta}:\mathbb{R}\rightarrow\mathbb{R}\), with parameters \(\theta\), can be interpreted as acting on self-adjoint matrices \(\mathbf{H}\) via their spectral decomposition. If \(\mathbf{H}=\mathbf{U}\mathbf{\Lambda}\mathbf{U}^{\mathbf{T}}\) with \(\mathbf{U}\) orthogonal and \(\mathbf{\Lambda}\) diagonal. Alternatively, if the function is analytic, it can be extended to a matrix function by converting its power series into a formal power matrix series,

\[f_{\theta}(\mathbf{H})=\mathbf{U}f_{\theta}(\mathbf{\Lambda})\mathbf{U}^{ \mathbf{T}}. \tag{9}\]

where \(c_{a}(\theta)\) are the weights of the power series of \(f_{\theta}\). An essential observation is that **any** continuous matrix function \(f_{\theta}\) preserves equivariance,

\[f_{\theta}(\mathbf{H}\circ g)=\boldsymbol{\rho}(g)f_{\theta}(\mathbf{H}) \boldsymbol{\rho}^{\ast}(g),\quad\forall g\in G. \tag{10}\]

The matrix function can be related to a high-order many-body equivariant function via the Cayley-Hamilton theorem. The eigendecomposition in Equation 9 is responsible for the non-locality of our approach. In practice, computing matrix functions is expensive as it requires diagonalization, scaling as \(n^{3}\) 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 \(\mathbf{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 \(f_{\theta}\), 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).

UpdateMultiple updates can be defined from the matrix function. All of the following updates differ fundamentally from the standard spectral GNN update, which is a filtering operation.

1. The **diagonal update** updates the state of each node with non-local features extracted from the diagonal blocks of the matrix function, \[\boldsymbol{h}_{ic}^{(t+1)}=\sum w_{c\bar{c}}^{(t)}f_{\theta}^{(t)}(\mathbf{H }_{\bar{c}}^{(t)})_{ii}\] (11) 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. Section 4.2.
2. The **dense update** utilizes the entire matrix, including all **off-diagonal terms**, to update the next matrix function, \[\boldsymbol{h}_{ic}^{(t+1)}=\sum_{\bar{c}}w_{c\bar{c}}^{(t)}f_{\theta}^{(t)}( \mathbf{H}_{\bar{c}}^{(t)}+f_{\theta}^{(t-1)}(\mathbf{H}_{\bar{c}}^{(t-1)}))_ {ii}\] (12) This update can not be performed using selected inversion techniques, but it can offer additional expressiveness guarantees.
3. The **sparse update:** uses only the parts of the matrix corresponding to the connected nodes to update the nodes and edges of the graph to obtain the matrix function in the next layer, \[\boldsymbol{h}_{ic}^{(t+1)}=\sum_{\bar{c}}w_{c\bar{c}}^{(t)}f_{\theta}^{(t)}( \mathbf{H}_{\bar{c}}^{(t)}+\{f_{\theta}^{(t-1)}(\mathbf{H}_{\bar{c}}^{(t-1)})_ {ij}\}_{j\in\mathcal{N}(i)})_{ii}.\] (13)

ReadoutThe readout phase is the same as the usual MPNN readout in Equation 6.
### Resolvent parameterization of matrix functions

The evaluation of the matrix function in Equation 9 is the practical bottleneck of our method. The cost of the evaluation depends on the choice of parameterization of the univariate function \(f_{\theta}\). For a general analytic \(\tilde{f}:\mathbb{C}\rightarrow\mathbb{C}\), resolvent calculus allows us to represent

\[\tilde{f}(\mathbf{H})=\oint_{\mathcal{C}}\tilde{f}(z)(z\mathbf{I}-\mathbf{H})^{- 1}\frac{dz}{2\pi i}, \tag{14}\]

where \(\mathcal{C}\) is a curve encircling the eigenvalues of \(\mathbf{H}\) and excluding any poles of \(\tilde{f}\). Approximating the contour integration with a quadrature rule with nodes \(z_{s}\) and weights \(w_{s}\) yields \(\tilde{f}(\mathbf{H})\approx\sum_{s}w_{s}\tilde{f}(z_{s})(z_{s}\mathbf{I}- \mathbf{H})^{-1}\), and merging \(w_{s}:=\tilde{w}_{s}\tilde{f}(z_{s})\) we arrive at the parameterization

\[f_{\theta}(\mathbf{H})=\sum_{s}w_{s}(z_{s}\mathbf{I}-\mathbf{H})^{-1}. \tag{15}\]

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 15 over standard usage of pole expansions in computing matrix functions (Higham, 2008) is that both the _weights_\(w_{s}\) and the _poles_\(z_{s}\) are now learnable parameters.

The derivation shows that in the limit of infinitely many pole-weight pairs \((z_{s},w_{s})\) 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 as 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 to avoid the spectrum of \(\mathbf{H}\), which is real since \(\mathbf{H}\) is assumed to be self-adjoint. Therefore, we choose adjoint pole weight pairs \((w_{s},z_{s})\) and \((w_{s}^{*},z_{s}^{*})\) to ensure that \(f_{\theta}\) is real when restricted to real arguments.

Linear scaling costThe pole expansion framework is the first key ingredient in linear scaling electronic structure methods Goedecker (1999) such as PEXSI Lin et al. (2009, 2013). The second ingredient is the selected inversion of the resolvent. Instead of computing the full inverse, \((z\mathbf{I}-\mathbf{H})^{-1}\), one first computes a sparse \(LDL^{*}\) factorization and then selectively computes only the diagonal entries of the resolvent. The bottleneck in this approach is the \(LDL^{*}\) factorization. For a full matrix, it scales as \(O(n^{3})\) operations and \(O(n^{2})\) memory. The complexity improves considerably for sparse matrices. Suppose that the sparsity pattern is \(d\)-dimensional, corresponding to a topologically \(d\)-dimensional graph; e.g. the cumulates in Section 5.1 are topologically one-dimensional despite being embedded in \(\mathbb{R}^{3}\). Using nested dissection ordering to minimize the fill-in, the cost of the \(LDL^{*}\) factorization reduces to \(O(n)\) for \(d=1\) (e.g., natural language processing and quasi-1D molecular systems such as carbon-nano-tubes); \(O(n^{3/2})\) operations and \(O(n\log n)\) memory for \(d=2\) (e.g., image recognition); and \(O(n^{2})\) operations and \(O(n^{4/3})\) memory for \(d=3\). The final step to reduce the computational cost of our architecture to _linear scaling_ is to replace the \(LDL^{*}\) factorization with an incomplete factorization, as proposed by Etter (2020) in the context of electronic structure calculations. They demonstrated that the \(LDL^{*}\) factorization inherits the decay of the resolvents and, therefore, an incomplete factorization can be performed with controllable error. However, this scheme remains a theoretical idea, and incorporating it into our architecture results in a compromise between computational cost and non-locality that requires significant further investigation.

### Expressivity of Matrix Function Networks

Non-local interaction

Purely formally, one may think of a matrix function \(\tilde{f}(\mathbf{H})\) as an infinite power series. This suggests that MFNs are inherently non-local and exhibit a convolution-like structure, similar to message-passing methods. This is 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 our method.

The degree of non-locality of the interaction can be precisely quantified. If \(\mathbf{H}\) has a finite interaction range, then the Combe-Thomas theorem Combes & Thomas (1973) implies that \(|(z\mathbf{I}-\mathbf{H})_{i,j}^{-1}|\leq Ce^{-\gamma_{s}d_{ij}}\), where \(\gamma_{z}=\mathrm{dist}(z,\sigma(\mathbf{H}))\) and \(d_{ij}\) is the length of the shortest path from node \(i\) to node
\(j\). Since we have taken \(\mathbf{H}\) self-adjoint with real spectrum, an estimate for \(\gamma_{z}\) is \(\Im z\). As a result, if we constrain the poles in the parameterization of \(f_{\theta}\) to be at \(\Im z_{s}=\gamma\), then the resulting matrix function will satisfy

\[\big{|}[f_{\theta}(\mathbf{H})]_{i,j}\big{|}\leq Ce^{-\gamma d_{ij}}. \tag{16}\]

Therefore, the degree of locality can be controlled by constraining \(\Im z_{s}\) to be some fixed value. While equation 16 only gives an upper bound, it is not difficult to see that it can in fact be attained for a suitable choice of matrix function. In practice, however, \(f_{\theta}\) can be seen as approximating an unknown target \(\tilde{f}\) with learnable weights \(w_{s}\). If \(\tilde{f}\) is analytic in a neighborhood of \(\sigma(\mathbf{H})\), then \(\big{|}[\tilde{f}(\mathbf{H})]_{i,j}\big{|}\leq Ce^{-\tilde{\gamma}d_{ij}}\), where \(\tilde{\gamma}\) measures how smooth \(\tilde{f}\) is (distance of poles of \(\tilde{f}\) to \(\sigma(\mathbf{H})\)). For our parameterization \(f_{\theta}\) we therefore expect the same behavior \(\gamma=\tilde{\gamma}\) in the pre-asymptotic regime of moderate \(d_{ij}\).

Geometric expressiveness

The WL-test quantifies a network's ability to distinguish non-isomorphic graphs. Its extension to graphs embedded in vector spaces is known as the geometric WL-test (Joshi et al., 2023). The expressivity in the following discussion adheres to these definitions. In the context of graphs in \(\mathbb{R}^{3}\), practical relevance is high. A direct equivalence exists between equivariant MPNNs with linear updates and infinite layers and a one-layer MFN due to the matrix function's product structure. For a one-layer MFN with a single matrix channel and two-body matrix entries, the expressiveness matches that of a two-body equivariant MPNN with **infinite layers** and **linear updates**. When matrix entries incorporate features beyond two-body, the one-layer MFN becomes more expressive than its two-body feature counterpart.

### Matrix normalization

Batch normalization (Ioffe and Szegedy, 2015) and layer normalization (Ba et al., 2016) are two techniques widely used in deep learning to stabilize training, improve convergence, and provide some form of regularization.

In the context of MFNs, a similar normalization strategy is needed to ensure stable training and improved convergence. Instead of normalizing the features directly as conventional normalization layers do, we normalize the eigenvalues of a set of H matrices with batch dimension \([\) batch, channels, \(n\), \(n\)], where \(n\) is the size of the matrix. The aim is 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,

\[\mathbb{E}(\Lambda)=\frac{\text{tr}(\mathbf{H})}{n},\quad\text{Var}(\Lambda) =\frac{\text{tr}(\mathbf{H}^{2})}{n-1}-\frac{\text{tr}(\mathbf{H})^{2}}{n(n-1)} \tag{17}\]

This normalization of eigenvalues ensures that the spectral properties of the graph used for representation in the MFN are effectively standardized, contributing to better training stability and convergence.

### Multivariate Matrix function

As MFNs extract features from multiple operators, it is useful to extend the univariate matrix function in Equation 9 to a multivariate matrix function. The space of multivariate functions \(\mathcal{F}(\mathcal{H}(\mathcal{G})\times...\times\mathcal{H}(\mathcal{G}))\) is isomorphic to the closure of the span of the tensor product of functions of single operators \(\mathcal{F}(\mathcal{H}(\mathcal{G}))^{\otimes n}\), We call 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,

\[f(\mathbf{H}_{1},...,\mathbf{H}_{n})=\frac{1}{\left(2\pi i\right)^{n}}\int \cdots\iint_{\partial D_{1}\times\cdots\times\partial D_{n}}\frac{f(z_{1}, \ldots,z_{n})}{\left(z_{1}\mathbf{I}-\mathbf{H}_{1}\right)\cdots\left(z_{n} \mathbf{I}-\mathbf{H}_{n}\right)}dz_{1}\cdots dz_{n} \tag{18}\]

Multivariate matrix functions create higher order features. A one-layer MFN with a matrix function of correlation order \(n\) is as expressive as a \((n+1)\)**-body** equivariant MPNN with **infinite layers** and **linear updates**. The space \(\mathcal{F}(\mathcal{H}(\mathcal{G}))^{\otimes n}\) is of very high dimension and it is possible to use standard compression techniques to find more efficient approximations such as tensor decomposition (Darby et al., 2023). The use of linear combinations of matrices in Equation 7 approximates some of this space. We leave further exploration of these kinds of network for future work.
### 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 to the action of the group of rotations, translations, reflections, and permutations. When these operators are expended on the basis of the irreducible representations of the rotation group, they exhibit a block structure (see Appendix A.1) in which each entry is labeled with representations \((ss,ps,pp,dd,...)\).

The most central of these operators is the Hamiltonian operator. The energy of a quantum system is related to the trace of a specific matrix function of the Hamiltonian,

\[E=\text{Tr}f(\mathbf{H}) \tag{19}\]

The Hamiltonian is usually computed as a fixed point of a self-consistent 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 to tight-binding methods that usually construct a single low-body Hamiltonian with physics-inspired functional forms but require several self-consistent iterations to converge.

## 5 Results

### Cumulenes: non-local 3D graphs

We compare the non-locality of MFNs to local MPNNs and global attention MPNNs using linear carbon chains, called cumulenes. Cumulenes are made up of double-bonded carbon atoms terminated with two hydrogen atoms at each end. Cumulenes exhibit pronounced non-local behavior as a result of strong electron delocalization. Small changes in chain length and relative angle between the terminating hydrogen atoms can result in large changes in the energy of the system, as visually represented in Figure 3. These structures are known to illustrate the limited expressivity of local models (Unke et al., 2021) and are 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 are capable of accurately extrapolating to longer chains, simultaneously capturing length and angle trends. In contrast, global attention models such as Spookynet (Unke et al., 2021), are unable to extrapolate to longer chains, highlighting the benefit of the matrix function formalism. For all MFN models in this section, we use MACE layers (Batatia et al., 2022) to form the matrix at each layer. We refer to the model as MFN (MACE). Details of the training set and the specific choice of parameterization of the matrix entries are included in the Appendix A.2.

**Trends with chain length and rotation** The lowest energy structure of cumulenes alternates between 90- and 0-degree angles for odd and even carbon atom counts, respectively. Consequently, varying the number of carbon atoms at a fixed angle unveils a distinctive zigzag pattern in the ground truth energy (Fig. 3 left panel). Although the local model, with two message passing layers, is trained on exactly these configurations, this system is beyond its expressivity, as can be seen by the linear trend for \(n_{c}>7\) in (Figure 3 left panel). In contrast, the invariant and equivariant MFN models perfectly reproduce density functional theory (DFT) energies, thanks to their inherent non-locality. To

Figure 2: **Block structure** of a Euclidean Operator, H, learnt in the MFNs. Each entry in H corresponds to a different product of representations of the group of rotations \((ss,sp,pp,...)\).
demonstrate the necessity of equivariance, we train models on the energy of a fixed size cumulene as a function of the dihedral angle between the hydrogen atoms. Figure 3 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 (Kovacs et al., 2023), complicating the assessment of model non-locality. The dataset introduced here is based on 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 a similar fashion with the same number of carbons (in-domain) and cumulenes of unseen length, not present in the dataset (out-domain 11,12 and 15,16). Table 1 shows that the MFN architecture significantly outperforms both local and attention-based models (Spookynet). Attention captures some non-locality, resulting in marginally lower errors on the train and in-domain test set. However, the learned non-locality does not generalize to larger molecules, obtaining energy and forces worse than those obtained with a simple local model. The structured non-locality of MFN enables generalization to larger unseen system sizes.

### Performance on pure graphs

In this section, we 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.2.

\begin{table}
\begin{tabular}{l c c c c c c c} \hline \hline
**Dataset** & \(n_{C}\) & \multicolumn{3}{c}{**E (meV/atom)**} & \multicolumn{3}{c}{**F (meV/A)**} \\  & \multicolumn{1}{c}{MACE} & SpoodyNet & MFN (MACE) & MACE & SpoodyNet & MFN (MACE) \\ \hline Train & 3-10,13,14 & 41.1 & 31.4 & **2.0** & 179.6 & 114.1 & **31.7** \\ Test (In Domain) & 3-10,13,14 & 41.8 & 30.8 & **2.6** & 205.6 & 162.3 & **34.0** \\ Test (Out Domain) & 11,12 & 16.5 & 31.4 & **0.9** & 108.5 & 116.2 & **22.5** \\ Test (Out Domain) & 15,16 & 12.0 & 23.4 & **2.6** & 72.1 & 87.6 & **37.7** \\ \hline \hline \end{tabular}
\end{table}
Table 1: **Guaranteed non-local cumulene dataset** containing rattled cumulene chains, with various chain lengths (\(n_{C}\)) and hydrogen dihedral angles (\(\phi\)). 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.

Figure 3: **Visualizing MFN expressivity on cumulene chains.** The left panel depicts energy trends with respect to cumulene chain length at a fixed angle \(\phi=5^{\circ}\). The right panel shows the DFT (ground truth) and the predicted energy as a function of the dihedral angle \(\phi\) between the hydrogen atoms for a cumulene chain containing 12 carbon atoms. Local many-body equivariant models (MACE) are only able to capture average trends, even though test configurations are included in the training set. Invariant MFNs (\(L=0\)) capture only the trends with respect to length, while equivariant MFNs (\(L=1\)) capture both non-local trends. All models have a cutoff distance \(r_{c}\) of 3Å, corresponding to the nearest neighbors, with two message-passing layers. The cutoff distance as well as MACE’s receptive field for the first carbon atom is annotated in the left panel.
**ZINC.** We use the default dataset splits for ZINC, aligned with the leaderboard baseline settings, with approximately 500K parameters set. Table 2 shows that MFN surpasses previous architectures, demonstrating the utility of learning various operators, even on pure graphs.

**TU Datasets.** We test our model on five TUDataset datasets involving both bioinformatics datasets (MUTAG, ENZYMES, PTC MR, and PROTEINS) and a social network dataset (IMDB-B). To ensure a fair comparison with baselines, we follow the standard 10-fold cross-validation and dataset split in table 3.

## 6 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 a computationally efficient solution. Our evaluations indicate state-of-the-art performance on 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.

### Reproducibility Statement

To ensure reproducibility and completeness, we include detailed descriptions of the model used, hyperparameters, and data sets in the Appendix. The ZINC and TU datasets are publicly available. We also attach our cumulene datasets in a supplementary. The code will be made public.

#### 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 reference 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.

\begin{table}
\begin{tabular}{l|c c c c c c c} \hline Method & GCN & GAT & MPNN & GT & SAN & Graphomer & PDF & MFN (GCN) \\ \hline MAE & 0.367\(\pm\)0.011 & 0.384\(\pm\)0.007 & 0.145\(\pm\)0.007 & 0.226\(\pm\)0.014 & 0.139\(\pm\)0.006 & 0.122\(\pm\)0.006 & 0.066\(\pm\)0.004 & **0.063\(\pm\)0.002** \\ \#para & 50\% & 531k & 481k & NA & 50\% & 48\% & 472k & 512k \\ \hline \end{tabular}
\end{table}
Table 2: 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.2.2.

\begin{table}
\begin{tabular}{l|c c c c c} \hline Method & MUTAG & ENZYMES & PTC\_MR & PROTEINS & IMDB-B \\ \hline GK & 81.52\(\pm\)2.11 & 32.70\(\pm\)1.20 & 55.65\(\pm\)0.5 & 71.39\(\pm\)0.3 & - \\ RW & 79.11\(\pm\)2.1 & 24.16\(\pm\)1.64 & 55.91\(\pm\)0.3 & 59.57\(\pm\)0.1 & - \\ PK & 76.0\(\pm\)2.7 & - & 59.5\(\pm\)2.4 & 73.68\(\pm\)0.7 & - \\ AWE & 87.87\(\pm\)9.76 & 35.77\(\pm\)5.93 & - & - & 74.45\(\pm\)5.80 \\ \hline PSCN & 88.95\(\pm\)4.4 & - & 62.29\(\pm\)5.7 & 75\(\pm\)2.5 & 71\(\pm\)2.3 \\ ECC & 76.11 & 45.67 & - & - & - \\ DGK & 87.44\(\pm\)2.72 & 53.43\(\pm\)0.91 & 60.08\(\pm\)2.6 & 75.68\(\pm\)0.5 & 66.96\(\pm\)0.6 \\ GraphSAGE & 85.1\(\pm\)7.6 & 58.2\(\pm\)6.0 & - & - & 72.35\(\pm\)3.5 \\ CapsGNN & 88.67\(\pm\)6.88 & 54.67\(\pm\)5.67 & - & 76.2\(\pm\)3.6 & 73.14\(\pm\)1.8 \\ GIN & 89.4\(\pm\)5.6 & - & 64.6\(\pm\)7.0 & 76.2\(\pm\)2.8 & 75.1\(\pm\)5.1 \\ \(k\)-GNN & 86.1 & - & 60.9 & 75.5 & 74.2 \\ IQN & 83.89\(\pm\)12.95 & - & 58.53\(\pm\)6.86 & 76.58\(\pm\)5.49 & 72.0\(\pm\)5.54 \\ PPGNN & 90.55\(\pm\)8.7 & - & 66.17\(\pm\)6.54 & **77.20\(\pm\)4.73** & 73.0\(\pm\)5.77 \\ GCN\({}^{2}\) & 89.39\(\pm\)1.60 & - & 66.84\(\pm\)1.79 & 71.71\(\pm\)1.014 & 74.80\(\pm\)2.01 \\ PDF & 89.01\(\pm\)4.35 & **73.50\(\pm\)6.39** & 68.36\(\pm\)8.38 & 76.28\(\pm\)5.1 & **75.60\(\pm\)2.69** \\ \hline MFN (GCN) & **91.5\(\pm\)7.35** & 72.9\(\pm\)7.55 & **68.9\(\pm\)8.09** & 76.18\(\pm\)4.07 & 74.1\(\pm\)1.04 \\ \hline \end{tabular}
\end{table}
Table 3: Results on TUDataset (Higher is better). Bold is best, and underlined second best within \(\pm 0.5\%\). Baselines are taken from (Yang et al., 2023) and model citations are in A.2.3
## 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.
