Neural networks for geospatial data
Analysis of geospatial data has traditionally been model-based, with a mean
model, customarily specified as a linear regression on the covariates, and a
covariance model, encoding the spatial dependence. We relax the strong
assumption of linearity and propose embedding neural networks directly within
the traditional geostatistical models to accommodate non-linear mean functions
while retaining all other advantages including use of Gaussian Processes to
explicitly model the spatial covariance, enabling inference on the covariate
effect through the mean and on the spatial dependence through the covariance,
and offering predictions at new locations via kriging. We propose NN-GLS, a new
neural network estimation algorithm for the non-linear mean in GP models that
explicitly accounts for the spatial covariance through generalized least
squares (GLS), the same loss used in the linear case. We show that NN-GLS
admits a representation as a special type of graph neural network (GNN). This
connection facilitates use of standard neural network computational techniques
for irregular geospatial data, enabling novel and scalable mini-batching,
backpropagation, and kriging schemes. Theoretically, we show that NN-GLS will
be consistent for irregularly observed spatially correlated data processes. We
also provide a finite sample concentration rate, which quantifies the need to
accurately model the spatial covariance in neural networks for dependent data.
To our knowledge, these are the first large-sample results for any neural
network algorithm for irregular spatial data. We demonstrate the methodology
through simulated and real datasets.
Wentao Zhan, Abhirup Datta
# Neural networks for geospatial data

###### Abstract

Analysis of geospatial data has traditionally been model-based, with a mean model, customarily specified as a linear regression on the covariates, and a covariance model, encoding the spatial dependence. While non-linear machine learning algorithms like neural networks are increasingly being used for spatial analysis, current approaches depart the model-based setup and cannot not explicitly incorporate the spatial covariance. We propose embedding neural networks directly within the traditional geostatistical models to accommodate non-linear mean functions while retaining all other advantages of the model-based setup including use of Gaussian Processes to explicitly model the spatial covariance, enabling inference on the covariate effect through the mean and on the spatial dependence through the covariance, and offering predictions at new locations via kriging. We propose _NN-GLS_, a new neural network estimation algorithm for the non-linear mean that explicitly accounts for the spatial covariance through generalized least squares (GLS), the same loss used in the linear case. We show that NN-GLS admits a representation as a special type of graph neural network (GNN). This connection facilitates use of standard neural network computational techniques for irregular geospatial data, enabling novel and scalable mini-batching, backpropagation, and kriging schemes. Theoretically, we show that NN-GLS will be consistent for irregularly observed spatially correlated data processes. To our knowledge this is the first asymptotic consistency result for any neural network algorithm for spatial data. We demonstrate the methodology through numerous simulations and an application to air pollution modeling.

_Keywords:_ Geostatistics, spatial statistics, Gaussian process, neural networks, machine learning, consistency.
Introduction

_Geostatistics_, the analysis of geocoded data, is traditionally based on stochastic process models which offer a coherent way to model data at any finite collection of locations while ensuring the generalizability of inference to the entire region. _Gaussian processes (GP)_ with a mean function capturing effects of covariates and the covariance function encoding the spatial dependence, is a staple for geostatistical analysis, offering theoretical guarantees and practical benefits. GP are flexible enough to model any smooth spatial surface, and can be specified parsimoniously with covariance functions using a very small set of parameters. The spatial covariance parameters offer insights into the smoothness and spatial properties of the response process (Stein, 1999). The finite dimensional realizations of a GP are multivariate Gaussian, thereby offering estimates of the mean and covariance parameters via convenient maximization of the Gaussian likelihood, and predictions at new locations by using conditional Gaussian distributions (see, e.g., Banerjee et al., 2014; Cressie and Wikle, 2015, for detailed exposition on GP models for spatial and spatio-temporal data). Also, computational roadblocks to using GP for large spatial data have been greatly mitigated by recent advances (see, Heaton et al., 2019, for a recent review of scalable GP approaches).

The mean function of a Gaussian process is often modeled as a linear regression on the covariates. The growing popularity and accessibility of machine learning algorithms such as neural networks, random forests, gradient boosted trees, capable of modeling complex non-linear relationships has heralded a paradigm shift. Practitioners are increasingly shunning models with parametric assumptions like linearity in favor of these machine learning approaches that can capture non-linearity and high-order interactions in a data-driven manner. The field of spatial statistics has not been insulated from this machine learning revolution. In particular, deep neural networks, a class of algorithms that has become almost synonymous with machine learning, have seen considerable recent adoption and
adaptation for geospatial data (see Wikle and Zammit-Mangion, 2023, for a comprehensive review).

Many of the machine learning based regression approaches assume independent observations, implicit in the choice of a sum of squared error loss (_ordinary least squares_ or _OLS_ loss) as the objective function used in estimating the algorithm parameters. Explicit encoding of spatial dependency via a covariance function, as is common in process-based geospatial models, is challenging within these algorithms. Current renditions of neural networks for spatial data circumvent this by using spatial co-ordinates or some transformations (like distances or basis functions) as additional covariates (Gray et al., 2022; Chen et al., 2020; Wang et al., 2019). These methods incorporate all spatial information directly into the mean, implicitly assuming that the error terms are independent. While this approach is very non-parametric conceptually, its explanatory power can be limited in practice. Many spatial features (basis functions) might be needed to capture the complex'shape' of spatial effect, resulting in a high-dimensional feature set (covariates) and triggering 'the curse of dimensionality'. Finding a balance between parsimony and explainability requires a careful design of spatial features that can explain complex spatial dependencies without drowning out the effect of the non-spatial covariates. Adding spatial features has also been considered for random forests (Hengl et al., 2018) and they have been shown to perform well only when the spatial signal is considerable compared to the covariate signal (Saha et al., 2023).

The second disadvantage of neural networks with added spatial features is that they obfuscate the separation of the spatial and non-spatial variations in the data. They model the mean as a joint function of covariates and space and do not offer an estimate of the covariate effect only, impeding understanding relationships between the response and the covariates. This is in stark contrast to the traditional GP based geospatial models, that model the covariate effect though the mean and the spatial variation through the covariance. This separation of the two sources of variation into first- and second-order effects naturally
induces parsimony via concepts like stationarity or isotropy for the covariance while still retaining the ability to model complex spatial dependence.

We propose a novel neural network algorithm to estimate non-linear means within the traditional Gaussian process models, explicitly accounting for the spatial dependence encoded in the GP covariance matrix. The core motivation comes from extension of the ordinary least squares (OLS) to generalized least squares (GLS) for linear models with dependent errors. For correlated data, GLS is more efficient than OLS according to the Gauss-Markov theorem. We modify the OLS loss used in neural networks to its GLS version. We refer to our algorithm as _NN-GLS_. We retain all advantages of the model-based framework, including separation of the covariate and spatial effects thereby allowing inference on both, parsimonious modeling of the spatial effect through the GP covariance function circumventing the need to create and curate spatial features, and seamlessly offering predictions at new locations via kriging. NN-GLS is compatible with any network architecture for the mean function and with any family of GP covariance function.

We note that the philosophy of GLS has been recently adopted for spatial analysis using tree-based machine learning algorithms like Random Forests (_RF-GLS_, Saha et al., 2023; Saha and Datta, 2023) and boosted trees (_GP-boost_, Sigrist, 2022; Iranzad et al., 2022). Forest and tree estimators use a brute force search to iteratively grow the regression trees, requiring multiple evaluation of the GLS loss within each step. This severely reduces the scalability of these approaches. RF-GLS also requires prior estimation of the spatial parameters, which are then kept fixed during the random forest estimation part.

