Modeling nonstationary spatial processes with normalizing flows

Technical Summary: Modeling Nonstationary Spatial Processes with Normalizing Flows

Problem Statement
Nonstationary spatial processes, characterized by statistical properties that vary across space, are prevalent in fields such as environmental science, geostatistics, and remote sensing. Traditional modeling approaches often rely on stationary Gaussian processes (GPs) with fixed covariance functions, which fail to capture complex spatial dependencies found in real-world datasets, potentially leading to biased inference and inaccurate predictions. While existing methods for modeling nonstationarity include spatially varying parameters, nonstationary Matérn covariance functions, and stochastic partial differential equations, they often require restrictive modeling assumptions or tailored constructions for specific problems.

A prominent alternative is the domain warping approach, originally proposed by Sampson and Guttorp (1992), which transforms a nonstationary process on a domain $D$ into a stationary process on a warped domain $D_1$ via an injective mapping $f: D \to D_1$. However, selecting an appropriate warping function is difficult. Previous deep learning-based approaches, such as the deep compositional spatial model by Zammit-Mangion et al. (2022), rely on user-selected warping units (e.g., radial basis functions, Möbius transformations). These methods are limited by the finite set of available units and are currently restricted to one- and two-dimensional spatial domains, hindering their application to higher-dimensional data.

Methodology
The authors propose a novel framework, GPNAF, which models nonstationary spatial processes using Neural Autoregressive Flows (NAFs). The core idea is to learn the injective warping function $f(\cdot)$ (denoted as $T(\cdot)$ in the paper) using a class of invertible neural networks capable of generating complex, high-dimensional transformations.

1. Neural Autoregressive Flows (NAFs): The framework utilizes NAFs, specifically the Deep Dense Sigmoidal Flow (DDSF) architecture. An NAF is constructed as a sequence of triangular maps $T(s; \vartheta)$, where each component of the output depends only on the preceding components of the input. This autoregressive structure ensures the map is invertible and injective, preventing space-folding issues.
2. Architecture: The warping function is parameterized by a neural network where the transformation parameters at each step are outputs of a "conditioner" network. To enforce the autoregressive property (where the $k$-th output depends only on inputs $s_1, \dots, s_{k-1}$), the authors employ binary masking techniques (Germain et al., 2015).
3. Model Formulation: The latent process $Y(\cdot)$ is modeled as a zero-mean Gaussian process on the warped domain $\tilde{D}$ with a stationary and isotropic covariance function $\tilde{C}_\phi$. The observed data $Z$ are noisy measurements of $Y$.
4. Estimation: The model parameters (warping parameters $\vartheta$ and stationary covariance parameters $\phi$) are estimated via a two-stage block-coordinate ascent procedure to maximize the log-likelihood. This involves alternating between optimizing the warping function and the stationary covariance parameters. Spatial prediction is performed using simple kriging on the warped domain.

Key Contributions

• High-Dimensional Flexibility: Unlike previous deep warping models limited to 1D and 2D, the proposed NAF-based framework is applicable to higher-dimensional spatial domains (e.g., 3D) with minimal changes to the underlying architecture.
• Automated Warping Unit Selection: The framework eliminates the need for users to manually select specific warping units (e.g., radial basis functions or Möbius transformations). The neural network learns the appropriate deformation directly from the data.
• Invertibility and Injectivity: By leveraging the theoretical properties of NAFs, the method guarantees an injective mapping, ensuring the warping function does not fold space onto itself, which is critical for valid spatial inference.
• Representational Capacity: The authors demonstrate that NAFs possess greater representational capacity than other commonly used spatial process models, capable of capturing complex anisotropic and nonstationary structures.

Results
The paper validates the GPNAF framework through simulation studies and a real-world application:

1. Simulation Studies:

   • Experiment 1 (SIWGP): Data was generated from a known spatial input warped Gaussian process. GPNAF achieved Mean Squared Prediction Error (MSPE) values closely matching the true model (GPdeepspat) and provided well-calibrated prediction intervals, outperforming stationary GPs, nonstationary Matérn GPs, and DeepKriging.
   • Experiment 2 (Spiral Warping): Data was generated using a complex spiral warping function. GPNAF significantly outperformed all competing methods in MSPE and provided the most calibrated prediction intervals, demonstrating its ability to capture intricate spatial distortions without prior physical knowledge.
   • Visualization: Standard error plots for GPNAF showed uncertainty dominated by spatially varying process properties rather than observation locations, unlike stationary models where uncertainty patterns mirrored observation density.

2. Real-World Application (Argo Floats):

   • The framework was applied to a 3D subset of the Argo Floats dataset (ocean temperature and salinity) in the Atlantic Ocean.
   • GPNAF was compared against stationary and nonstationary Matérn models and DeepKriging. GPNAF achieved the lowest MSPE (0.17) and tightest prediction intervals among the models, while maintaining a 95% prediction interval coverage probability (PICP) close to the nominal level.
   • The learned warping function revealed spatial deformations consistent with ocean surface currents: regions with strong currents (rapid temperature changes) were stretched in the warped domain, while regions with weak currents were compressed. This effectively homogenized the varying dependence structures, allowing a stationary model to be used in the transformed space.

Significance and Claims
The paper claims that GPNAF offers a flexible, data-driven approach to modeling complex, nonstationary, and anisotropic spatial processes. By utilizing normalizing flows, the method overcomes the limitations of previous warping models regarding dimensionality and the need for manual specification of warping units. The authors assert that this framework is particularly valuable for high-dimensional applications in geosciences and environmental sciences where spatial variability is complex and difficult to parameterize a priori. The successful application to 3D oceanographic data highlights the utility of the framework in real-world scenarios involving heterogeneous thermal variations. The authors note that while the current approach is fully data-driven, future work could incorporate physics-guided constraints or extend the framework to spatio-temporal settings.