Spatial Bayesian Neural Networks

• SBNNs are probabilistic neural architectures that integrate explicit spatial embeddings and Bayesian priors to model stochastic spatial processes.
• They employ spatial embedding layers using radial basis functions and spatially varying priors for network weights to capture spatial heterogeneity.
• SBNNs use calibration techniques like Wasserstein optimization and auxiliary networks to ensure accurate uncertainty quantification and predictive performance.

Spatial Bayesian Neural Networks (SBNNs) are a class of probabilistic neural architectures that explicitly incorporate spatial structure into Bayesian neural networks to model stochastic spatial (and often spatio-temporal) processes. SBNNs provide a framework for learning spatial fields, capturing spatial heterogeneity, and quantifying predictive uncertainty by placing Bayesian priors on neural network weights, introducing spatial embeddings, and calibrating the model against the finite-dimensional properties of spatial data or simulators. The versatility of SBNNs allows them to approximate classic spatial models (e.g., Gaussian processes, max-stable processes), generalize to irregular spatial domains, and scale to real-world applications in geophysics, climate, and beyond (Zammit-Mangion et al., 2023, McDermott et al., 2017, Sainsbury-Dale et al., 2023, Aich et al., 29 May 2025).

1. Mathematical and Structural Foundations

A canonical SBNN comprises a neural random field $Y(s) = f(s; \theta)$ over a spatial domain $D \subset \mathbb{R}^d$, where $\theta$ denotes the (random) weights and biases drawn from a specified prior $p(\theta)$ (Zammit-Mangion et al., 2023). The network input $s$ is typically passed through a spatial embedding layer $\phi(s)$, constructed using a basis of radial basis functions (RBFs) centered at coordinates $\{\xi_k\}$, with

$$\phi_k(s; \tau) = \exp(-\|s - \xi_k\|^2 / \tau^2), \quad k=1,\dots,K,$$

providing a spatial encoding that captures locality and facilitates smooth spatial dependence (Zammit-Mangion et al., 2023, Aich et al., 29 May 2025).

After embedding, the resultant features are processed through multiple hidden layers, each with Bayesian weight priors, potentially spatially varying:

$$\theta_{l, i} \sim \mathcal{N}(\mu_{l,i}, \sigma_{l,i}^2),$$

where $\mu_{l,i}, \sigma_{l,i}$ can be either fixed or functions of $s$ (parameterized via the spatial basis) (Zammit-Mangion et al., 2023). The overall architecture provides a generative model for spatial fields, where each realization of $\theta$ yields a spatial process $Y(s)$ with tractable finitedimensional marginals.

Extensions include spatially structured priors (e.g., GMRFs on spatial intercepts), copula-driven initialization for non-Gaussian dependence (Aich et al., 29 May 2025), embedding-lagged states for spatio-temporal models (McDermott et al., 2017), and graph-based architectures for irregular spatial arrangements (Sainsbury-Dale et al., 2023).

2. Spatial Bayesian Hierarchies and Prior Specification

Prior choices in SBNNs are core to capturing spatial structure and regularizing the highdimensional parameter space. Several paradigms emerge:

• Standard Gaussian Priors are placed on weights/biases for tractable propagation and as a universal starting point. In "prior-per-layer" (IL) and "prior-per-parameter" (IP) variants, hyperparameters are either shared across layers or learned individually for each parameter, with spatially varying generalizations (VL, VP) in which prior mean and variance are expanded on the embedding basis: $\mu_l(s) = \alpha_l^T \phi(s)$, $\sigma_l(s) = \mathrm{softplus}(\beta_l^T \phi(s))$ (Zammit-Mangion et al., 2023).
• Copula-based Priors arise when dependence structure outside the Gaussian field is essential. The A2 Copula SBNN introduces an Archimedean copula (generator inverse: $\varphi_{A_2}^{-1}(t; \theta)$) to jointly initialize weights in fully connected layers, allowing dual-tail (upper and lower extreme) dependence, as encoded by copulas, to propagate from weights to outputs (Aich et al., 29 May 2025).
• Stochastic Search Variable Selection (SSVS) Priors implement sparsity by mixing truncated-normal distributions with binary indicators, particularly suited for high-dimensional RNN weight matrices and for feature selection (McDermott et al., 2017).
• Graph-based Spatial Representation bypasses explicit spatial coordinates using message passing on graphs (e.g., GNNs with edge distance weighting), making SBNNs applicable to arbitrary or irregular spatial supports (Sainsbury-Dale et al., 2023).