NN-GLS avoids both these issues by offering a representation of the algorithm as a special type of _graph neural network (GNN)_. We show that NN-GLS using any multi-layer neural network architecture for the mean and a Nearest Neighbor Gaussian Process (NNGP, Datta et al., 2016_a_) for the covariance is a GNN with additional graph-convolution layers based on the nearest-neighbor graph and with graph-weights derived from kriging.
Leveraging this representation of the model as a GNN, we can exploit the various computing techniques used to expedite neural networks. This includes using _mini-batching_ or _stochastic gradients_ to run each iteration of the estimation on only a subset of the data in each cycle, as well as use of _backpropagation_ to compute the gradients. We provide novel scalable mini-batching, backpropagation and kriging schemes using the GNN framework Also, spatial parameters are now simply weight parameters in the network and updated throughout the training. This connection of NN-GLS to GNN, demonstrating use of GNN for irregular spatial data, is of independent importance for other types of spatial analysis.

Theoretically, there is little asymptotic theory supporting current neural networks approaches for spatially correlated data. Even in the case of i.i.d. errors, the theory of neural networks, while being discussed for a long time (Hornik et al., 1989), is still in its nascent stages. Shen et al. (2019)'s recent work provides a framework to study asymptotics of one-layer neural network estimator class. A main contribution of our article is proving the consistency of NN-GLS when the observations have spatially correlated errors arising from a stochastic process. We first extend Shen et al. (2019)'s framework to prove general results on the existence and the consistency of neural networks using GLS loss for dependent data processes. The novelty in the consistency results is accommodating spatial dependency in the data which prevents the traditional techniques (such as symmetrization and Hoeffding-type bound) from being used. Consequently, direct application of Shen et al. (2019)'s results is inapplicable to our case. As a solution, we design a new Orlicz norm in the function space for the successful application of uniform laws of large numbers.

We then apply the general results to prove consistency of NN-GLS for irregular spatial data designs and popular GP models like the Matern Gaussian process. We also show that consistency holds when using NNGP for the GLS loss. To our knowledge, these are the first theoretical result for neural networks under spatial dependence and the first examples of consistency of any machine learning approach to estimating a non-linear mean of a GP
for irregular spatial designs.

The rest of the manuscript is organized as follows. In section 2, we review process-based geostatistical models and neural networks. Section 3 describes the idea of combining the two perspectives. Section 4 formally proposes the algorithm NN-GLS, depicts its connection with graph neural networks (GNN), and exploits this representation to provide scalable estimation and prediction algorithms. Section 5 presents the theoretical results. Section 6 and 7 respectively demonstrate the performance of NN-GLS in simulated and real data.

## 2 Preliminaries

### Process-based geostatistical modeling

Consider spatial data collected at locations \(s_{i}\), \(i=1,\ldots,n\), comprising of a covariate vector \(\mathbf{X}_{i}:=\mathbf{X}(s_{i})\in\mathbb{R}^{d}\) and a response \(Y_{i}:=Y(s_{i})\in\mathbb{R}\). Defining \(\mathbf{Y}=(Y_{1},\ldots,Y_{n})^{\prime}\) and the \(n\times d\) covariate matrix \(\mathbf{X}\) similarly, the spatial linear model is given by

\[\mathbf{Y}\sim\mathcal{N}(\mathbf{X}\mathbf{\beta},\mathbf{\Sigma}(\mathbf{\theta})) \tag{1}\]

where \(\mathbf{\Sigma}:=\mathbf{\Sigma}(\mathbf{\theta})\) is a \(n\times n\) covariance matrix, parameterized by \(\mathbf{\theta}\).

A central objective in spatial analysis is to extrapolate inference beyond just the data locations to the entire continuous spatial domain. Stochastic processes are natural candidates for such domain-wide models. In particular, (1) can be viewed as a finite sample realization of a Gaussian process (GP)

\[Y(s)=\mathbf{X}(s)^{\top}\mathbf{\beta}+\epsilon(s),\epsilon(\cdot)\sim GP(0,\Sigma( \cdot,\cdot)). \tag{2}\]

where \(\epsilon(s)\) is a zero-mean Gaussian process modeling the spatial dependence via the co
variance function \(\Sigma(\cdot,\cdot)\) such that \(\Sigma(s_{i},s_{j})=Cov(Y(s_{i}),Y(s_{j}))\). Often, \(\epsilon(\cdot)\) can be decomposed into a latent spatial GP and a non-spatial (random noise) process. This results in the variance decomposition \(\mathbf{\Sigma}=\mathbf{C}+\tau^{2}\mathbf{I}\) where \(\mathbf{C}\) is the covariance matrix corresponding to the latent spatial GP and \(\tau^{2}\) is the variance for the noise process. One can impose plausible modeling assumptions on the nature of spatial dependence like stationarity (\(C(s_{i},s_{j})=C(s_{i}-s_{j})\)) or isotropy (\(C(s_{i},s_{j})=C(\|s_{i}-s_{j}\|)\)) to induce parsimony, thereby requiring a very-low dimensional parameter \(\mathbf{\theta}\) to specify the covariance function.

Gaussian processes have become a staple in geostatistical analysis for multitudinous reasons. It offers a coherent way to jointly model both observed data at a finite set of locations and unobserved data anywhere in the continuous spatial domain. They are non-parametric in the sense of being able to model any smooth spatial surface with suitable choice of covariance functions (Van Der Vaart and Van Zanten, 2011). Yet GP are parsimonious, specified by a parametric covariance function with a handful of parameters. From a practical perspective, finite-dimensional realizations of a GP are simply multivariate Gaussian leading to the data likelihood of the form (1). This effectuates simple estimation of the parameters \(\mathbf{\beta}\) and \(\mathbf{\theta}\) by maximum likelihood process. Additionally, predictions at new locations (kriging) simply reduces to obtaining the conditional Gaussian distributions.

### Neural networks

Artificial neural network (ANN) or simply, neural networks (NN) are widely used to model non-parametric regression of the form \(E(Y_{i})=f(\mathbf{X}_{i})\) where \(Y_{i}\) represents the univariate response and \(\mathbf{X}_{i}\) represents \(d\)-dimensional covariate vector. The unknown non-linear regression function \(f(\cdot)\) is approximated using a specific family of functions, which is the feed-forward part of a NN. Feed-forward neural networks, also referred to as a _multi-layer perceptrons (MLP)_, describes how the inputs \(\mathbf{X}_{i}\) are translated into outputs \(Y_{i}\). It is spec
ified via a sequence of hidden layers composed of a large number of nodes, connection between layers is specified via weights (as the contributions from each node) and activation (link) functions (to model non-linearity). Mathematically, an \(L\)-layer NN can be described as,

\[\begin{split}\mathbf{A}_{i}^{(0)}=\mathbf{X}_{i},&\mathbf{Z}_{i}^ {(l)}=\mathbf{W}_{(l)}^{\top}\mathbf{A}_{i}^{(l-1)},\,\mathbf{A}_{i}^{(l)}=\mathbf{g}_{l}( \mathbf{Z}_{i}^{(l)}),\,l=1,\ldots,L\\ & O_{i}=\mathbf{W}_{(L+1)}^{\top}\mathbf{A}_{i}^{(L)},\,f(\mathbf{X}_{i} )=O_{i},\,i=1,\ldots,n\end{split} \tag{3}\]

where for layer \(l\), \(\mathbf{A}^{(l)}\) represents the \(d_{l}\) nodes, \(\mathbf{W}_{(l)}\)'s are the weight matrix, \(\mathbf{Z}_{i}^{(l)}\)'s are hidden features as a linear combination of nodes with weights \(\mathbf{W}_{(l)}\)'s, and \(\mathbf{g}_{(l)}(\mathbf{Z}_{i}^{(l)})\) denotes the non-linear activation functions \(g_{l}(\cdot)\) (e.g., sigmoid function, ReLU function) applied to each component of \(\mathbf{Z}_{i}^{(l)}\). The final layer \(O_{i}\) is called the _output layer_ and gives the modeled mean of the response, i.e., \(O_{i}=f(\mathbf{X}_{i})=E(Y_{i})\). The choice of the activation functions is made before the training, and the weight matrices are estimated during the training process. Parameter estimation in NN is done using back-propagation, where estimates of weight matrices are updated based on the current model's fits on the training data. For regression, the fit accuracy is measured by the mean squared error (MSE) or ordinary least squares (OLS) loss

