Scalable Spatiotemporal Prediction with Bayesian Neural Fields
Feras Saad Jacob Burnim Colin Carroll Brian Patton Urs Köster Rif A. Saurous Matthew Hoffman
Spatiotemporal datasets, which consist of spatially-referenced time series,
are ubiquitous in many scientific and business-intelligence applications, such
as air pollution monitoring, disease tracking, and cloud-demand forecasting. As
modern datasets continue to increase in size and complexity, there is a growing
need for new statistical methods that are flexible enough to capture complex
spatiotemporal dynamics and scalable enough to handle large prediction
problems. This work presents the Bayesian Neural Field (BayesNF), a
domain-general statistical model for inferring rich probability distributions
over a spatiotemporal domain, which can be used for data-analysis tasks
including forecasting, interpolation, and variography. BayesNF integrates a
novel deep neural network architecture for high-capacity function estimation
with hierarchical Bayesian inference for robust uncertainty quantification. By
defining the prior through a sequence of smooth differentiable transforms,
posterior inference is conducted on large-scale data using variationally
learned surrogates trained via stochastic gradient descent. We evaluate BayesNF
against prominent statistical and machine-learning baselines, showing
considerable improvements on diverse prediction problems from climate and
public health datasets that contain tens to hundreds of thousands of
measurements. The paper is accompanied with an open-source software package
(https://github.com/google/bayesnf) that is easy-to-use and compatible with
modern GPU and TPU accelerators on the JAX machine learning platform.