These prior formulations are crucial in ensuring that the Bayesian machinery can learn complex spatial dependence, regularize overfitting, and inject interpretable domain constraints.

3. Calibration and Inference Methodologies

Calibration in SBNNs seeks to align the finite-dimensional distributions induced by the network with the empirical or simulator-driven target process. The dominant approach is Wasserstein (optimal transport) calibration, which minimizes

$$W_1(\psi) = \sup_{\varphi: \mathrm{Lip}(\varphi)\leq 1} \Big[ \mathbb{E}_\psi \varphi(Y) - \mathbb{E} \varphi(\tilde X) \Big],$$

where $\psi$ are the prior hyperparameters, and $\varphi$ is parameterized by a neural network "critic" with a gradient penalty to enforce Lipschitzness (Zammit-Mangion et al., 2023). Two-stage adversarial optimization iterates inner (critic) and outer (SBNN prior) loops until empirical Wasserstein distance converges.

Training objectives may combine mean-squared error, Wasserstein loss, moment matching, and correlation penalties (as in A2-SBNN), to achieve distributional, marginal, and spatial alignment with the target field (Aich et al., 29 May 2025). For SBNNs on irregular graphs, the loss is the empirical Bayes risk (e.g., mean-absolute-error for point estimation, quantile loss for marginal intervals), and calibration is amortized over multiple spatial configurations (Sainsbury-Dale et al., 2023).

Posterior inference, once calibration is complete, proceeds via Hamiltonian Monte Carlo, stochastic-gradient HMC, or variational Bayes on the restricted posterior

$$p(\theta \mid Z, \psi^*) \propto p(Z \mid \theta) \, p(\theta \mid \psi^*),$$

where $Z$ are noisy field observations. This allows sampling-based prediction and uncertainty quantification (Zammit-Mangion et al., 2023).

4. Uncertainty Quantification and Predictive Analytics

SBNNs explicitly decompose and propagate predictive uncertainty. The predictive variance for future values is decomposed into aleatoric (data-level) and epistemic (model-level) sources:

$$\mathrm{Var}(y_{T+1} \mid \mathbf{y}_{1:T}) = \mathbb{E}_{\theta \mid y}\left[ \mathrm{Var}(y_{T+1} \mid \theta) \right] + \mathrm{Var}_{\theta \mid y}\left[ \mathbb{E}(y_{T+1} \mid \theta) \right],$$

enabling full Bayesian uncertainty intervals (McDermott et al., 2017). In practice, predictive distributions are constructed via Monte Carlo draws of $(\theta^{(s)}, h^{(s)})$ and sampling new field values.

For complex models (e.g., GNN-based SBNNs), uncertainty is quantified by training auxiliary networks to produce credible intervals—for each parameter, lower and upper quantiles are regressed using quantile loss, with coverage empirically validated (Sainsbury-Dale et al., 2023). The approach is extended to process-level uncertainty, capturing rare spatial extremes and co-movements when copula-based priors are invoked (Aich et al., 29 May 2025). MC-Dropout and spatial dropout variants provide efficient empirical uncertainty quantification, especially for hardware-optimized or binary SBNNs (Ahmed et al., 2023).

5. Extensions to Irregular, Non-Gaussian, and Hardware-efficient Domains

SBNNs generalize beyond regular grids and Gaussian dependence. Graph-based SBNNs use message passing over edge-weighted spatial proximity graphs, achieving invariance to permutations, variable sample sizes, and unseen spatial networks (Sainsbury-Dale et al., 2023). Copula-driven SBNNs, such as A2-SBNN, directly encode and calibrate dual-tail spatial dependencies, outperforming standard Gaussian SBNNs under heavy-tailed and extreme co-movement settings (Aich et al., 29 May 2025).