\[\sum_{i=1}^{n}(Y_{i}-f(\mathbf{X}_{i}))^{2}. \tag{4}\]

The training process iteratively switches between feed-forward (obtaining the fits given the current estimates of the weights) and back-propagation (updating the weights given the current fits). Estimation is expedited by mini-batching where the data are split into smaller and disjoint mini-batches and at each iteration the loss (4) is approximated by restricting to one of the mini-batches, and cycling among the mini-batches over iterations.
Neural networks for Gaussian process models

The two paradigms reviewed in Section 2 are complimentary in their scope. The popularity of the geospatial linear models (Section 2.1) is owed to their simplicity, interpretability, and parsimony - separating the covariate effect and the spatial effects, modeling the former through a linear mean, and the latter through the GP covariance matrix. However, it relies on the strong assumption of a linear covariate effect. Neural networks can estimate arbitrary non-linear covariate effect. However, implicit in the usage of the OLS loss (4) for NN is the assumption that the data units are independent. This is violated for spatial data, where the error process is a dependent stochastic process as in (2).

We bridge the two paradigms by allowing the mean of the GP model to be non-linear in the covariates and proposing a novel NN algorithm to estimate it. We extend (2) to

\[Y_{i}=f(\mathbf{X}_{i})+\epsilon(s_{i});\epsilon(\cdot)\sim GP(0,\Sigma(\cdot, \cdot)). \tag{5}\]

Estimating \(f\) in (5) using neural networks (3) needs to account for the spatial dependence modeled via the GP covariance. Using the OLS loss (4) ignores this covariance. We now extend NN to explicitly accommodate the spatial covariance in its estimation.

### NN-GLS: Neural networks using GLS loss

Let \(\mathbf{f}(\mathbf{X})=\left(f(\mathbf{X}_{1}),f(\mathbf{X}_{2}),\cdots,f(\mathbf{X}_{n})\right)^ {\top}\). Then the data likelihood for the non-linear spatial model (5) is

\[\mathbf{Y}\sim N(f(\mathbf{X}),\mathbf{\Sigma}). \tag{6}\]

This is a direct generalization of the spatial linear model (1). It is well known that in models like (1) for correlated outcomes, an OLS loss is replaced by a _generalized least squares (GLS)_ loss using the covariance matrix \(\mathbf{\Sigma}\). To estimate \(f\) in (6) using NN, we
naturally replace the OLS loss (4) with the GLS loss

\[\mathcal{L}_{n}(f)=\frac{1}{n}(\mathbf{Y}-\mathbf{f(X)})^{\top}\mathbf{Q}(\mathbf{Y}-\mathbf{f(X) }), \tag{7}\]

accounting for the spatial dependency with the working precision matrix \(\mathbf{Q}\), which equals \(\mathbf{\Sigma}^{-1}\) or, more practically, an estimate of it.

We refer to the neural network estimation using the GLS loss (7) as _NN-GLS_. Conceptually, generalizing NN to NN-GLS is well-principled as minimizing the GLS loss (7) with \(\mathbf{Q}=\mathbf{\Sigma}^{-1}\) is equivalent to obtaining a maximum likelihood estimate of \(f\) in (6). In practice however, for spatial dependence modeled using GP, the GLS loss ushers in multiple computational issues for both mini-batching and backpropagation, techniques fundamental to the success of NN. The back-propagation equations corresponding to the GLS loss will involve computing an inverse of the dense \(n\times n\) matrix \(\mathbf{\Sigma}\) which requires \(O(n^{2})\) storage and \(O(n^{3})\) time. This is repeated multiple times as the parameters \(\mathbf{\theta}\) specifying \(\mathbf{\Sigma}\) are updated. These computing needs are infeasible for even moderate \(n\).

Furthermore, mini-batching involves partitioning the data into small subsets. Sensitivity to the choice of partitioning can be ignored for data with i.i.d. errors as under random partitioning the losses for the mini-batches are i.i.d. and are unbiased estimates of the loss for the full data. Partitioning of spatial data is challenging due to dependence across data units. Also, while the OLS loss (4) is naturally amenable to mini-batching due to its additive nature, this advantage is lost when switching to GLS loss as the quadratic form in (7) usually do not admit such additive decomposition for typical choices of the GP covariance function. In the next Section, we develop an algorithm _NN-GLS_ with a specific class of GLS loss that mitigates these issues and offers a pragmatic approach to use NN for GP models.
NN-GLS as Graph Neural Network

We offer a representation of NN-GLS as a special graph neural network (GNN) with OLS loss, a connection that allows developing scalable mini-batching and backpropagation algorithms for NN-GLS. We propose choosing \(\mathbf{Q}\) as the precision matrix from a Nearest Neighbor Gaussian Process (NNGP, Datta et al., 2016_a_). NNGP provides a sparse approximation to the dense full GP precision matrix \(\mathbf{\Sigma}^{-1}\) without requiring any large matrix operations. NNGP creates a _directed acyclic graph (DAG)_ based on pairwise distances among the \(n\) data locations, such that each node (location) has at most \(m\ll n\) directed (nearest) neighbors. Letting \(N(i)\) be the set of neighbors of location \(s_{i}\) in the DAG, NNGP yields a precision matrix \(\widetilde{\mathbf{\Sigma}}^{-1}=(\mathbf{I}-\mathbf{B})^{\top}\mathbf{F}^{-1}( \mathbf{I}-\mathbf{B})\)(Finley et al., 2019; Datta et al., 2016\(b\); Datta, 2022) where \(\mathbf{B}\) is a strictly lower triangular matrix and \(\mathbf{F}\) is a diagonal matrix with

\[\begin{split}\mathbf{B}_{i,N(i)}&=\mathbf{\Sigma} \big{(}i,N(i)\big{)}\mathbf{\Sigma}\big{(}N(i),N(i)\big{)}^{-1},\,B_{ij}=0\text{ elsewhere, and}\\ \mathbf{F}_{ii}&=\mathbf{\Sigma}_{ii}-\mathbf{\Sigma}\big{(} i,N(i)\big{)}\mathbf{\Sigma}\big{(}N(i),N(i)\big{)}^{-1}\mathbf{\Sigma}\big{(}N(i),i \big{)}.\end{split} \tag{8}\]

NNGP precision matrices can be obtained using only inversion of \(n\) small matrices of size \(m\times m\). As \(m\ll n\), it requires total \(O(n)\) time and storage. We note that basis functions derived from NNGP has been used as additional features to capture spatial structure in neural networks (Wang et al., 2019). This differs from our approach of using NNGP to directly model the spatial covariance, analogous to the practice in spatial linear models. We detail below how choosing \(\mathbf{Q}\) in the GLS loss (7) to be the NNGP precision matrix \(\widetilde{\mathbf{\Sigma}}^{-1}\) enables fast mini-batching and backpropagation for the NN parameter estimation.

We first offer a representation of NN-GLS using NNGP precision matrix in the GLS loss as a graph neural network using OLS loss. This connection will enable NN-GLS to retain the computational advantages of OLS loss. A GLS loss can be viewed as an OLS loss for the decorrelated response \(\mathbf{Y}^{*}=\mathbf{Q}^{\frac{1}{2}}\mathbf{Y}\), where \(\mathbf{Q}^{\frac{1}{2}}\) is the Cholesky factor of \(\mathbf{Q}=\mathbf{Q}^{\frac{\top}{2}}\mathbf{Q}^{\frac{1}{2}}\).
Hence, decorrelation is simply a linear operation. A convenience of choosing \(\mathbf{Q}=\widetilde{\mathbf{\Sigma}}^{-1}\), the NNGP precision matrix, is that decorrelation becomes a convolution in the nearest neighbor DAG. To elucidate, note that \(\mathbf{B}_{i,N(i)}\) defined in (8) denotes the kriging weights for predicting \(Y_{i}\) based on its directed nearest neighbors \(\mathbf{Y}_{N(i)}\) using a GP with covariance \(\Sigma(\cdot,\cdot)\). Similarly, \(F_{ii}\) in (8) is the corresponding nearest neighbor kriging variance. Letting \(N^{*}[i]=N(i)\cup\{i\}\) denote the graph neighborhood for the \(i^{th}\) node and defining weights

