Neural amortized kriging for scalable Gaussian process inference

Oregon Convention Center 

Spatial statistics often leverages the flexibility and interpretability of Gaussian processes to predict values at unseen spatial locations through Kriging. Unfortunately, determination of Kriging weights relies on the inversion of the process' covariance matrix, creating a computational bottleneck for large spatial datasets. We propose neural amortized Kriging that uses feed-forward neural networks (FFNNs) to learn a mapping from scaled spatial location coordinates and covariance function parameters to Kriging weights and the spatial variance. The FFNNs are trained on synthetic data, and the Vecchia approximation is used to ensure scalability to large spatial datasets. Since the FFNNs do not require a matrix inversion step for predictions, our approach bypasses the bottleneck of Gaussian processes entirely. We demonstrate significant speedup over existing frequentist methods with comparable estimation and prediction errors through simulation studies and using the Jason-3 windspeed dataset.