Binary SBNNs with spatial dropout (Spatial-SpinDrop) use per-feature-map dropout masks generated by the intrinsic stochasticity of spintronic memory (STT-MRAM), drastically reducing inference energy and hardware overhead compared to per-activation dropout. Empirical results show up to $9\times$ reduction in hardware modules and $94\times$ improvement in energy efficiency, while maintaining uncertainty calibration and predictive accuracy within $1\%$ of floating-point baselines (Ahmed et al., 2023).

6. Applications, Advantages, and Limitations

SBNNs are deployed in surrogate modeling for computationally expensive simulators (e.g., geophysical or climate fields), uncertainty-aware inference in data-driven non-stationary settings, and high-throughput parameter estimation for spatial processes (Gaussian, max-stable, lognormal) (Zammit-Mangion et al., 2023, Sainsbury-Dale et al., 2023). Key advantages include:

• Universal process approximation for stationary, nonstationary, non-Gaussian, and extreme-value spatial targets.
• Decoupling model calibration from explicit parametric selection of spatial covariance or dependence class.
• Broad scalability via GPU acceleration, amortized training, and inference variants suitable for large spatial domains.
• Efficient uncertainty quantification, including for hardware (edge or IoT) deployments.

Limitations are mainly the need for extensive independent replicates or simulations for calibration, high hyperparameter count and storage cost in fine-grained spatial prior variants, interpretability bottlenecks relative to mechanistic models, and, in some architectures, limited support for irregular or dynamic spatial supports (although GNN-based and copula-based approaches remove this restriction) (Zammit-Mangion et al., 2023, Aich et al., 29 May 2025). Extending SBNNs to spatio-temporal graphs, censoring/extremal processes, missing data, and richer copula structures remain active directions.

7. Comparative Performance and Empirical Benchmarks

Empirical evaluations demonstrate that SBNN variants surpass non-spatial BNNs in matching target spatial field distributions, with the advantage magnified for non-stationary or non-Gaussian data. For instance, calibrated SBNNs recover Gaussian process covariances and marginals, accurately represent lognormal and extreme-value processes, and outperform echo-state network ensembles and composite likelihood in out-of-sample mean squared prediction error (MSPE), CRPS, and credible interval coverage (McDermott et al., 2017, Zammit-Mangion et al., 2023, Sainsbury-Dale et al., 2023, Aich et al., 29 May 2025).

The table below summarizes representative performance improvements from SBNN architectures:

Task/Model	SBNN Metric (Best)	Baseline (Best)
GP Parameter Estimation	RMSE: 0.050, Coverage: 94%	ML: RMSE 0.049
Max-stable: Schlather	RMSE: 0.07	PL: RMSE 0.36
Extreme Dependence (θ=9)	RMSE: 0.076, Corr: 0.94	Gaussian SBNN: 0.095, 0.91
Hardware (LeNet-5, MNIST)	Energy: 0.68μJ/image	RRAM: 9.3μJ, FPGA: 18.97μJ

These results substantiate the efficiency, flexibility, and calibration capability of SBNN models across methodological and applied axes.

---

References:

• "Spatial Bayesian Neural Networks" (Zammit-Mangion et al., 2023)
• "Bayesian Recurrent Neural Network Models for Forecasting and Quantifying Uncertainty in Spatial-Temporal Data" (McDermott et al., 2017)
• "Neural Bayes Estimators for Irregular Spatial Data using Graph Neural Networks" (Sainsbury-Dale et al., 2023)
• "A2 Copula-Driven Spatial Bayesian Neural Network For Modeling Non-Gaussian Dependence: A Simulation Study" (Aich et al., 29 May 2025)
• "Spatial-SpinDrop: Spatial Dropout-based Binary Bayesian Neural Network with Spintronics Implementation" (Ahmed et al., 2023)