\[\mathbf{v}_{i}^{\top}=\frac{1}{\sqrt{F_{ii}}}(1,-\mathbf{B}_{i,N(i)}), \tag{9}\]

we can write \(Y_{i}^{*}=\mathbf{v}_{i}^{\top}\mathbf{Y}_{N^{*}[i]}\). Thus the decorrelated responses \(Y_{i}^{*}\) are simply convolution over the DAG used in NNGP with the graph convolution weights \(\mathbf{v}_{i}\) defined using kriging. Similarly, one can define the decorrelated output layer \(O_{i}^{*}=\mathbf{v}_{i}^{\top}\mathbf{O}_{N^{*}[i]}\) using the same graph convolution, where \(O_{i}\) is the output layer of the neural network \(f\) (see (3))

The decorrelation step makes NN-GLS a special type of graph neural network (GNN) as depicted in Figure 1. In GNN, typically, the input observations are represented on a graph and the locality information is aggregated using convolution layers based on the graph structure (graph convolution). For NN-GLS, both the inputs \(\mathbf{X}_{i}\) and the responses \(Y_{i}\) are graph-valued objects as they both correspond to the locations \(s_{i}\), which are the nodes of the nearest neighbor DAG. First, the inputs \(\mathbf{X}_{i}\)'s are passed through the feed forward NN (or multi-layer perceptron) to produce the respective output layer of \(O_{i}\)'s. This is a within-node operation, and any architecture can be used (number of layers, number of nodes within each layer, sparsity of connections, choice of activation functions). Subsequently, the output layer \(O_{i}\)'s from the MLP are passed through an additional graph-convolution to create the decorrelated output layer of \(O_{i}^{*}\)'s using the weights \(\mathbf{v}_{i}\) (9). This layer is matched, using the OLS loss, to the decorrelated response layer of \(Y_{i}^{*}\)', created from the \(Y_{i}\)'s using the same graph convolution. Thus fitting a GLS loss to a NN is simply fitting an OLS loss
to a new NN with two additional decorrelation layers at the end of the MLP.

To summarize, there are two ingredients of NN-GLS: a feed forward NN or MLP (using any architecture) for intra-node operations; and a sparse DAG among the locations for incorporation of spatial correlation via inter-node graph convolutions. Information about the mutual distances among the irregular set of locations are naturally incorporated in the kriging-based convolution weights. We now discuss how this formulation of NN-GLS as a GNN with OLS loss helps leverage the traditional strategies to scale computing for NN.

Figure 1: NN-GLS as a graph neural network with two decorrelation layers
### Mini-batching

Using the NNGP precision matrix \(\widetilde{\mathbf{\Sigma}}^{-1}\) as \(\mathbf{Q}\) simplifies the GLS loss (5) to be \(\mathcal{L}_{n}=\sum_{i=1}^{n}(Y_{i}^{*}-O_{i}^{*})^{2}\) where \(Y_{i}^{*}\)'s correspond to the decorrelated response layer and \(O_{i}^{*}\)'s are the decorrelated output layer in the GNN (Figure 1). Leveraging the additivity of this OLS loss, we can write \(\mathcal{L}_{n}=\sum_{b=1}^{B}\mathcal{L}_{b,n}\) where \(L_{b,n}=\sum_{i\in S_{b}}(Y_{i}^{*}-O_{i}^{*})^{2}\), \(S_{1},\ldots,S_{B}\) being a partition of the data-locations each of size \(K\). The \(Y_{i}^{*}\)'s are uncorrelated and identically distributed (exactly under NNGP distribution, and approximately if the true distribution is a full GP), so the loss \(L_{b,n}\) corresponding to the mini-batch \(S_{b}\) are approximately i.i.d. for \(b=1,\ldots,B\). Hence, parameter estimation via mini-batching can proceed like the i.i.d. case. The only additional computation we incur is during the graph convolution as obtaining all \(O_{i}^{*}\) for the mini-batch \(S_{b}\) involves calculating \(O_{i}\)'s for the neighbors of all units \(i\) included in \(S_{b}\).

### Back-propagation

Gradient descent or back-propagation steps for NN-GLS can be obtained in closed form and without any large matrix inversions. We provide the back-propagation equations for a single-layer network with \(O_{i}=\mathbf{\beta}^{\top}\mathbf{A}_{i}\), where \(\mathbf{\beta}\) and \(\mathbf{A}_{i}=\mathbf{g}(\mathbf{Z}_{i})\) are \(d_{1}\) dimensional, \(d_{1}\) being the number of hidden nodes, \(g\) is the known activation (link) function. Here \(\mathbf{Z}_{i}=\mathbf{W}^{\top}\mathbf{X}_{i}\) is the hidden layer created using weights \(\mathbf{W}^{\top}=(W_{rj})\) is the weight-matrix (with rows \(\mathbf{w}_{r}^{\top}\)). Denote \(z_{ir}=\mathbf{w}_{r}^{\top}X_{i}\), \(a_{ir}=g(z_{ir})\), \(a_{ir}^{*}=\mathbf{v}_{i}^{\top}\mathbf{a}_{N^{*}[i]r}\), \(\delta_{i}=-2(y_{i}-\mathbf{\beta}^{\top}A_{i})\) and \(\delta_{i}^{*}=\mathbf{v}_{i}^{\top}\delta_{N^{*}[i]}\).
Then we have the following customized back-propagation updates for NN-GLS:

\[\begin{split}\beta_{r}^{(t+1)}&=\beta_{r}^{(t)}-\gamma _{t}\sum_{i\in S_{b(t)}}\delta_{i}^{*}a_{ir}^{*}\\ w_{rj}^{(t+1)}&=w_{rj}^{(t)}-\gamma_{t}\beta_{r} \sum_{i\in S_{b(t)}}\delta_{i}^{*}\left(\mathbf{v}_{i}^{\top}(\mathbf{g}^{\prime}(\mathbf{ z}_{N^{*}[i]r})\odot\mathbf{X}_{N^{*}[i]j})\right),\\ \theta_{c}^{(t+1)}&=\theta_{c}^{(t)}+\frac{\gamma _{t}}{2}\sum_{i\in S_{b(t)}}\delta_{i}^{*}\left(\left(\frac{\partial\mathbf{v}_{i}} {\partial\theta_{c}}\right)^{\top}\delta_{N^{*}[i]}\right)\end{split} \tag{10}\]

Here \(\gamma_{t}\) is the learning rate and \(S_{b(t)}\) is the mini-batch for the \(t^{th}\) iteration, and \(g^{\prime}\) is the derivative of \(g\). These custom and scalable gradient descent steps for NN-GLS can also be conducted using off-the-shelf neural network software. One can simply input the mini-batch loss for NN-GLS to obtain scalable gradient descent updates for the NN parameters using numerical, automatic, or symbolic differentiation.

Note that in the GNN representation, the spatial covariance parameter \(\mathbf{\theta}\) are just the parameters of the convolution weights \(\mathbf{v}_{i}\). Hence, the updates for \(\mathbf{\theta}\) can be just absorbed as a back-propagation step as shown in (10). Alternatively, the spatial parameters \(\mathbf{\theta}\) can also be updated, given the current estimate of \(f\), by maximizing the NNGP log-likelihood

\[\sum_{i}\Big{(}(Y_{i}^{*}(\mathbf{\theta})-O_{i}^{*}(\mathbf{\theta}))^{2}+\log F_{ii} (\mathbf{\theta})\Big{)}. \tag{11}\]

### Spatial predictions

Using the GNN representation, subsequent to estimating \(\widehat{f}(\cdot)\), predictions at a new location can be obtained seamlessly. Given a new location \(s_{0}\) and covariates \(\mathbf{X}_{0}:=\mathbf{X}(s_{0})\), first \(\mathbf{X}_{0}\) is passed through the trained feed-forward part (MLP) of the GNN to obtain the output \(O_{0}=\widehat{f}(\mathbf{X}_{0})\) (see Figure 1). Next, the new location \(s_{0}\) is added as a new node to the DAG, and let \(N(0)\) be its set of \(m\) neighbors on the DAG, \(N^{*}[0]=\{s_{0}\}\cup N(0)\), and
define the graph weights \(\mathbf{v}_{0}\) similar to (9). Then using the graph convolution step, we obtain the decorrelated output \(O_{0}^{*}=\mathbf{v}_{0}^{\top}\mathbf{O}_{N^{*}[0]}\). As the decorrelated output layer \(O^{*}\) is the model for the decorrelated response layer \(Y^{*}\) in the GNN, we have \(\widehat{Y}_{0}^{*}=O_{0}^{*}\). Finally, as \(Y_{0}^{*}=\mathbf{v}_{0}^{\top}\mathbf{Y}_{N^{*}[0]}\) is the graph-convoluted version of \(Y_{0}\), we need to deconvolve \(\widehat{Y}_{0}^{*}\) over the DAG to obtain \(\widehat{Y}_{0}\). This leads to the final prediction equation

\[\widehat{Y}_{0}=\sqrt{F_{0,0}}(\widehat{Y}_{0}^{*}+\mathbf{B}_{0,N(0)}^{\top} \mathbf{Y}_{N(0)}). \tag{12}\]

In the spatial linear model (1), predictions at new locations are performed using kriging. It is easy to verify that the prediction (12) is exactly same as the \(m\)-nearest neighbor kriging predictor for the spatial non-linear model (6), i.e.,

\[\widehat{Y}_{0}=\widehat{f}(\mathbf{X}_{0})+\mathbf{\Sigma}(s_{0},N(0))\mathbf{\Sigma}(N(0 ),N(0))^{-1}(\mathbf{Y}_{N(0)}-\widehat{\mathbf{f}}_{N(0)}). \tag{13}\]

Thus the GNN architecture for NN-GLS offers a simple and coherent way to obtain kriging predictions for the spatial non-linear model. We summarize all the steps of _NN-GLS_ using GNN in Algorithm 1.

## 5 Theory

There is currently no theory on asymptotic properties of neural networks (or their spatial renditions reviewed in the Introduction) under spatial dependence. Even in a non-spatial context theory of neural networks is very much in its nascent stage despite the foundational result of Hornik et al. (1989) showing that NN are universal approximators of any continuous function. Recently, Shen et al. (2019), proved the consistency of one-layer neural networks under OLS loss for i.i.d. data without any dependence.

We prove asymptotic consistency of NN-GLS in estimating the non-linear regression
function in a GP model by minimizing the GLS loss (7). We generalize the framework of Shen et al. (2019) to accommodate the spatially correlated error process, use of a GLS loss, and arbitrary irregular spatial designs. To our knowledge, this is the first asymptotic study of Neural Networks under spatial dependence.
### Notations and Assumptions

We first specify the notations and assumptions for the theoretical study. Let \(\mathbb{R}\) and \(\mathbb{N}\) denote the set of real numbers and natural numbers respectively. \(\|\cdot\|_{p}\) denotes the \(\ell_{p}\) norm for vectors or matrices, \(0<p\leq\infty\). Given the covariates \(\mathbf{X}_{1},\mathbf{X}_{2},\cdots,\mathbf{X}_{n}\), for any function \(f\), we define the norm \(\|f\|_{n}^{2}=\dfrac{1}{n}\sum_{i=1}^{n}f^{2}(X_{i})\). Given a \(n\times n\) matrix \(\mathbf{A}\), \(\lambda(\mathbf{A})\) denotes its eigenspace, \(\lambda_{\max}=\sup\{\lambda(\mathbf{A})\}\), and \(\lambda_{\min}=\inf\{\lambda(\mathbf{A})\}\). A sequence of numbers \(\{a_{n}\}_{n\in\mathbb{N}}\) is \(O(b_{n})\) (\(o(b_{n})\)) if the sequence \(\{|a_{n}/b_{n}|\}_{n\in\mathbb{N}}\) is bounded from above (goes to zero) as \(n\to\infty\). Random variables (distributions) \(X\sim Y\) means \(X\) and \(Y\) have the same distribution.

We first specify Assumptions on the data generation process.

**Assumption 1** (Data generation process).: The data \(Y_{i}:=Y(s_{i}),i=1,\ldots,n\) is generated from a Gaussian Process with a non-linear mean, i.e., \(Y_{i}=f_{0}(\mathbf{X}_{i})+\epsilon_{i}\), where the \(f_{0}(\cdot)\) is a continuous function and the error process \(\{\epsilon_{i}\}\) is a GP such that the maximum eigenvalue of the covariance matrix \(\mathbf{\Sigma}_{n}=Cov(\mathbf{Y})\) is uniformly upper-bounded in \(n\).

Assumption 1 imposes minimal restrictions on the data generation process. The mean function \(f_{0}\) is allowed to be any continuous function. The restriction on the spectral norm of the GP covariance matrix is tied to the spatial design. We show later in Propositions 1 and 2 that this assumption is satisfied for common GP covariance choices for any irregular set of locations in \(\mathbb{R}^{2}\) separated by a minimum distance.

We next state assumptions on the analysis model, i.e., the neural network family and the GLS working precision matrix. We consider a one-layer neural network class:

\[\mathcal{F}_{0}:=\left\{\alpha_{0}+\sum_{i=1}^{n}\alpha_{i}\sigma(\mathbf{W}_ {i}^{\top}\mathbf{X}+\mathbf{w}_{0i}),\ \mathbf{W}_{i}\in\mathbb{R}^{d\times p},\ \mathbf{w}_{0i}\in\mathbb{R}^{d},\ \alpha_{i}\in\mathbb{R}\right\}. \tag{14}\]

This formulation is equivalent to setting \(L=1\), \(g_{1}(\cdot)\) as the sigmoid function \(\sigma(\cdot)\) and \(g_{2}(\cdot)\) as identity function in (3). It is easy to see that \(\mathcal{F}_{0}\) can control the approximation error, as
one-layer NN are universal approximators. However, this class of functions \(\mathcal{F}_{0}\) can be too rich to control the estimation error. A common strategy to circumvent this is to construct a sequence of increasing function classes, also known as sieve, to approximate \(\mathcal{F}_{0}\), i.e.,

\[\mathcal{F}_{1}\subseteq\mathcal{F}_{2}\subseteq\cdots\subseteq\mathcal{F}_{n} \subseteq\mathcal{F}_{n+1}\subseteq\cdots\subseteq\mathcal{F}_{0}.\]

With careful trade-off of the complexity of the function classes, it's possible to control the estimation error (in terms of the covering number of \(\mathcal{F}_{n}\)) using some suitable _uniform-law-of-large-number (ULLN)_ while still being able to keep the approximation error in check. Following Shen et al. (2019) we consider the sieve given below.

**Assumption 2** (Function class).: The mean function \(f\) is modeled to be in the NN class

\[\begin{split}\mathcal{F}_{n}=&\Bigg{\{}\alpha_{0}+ \sum_{j=1}^{r_{n}}\alpha_{j}\sigma(\boldsymbol{\gamma}_{j}^{\top}\boldsymbol{X }+\gamma_{0,j}):\boldsymbol{\gamma}_{j}\in\mathbb{R}^{d},\alpha_{j},\gamma_{0, j}\in\mathbb{R},\\ &\sum_{j=0}^{r_{n}}|\alpha_{j}|\leq V_{n}\text{ for some }V_{n}>1/K,\,\max_{1\leq j \leq r_{n}}\sum_{i,j=0}^{d}|\gamma_{i,j}|\leq M_{n}\text{ for some }M_{n}>0\Bigg{\}},\end{split} \tag{15}\]

where \(r_{n},V_{n},M_{n}\to\infty\) as \(n\to\infty\) and satisfies the scaling

\[r_{n}V_{n}^{2}\log V_{n}r_{n}=o(n). \tag{16}\]

Here the activation function \(\sigma(\cdot)\) is a Lipschitz function on \(\mathbb{R}\) with range \([-r_{b},r_{b}]\) and Lipschitz constant \(K\). (For sigmoid function, \(r_{b}=1\) and \(K=1/4\)).

Hornik et al. (1989) has shown that \(\cup_{n}\mathcal{F}_{n}\) is dense in the continuous function class \(\mathcal{F}_{0}\) which will control the approximation error. The estimation error will depend on the covering number for this class which can be controlled under the scaling rate (16).

Finally, to guarantee regularity of the GLS loss (7) used for estimating the NN function
\(f\), we require the following conditions on the working precision matrix \(\mathbf{Q}\).

**Assumption 3** (Spectral interval for the working precision matrix.).: All eigenvalues \(\lambda\) of \(\mathbf{Q}\) lie in \((M_{low},M_{high})\) for universal constants \(M_{low},M_{high}>0\).

The lower bound ensures that a small GLS loss implies a small estimation error. The upper bound ensures an uniform Lipschitz continuity of the GLS loss function and the consistency of the loss function (Lemma S2) using empirical process results, which serves as a key part of the theory. We show in Propositions 1 and 2 how this Assumption is satisfied when \(\mathbf{Q}\) is either the true GP precision matrix \(\mathbf{\Sigma}_{n}^{-1}\) or its NNGP approximation.

### Main results

We first provide general results on the existence and consistency of neural networks estimators minimizing a GLS loss for dependent data. In Section 5.3 we apply these results to establish consistency of NN-GLS for common choices of GP covariance and spatial designs. The expected value of the GLS loss (7) is given by:

\[\begin{split} L_{n}(f)=\mathbb{E}\big{[}\mathcal{L}_{n}(f)\big{]} &=\mathbb{E}\left[\frac{1}{n}(\boldsymbol{y}-\boldsymbol{f}(x) )^{\top}\mathbf{Q}(\boldsymbol{y}-\boldsymbol{f}(x))\right]\\ &=\frac{1}{n}\mathbb{E}\big{[}\boldsymbol{\epsilon}^{\top}\mathbf{ Q}\boldsymbol{\epsilon}\big{]}+\frac{1}{n}(\boldsymbol{f}_{0}(x)-\boldsymbol{f}(x))^{ \top}\mathbf{Q}(\boldsymbol{f}_{0}(x)-\boldsymbol{f}(x))\end{split} \tag{17}\]

It is evident from above that \(f_{0}\) naturally minimizes \(L_{n}(f)\), while NN-GLS tries to minimize \(\mathcal{L}_{n}(f)\). We first show that such a minimizer exists in the sieve class \(\mathcal{F}_{n}\).

**Theorem 1** (Existence of sieve estimator).: For any \(n\), given data \((Y_{i},\boldsymbol{X}_{i},s_{i}),i=1,\cdots,n\) described in the model (5) and a prefixed working precision matrix \(\mathbf{Q}\), with the function classes \(\mathcal{F}_{n}\) defined in (15), there exists a sieve estimator \(\widehat{f}_{n}\) such that

\[\widehat{f}_{n}=\text{argmin}\{\mathcal{L}_{n}(f):f\in\mathcal{F}_{n}\}. \tag{18}\]
The proof of this and all subsequent theoretical results are provided in the Supplementary materials. The existence result ensures that a sieve estimator in the class of neural networks that minimizes the GLS loss is well defined. It is then natural to study its asymptotic consistency, as we do in the next result.

**Theorem 2** (Consistency).: Consider dependent data generated as \(Y_{i}=f_{0}(\mathbf{X}_{i})+\epsilon(s_{i})\) where \(f_{0}\) is a continuous function, \(\epsilon(\cdot)\) is a Gaussian process. Then under Assumptions 1, 2, 3, the NN-GLS estimate \(\widehat{f}_{n}\) (18) minimizing the GLS loss \(\mathcal{L}_{n}(f)\) in (7) is consistent in the sense \(\|\widehat{f}_{n}-f_{0}\|_{n}\xrightarrow{p}0\).

To our knowledge, this is the first result on consistency of any type of neural network under a dependent error process. We rely on very mild assumptions on the function class, and the covariance matrices of the data generation and analysis models and show in Section 5.3 how these are satisfied for typical GP covariances and irregular spatial data designs. Also, note that this general result does not restrict the nature of the dependence to be spatial. Hence, while spatial applications is the focus of this manuscript, Theorem 2 can be used to establish consistency of neural networks for time-series, spatio-temporal, network, or other types of structured dependence.

Theorem 2 generalizes the analogous consistency result of Shen et al. (2019) from i.i.d. data and OLS loss to dependent error processes and use of a GLS loss. Consequently, the proof addresses challenges that do not arise in the i.i.d. case. The spatial dependency makes the standard Rademacher randomization fail and preventing use the standard Uniform law of large number (ULLN) result. We overcome the difficulty by delving into the construction of normed functional space consisting of random processes. Such a space is the basis of applying a suitable ULLN. We introduce a new Orlicz norm to adjust for data dependence and use of the GLS loss, and are able to construct a ULLN for our dependent setting by showing that the empirical process is well behaved with respect to this Orlicz norm.
### Matern Gaussian processes

We now establish consistency of NN-GLS for common GP covariance families, spatial data designs, and choice of working precision matrices. The main task in applying the general consistency result (Theorem 2) for these specific settings is verifying compliance to the regularity assumptions - i.e., the spectral bounds on the true Gaussian process covariance (Assumption 1) and on the working precision matrix (Assumptions 3).

We provide consistency results of NN-GLS for spatial data generated from the Matern Gaussian process. The Matern covariance function is specified as \(C(\|s_{i}-s_{j}\|_{2})=C(\|h\|_{2})=\sigma^{2}\frac{2^{1-\nu}\big{(}\sqrt{2}\phi \|h\|_{2}\big{)}^{\nu}}{\Gamma(\nu)}\mathcal{K}_{\nu}\big{(}\sqrt{2}\phi\|h\|_ {2}\big{)}\). This is the predominant choice of covariance family in geostatistical models due to the interpretability of its parameters as \(\sigma^{2}\) is the marginal spatial variance, \(\phi\) controls the rate of decay of spatial correlation, \(\nu\) controls the smoothness of the underlying process (Stein, 1999). Closed form representations are available for several special or limiting cases such as the exponential (\(\nu=1/2\)) or Gaussian (\(\nu\to\infty\)) covariance functions. Our first result considers data generated from a class of GP that contains the Matern family and where the working precision matrix is the true GP precision matrix.

**Proposition 1**.: Consider data generated from a spatial process \(Y_{i}=f_{0}(\mathbf{X}_{i})+\epsilon(s_{i})\) at locations \(s_{1},\ldots,s_{n}\) in \(\mathbb{R}^{2}\), where \(f_{0}\) is continuous, \(\epsilon(s_{i})\) is a Gaussian process with covariance function \(\Sigma(s_{i},s_{j})=C(s_{i},s_{j})+\tau^{2}\delta(s_{i}=s_{j})\) with \(\tau^{2}>0\), \(\delta\) is the Kronecker delta, and \(C(s_{i},s_{j})=C(\|s_{i}-s_{j}\|)\) is a covariance of a stationary spatial GP such that \(C(u)=o\big{(}u^{-(2+\kappa)}\big{)}\) for some \(\kappa>0\). Suppose the data locations are separated by a minimum distance \(h>0\), i.e., \(\|s_{i}-s_{j}\|\geq h,\ \forall i\neq j\). Let \(\mathbf{\Sigma}_{n}=\mathbf{C}+\tau^{2}\mathbf{I}\) denote the covariance matrix of \(\mathbf{Y}=\left(Y(s_{1}),\ldots,Y(s_{n})\right)^{\top}\), where \(\mathbf{C}=\left(C(s_{i}-s_{j})\right)\). Then the NN-GLS estimator \(\widehat{f}_{n}=\operatorname{argmin}\{\mathcal{L}_{n}(f):f\in\mathcal{F}_{n}\}\) using \(\mathbf{Q}=\mathbf{\Sigma}_{n}^{-1}\) is consistent in the sense \(\|\widehat{f}_{n}-f_{0}\|_{n}\xrightarrow{p}0\).

Proposition 1 demonstrates how consistency is achieved for NN-GLS in Matern GP models with minimal assumptions on the data generation process. The decay rate on the
spatial covariance \(C(u)=o\big{(}u^{-(2+\kappa)}\big{)}\) is satisfied by the Matern family (Abramowitz and Stegun, 1948). The proposition holds for any irregular spatial design in \(\mathbb{R}^{2}\) meeting the restriction of the locations being separated by a minimum distance. As the sample size \(n\) grows this is essentially equivalent to considering the increasing domain paradigm that is commonly adopted in spatial asymptotics as parameters are not identifiable if data are collected densely in a bounded spatial domain (Zhang, 2004).

Propositions 1 describes the case where true covariance structure is known. In that case, it's possible to directly use the inverse of the covariance matrix as the working precision matrix in the GLS loss. However, this is often infeasible for multiple reasons. First, the true covariance parameters are usually unknown and the working covariance matrix will typically use different (estimated) parameter values. Computationally, GLS loss using the full Matern GP covariance matrix will require \(O(n^{3})\) time and \(O(n^{2})\) storage which are not available even for moderate \(n\). The next proposition introduces a more pragmatic result proving consistency of NN-GLS for data generated from Matern GP but when using a working precision matrix derived from NNGP (as described in Section 4) and with parameter values different from the truth.

**Proposition 2**.: Consider data generated as in Proposition 1 from a Matern GP with parameters \(\mathbf{\theta}_{0}=(\sigma_{0}^{2},\phi_{0},\nu_{0},\tau_{0}^{2})\) at locations separated by a distance \(h>0\). Let \(\mathbf{Q}\) be the NNGP precision matrix based on a Matern covariance with parameters \(\mathbf{\theta}=(\sigma^{2},\phi,\nu,\tau^{2})\) and using neighbor sets of maximum size \(m\) with each location appearing in at most \(M\) many neighbor sets. Then there exists some \(K>0\) such that for any \(\phi>K\), the NN-GLS estimator \(\widehat{f}_{n}=\text{argmin}\{\mathcal{L}_{n}(f):f\in\mathcal{F}_{n}\}\) using \(\mathbf{Q}\) is consistent in the sense \(\|\widehat{f}_{n}-f_{0}\|_{n}\xrightarrow{p}0\).

Proposition 2 proves consistency of NN-GLS for Matern GP when using NNGP working precision matrices. This is the actual choice of \(\mathbf{Q}\) used in our algorithm as it can be represented as a GNN thereby facilitating a scalable implementation, as described in Section
4. The result allows the spatial parameter used in the working covariance to be different from the truth, but requires two assumptions. The spatial decay parameter \(\phi\) for the working precision matrix \(\mathbf{Q}\) needs to be sufficiently large to ensure \(\mathbf{Q}\) is not close to being singular which will lead to numerical instability. We note that this is not a restriction on the data generation process but just on the working precision matrix and can be chosen by the user (although it is not required in practice). Similarly, the restriction on each location appearing in at most a fixed number of neighbor sets, although typically not enforced in NNGP, is usually satisfied in all but very pathological designs.

To our knowledge, Propositions 1 and 2 on consistency of NN-GLS are the first examples of consistency of any machine learning based approach to estimating a non-linear mean of a Matern GP for irregular spatial designs. The only similar result in the literature is consistency of RF-GLS, a GLS based random forest approach of Saha et al. (2023), for Matern GP. However, their result relies on a one-dimensional regular lattice design, restricts the true Matern process smoothness to be half-integers, and needs the covariates \(\mathbf{X}_{i}\) to be i.i.d. Our result is valid for any irregular spatial designs in the two-dimensional space - the most typical setting of spatial data collection. The result also holds for any true parameter values of the Matern process, and do not impose any assumption on the covariates.

## 6 Simulation study

We conduct extensive simulation experiments to study the advantages of NN-GLS over other existing methods in terms of both prediction and estimation. The data are simulated from the non-linear spatial GP model (5) with two choices for the non-linear mean function \(f_{1}(x)=10\sin(\pi x)\) and \(f_{2}(x)=\dfrac{1}{6}(10\sin(\pi x_{1}x_{2})+20(x_{3}-0.5)^{2}+10x_{4}+5x_{5})\) (Friedman function, Friedman, 1991). The error process is an exponential GP with a nugget. We consider
3 choices for each spatial parameter -- variance \(\sigma^{2}\), spatial decay \(\phi\), and error variance (nuggets) ratio \(\tau^{2}/\sigma^{2}\), thereby providing 27 combinations in total. For each setting, we perform 100 independent experiments including data generation, model fitting, and evaluation. In the data-generation process, covariates \(\mathbf{X}\) and coordinates \(s\) are independently sampled from \(\text{Unif}\big{(}[0,1]^{d}\big{)}\) and \(\text{Unif}\big{(}[0,10]^{2}\big{)}\).

We consider 4 candidate neural network approaches for comparison to NN-GLS - NN without spatial consideration (NN-nonspatial), kriging after simple NN (NN-nsp+kriging), NN using the spatial coordinates as two additional inputs (NN-latlon), and NN using spline basis as additional inputs (NN-splines). For NN-splines we use the _Deepkriging_ algorithm of Chen et al. (2020). Among these, NN-GLS and NN-nonspatial are capable for both estimation and prediction, NN-nsp+kriging, NN-latlon, and NN-splines are only designed for prediction.

We first evaluate the estimation performance of NN-GLS and NN-nonspatial in terms of Mean Integrated Square Error (MISE) for the estimate \(\widehat{f}\). Figure 2 (a) gives the comparison of estimation performance for different choice of nugget variance. NN-GLS consistently out

Figure 2: Comparison between methods on estimation of the mean function \(f_{2}\).
performs NN-nonspatial. In addition, Figure 2 (b) shows that NN-GLS has more significant advantage over non-spatial neural networks when \(\phi\) is small. This is expected as for small \(\phi\), there is strong spatial correlation in the data, so the performance of NN-nonspatial suffers on account of ignoring this spatial information. The deterioration in performance of NN-nonspatial over NN-GLS is lesser for large \(\phi\), as there is only weak spatial correlation in the data.

We compare prediction performances using the Relative Mean Squared Error (RMSE) on the test set, obtained by standardizing the MSE by the variance of the response so that the quantity can be compared across different experiments. In Figure 3, we present the prediction results from (a) a correctly specified model with the data generated as (5) using the exponential GP, and (b) a misspecified model for NN-GLS, where the spatial dependence is not from a GP but from a fixed smooth spatial surface (see details in section S2.6). The non-spatial NN, unsurprisingly, offers the worst prediction performance, on account of not using spatial information either in estimation or prediction. However, NN-nsp+kriging, which

Figure 3: (a) Prediction performance comparison when the mean function is \(f_{0}=f_{2}\) and spatial error randomly generated across repetitions; (b) Prediction performance comparison when the mean function is \(f_{0}=f_{1}\) and spatial error generated across from a fixed surface;
uses the estimates of the mean from NN-nonspatial but performs kriging on the residuals, seem to consistently perform better than NN-splines and NN-latlon. This demonstrates the limitation of using just the spatial co-ordinates or some spatial basis functions as additional features in neural networks. The choice and the number of these basis functions may need to be tuned carefully to optimize performance for specific applications. NN-GLS circumvents the need to introduce basis functions by parsimoniously modeling the spatial dependence directly through the GP covariance, as is done in geostatistical linear models, and performs better than these methods under both settings.

Section S2 of the Supplementary materials presents all the other simulation results, including estimation and prediction performance comparison for both choices of functions and all choices of the parameters (Sections S2.1 and S2.2). NN-GLS consistently outperforms NN-nonspatial for estimation of the non-linear covariate effect, and is the best or competitive with the best method for prediction. In Section S2.3, we investigate if the estimation of the spatial parameters has an impact on the performance of NN-GLS by comparing it to _NN-GLS (oracle)_ which uses the true spatial parameters. We observe that NN-GLS's performance is quite the same as that of NN-GLS (oracle), since it provides an accurate estimation of spatial parameters. NN-GLS also performs well for a higher dimensional mean function (of 15 covariates) (Section S2.4). We finally assess robustness of NN-GLS to model misspecification, including misspecification of the GP covariance (Section S2.5), and complete misspecification of the spatial dependence, being not generated from a GP (Section S2.6). For both estimation and prediction, NN-GLS performs as the best or comparably with the best method for both cases of misspecification.

Theoretically, NN-GLS requires linear running time by using NNGP covariance and the GNN framework (see Section 4). Figure 4 plots running times against sample sizes in log-log scale for the following different components of the NN-GLS algorithm: 't-NN-GLS-train' represents the step of neural network's training; 't-spatial-est' represents the step of
estimating spatial parameters and 't-total' gives the total run time. We find that the slopes of the curves are close to 1 (dashed lines), corresponding to linear run-times. This verifies the theoretical computational complexity.

## 7 Real data example

The real data example considered here is spatial analysis of PM\({}_{2.5}\) (fine particulate matter) data in the continental United States. The regulatory PM\({}_{2.5}\) data comes from US environmental protection agency's (EPA) air quality system (AQS). The AQS collects daily air quality data through widely distributed monitors and validates it through a quality assurance procedure. Since PM\({}_{2.5}\) is known to have prominent spatial patterns, we use NN-GLS, modeling the PM\({}_{2.5}\) concentrations across the U.S. as the correlated response, and using other meteorological variables as covariates. The meteorological data are obtained from the National Centers for Environmental Prediction's (NCEP) North American

Figure 4: Log running times for different components of the NN-GLS algorithm as functions of the (log) sample size when the mean functions are (a) \(f_{1}\) and (b) \(f_{2}\). The dashed lines represent linear running time.
Regional Reanalysis (NARR) product, which generates reanalyzed data for temperature, wind, moisture, soil, and dozens of other parameters. with spatial resolution of about \(32\times 32\) km. Following a similar analysis in Chen et al. (2020), we consider the daily PM\({}_{2.5}\) data on June 5th, 2019 for this analysis (other dates are considered in the Supplementary materials). We have PM\({}_{2.5}\) concentration (\(\mu g/m^{3}\)) from 841 stations across the states and six meteorological variables provided at 7706 grid cells from NARR. The six variables are: precipitation accumulation, air temperature, pressure, relative humidity, west and north wind speed. Since the coordinates of the two data sets are different, the spatial resolution of NARR data is retained and the PM\({}_{2.5}\) data are matched onto the grid by averaging in the grid cell. Grid cells without any EPA monitor are removed and there are 604 data points left for the downstream analysis.

Figure 5 (a) demonstrates the PM\({}_{2.5}\) distribution on the date (smoothed by inverse distance weighting interpolation), as well as the nationwide EPA monitor's distribution (the orange spots). The spatial nature of PM\({}_{2.5}\) is evident from the map. We consider the same 5 methods from Section 6: NN-nonspatial, NN-GLS, NN-latlon, NN-nsp+Kriging, and NN-splines. To evaluate the prediction performance, we randomly take 80% of the data as a training set and the rest part as a testing set, train the model and calculate the

Figure 5: PM\({}_{2.5}\) data analysis.
RMSE of the prediction on the testing set. The procedure is repeated for 100 times. The performance of each method is shown in figure 5 (b). We find that NN-GLS has the lowest average RMSE with NN-splines being the next best. This trend is consistent for other choices of dates and for different ways of splitting data into train and test sets (see Section S3 of the Supplementary materials).

As discussed before, unlike most of the other methods which only offer predictions, NN-GLS provides a direct estimate of the effect of the meteorological covariates on PM\({}_{2.5}\), specified through the mean function \(f(\mathbf{X})\).However, \(\mathbf{X}\) is six-dimensional in this application precluding any convenient visualization of the function \(f(\mathbf{X})\). Hence, we use partial dependency plots (PDP) - a common tool in the machine learning world for visualizing the marginal effect of one or two features on the response. In figure 6, we present the PDP of PM\({}_{2.5}\) on temperature and wind. We see clear non-linear effects in both, thereby demonstrating the need to move beyond the traditional linear spatial models and consider geostatistical models with non-linear means, like NN-GLS, for these types of spatial analysis.

Figure 6: Partial dependency plots showing the marginal effects of temperature and vertical (north) wind.
Discussion

In this work, we showed how neural networks can be embedded directly within traditional geostatistical Gaussian process models, creating a hybrid machine learning statistical modeling approach that balances parsimony and flexibility. Compared to existing renditions of NN for spatial data which need to create and curate many spatial basis functions as additional features, NN-GLS simply encodes spatial information explicitly and parsimoniously through the GP covariance. It reaps all benefits of the model-based framework including separation of the non-spatial and spatial structures into the mean and the covariance, use of GLS loss to account for spatial dependence while estimating the non-linear mean using neural networks, and prediction at new locations via kriging.

We show that NN-GLS using Nearest Neighbor Gaussian Process models can be represented as a graph neural network. The resulting GLS loss is equivalent to adding a graph convolution layer to the output layer of a OLS-style neural network. This ensures the computational techniques used in standard NN optimization like mini-batching and back-propagation can be adapted for NN-GLS, resulting in a linear time algorithm. Also, kriging predictions can be obtained entirely using the GNN architecture using graph convolution and deconvolution.

As was expected in our original motivation, NN-GLS should be more efficient than NN, which comes analogously from Gaussian-Markov theorem's result on GLS and OLS for linear models. Both our simulation and real-data studies under various settings coincide with the theoretical expectation and illustrate promising applications of our method.

In theory, we prove general results on the existence and consistency of neural networks using GLS loss for spatial processes. We show how the approach will be consistent for estimating the non-linear mean of data generated from a Matern GP observed over an irregular set of locations. There is a gap between the theoretical setup and the actual
implementation of NN-GLS, as the theory relies on a restricted class of neural networks and does not consider steps like mini-batching and backpropagation used in practice. However, even in a non-spatial context, this gap between the practice and theory of NN is yet to be bridged. To the best of our knowledge, we provide the first theoretical results for any neural network approach under spatial dependency. Even though their assumptions may not be the most general in practice, the theoretical results provide us with a valuable understanding of the algorithm and guarantee its effectiveness in most cases, which is consistent with our observations in simulations and real data analysis.

The connection between NN-GLS and GNN is of independent importance. It demonstrates how GNN, with spatially-informed convolution weights, can be used for irregular spatial data, as it is equivalent to a GP model with neural network mean and NNGP covariance. Part of our future work will be to extend this framework and consider more general graph deconvolution layers, and toother types of irregular spatial data where graphical models are already in use, like multivariate spatial outcomes using inter-variable graphs (Dey et al., 2022) or areal data graphs (Datta et al., 2019). We also primarily focused on modeling the mean as a function of the covariates using a rich non-linear family, i.e., the neural network class, while using stationary covariances to model the spatial structure. However, non-stationarity can be easily accommodated in NN-GLS either by including few lower order (slowly varying) basis functions in the mean or using the GLS loss with a non-stationary covariance function. For example, Zammit-Mangion et al. (2021) proposed rich classes of non-stationary covariance functions using transformations of the space modeled with neural networks. In the future, we will explore extensions of NN-GLS to accommodate such non-stationary GP covariances. We will also develop an efficient, publicly available user-friendly software, to make NN-GLS broadly accessible for analysis of spatial data.
## Acknowledgement

This work was partially supported by National Science Foundation (NSF) Division of Mathematical Sciences grant DMS-1915803 and National Institute of Environmental Health Sciences (NIEHS) grant R01 ES033739.
