Google Research Open-Sources AutoBNN for Probabilistic Time Series Forecasting

Google Research released AutoBNN, an open-source package written in JAX that automates the discovery of interpretable time series forecasting models, according to a blog post from the research team.

The tool combines the interpretability of traditional Bayesian approaches—like Gaussian processes—with the scalability of neural networks. Traditional Bayesian methods often require domain experts to manually select model components and can be computationally expensive, while neural networks lack interpretability and don't produce reliable confidence intervals.

AutoBNN is based on research using Gaussian processes with learned kernel structures. The kernel function encodes assumptions about the data being modeled, such as trends, periodicity, or noise. AutoBNN replaces GPs with Bayesian neural networks (BNNs), which maintain a probability distribution over weights rather than fixed weights, capturing uncertainty in predictions.

The package is available within TensorFlow Probability and implements base kernels (Linear, Quadratic, Periodic, Matérn, ExponentiatedQuadratic) and operators (Addition, Multiplication, ChangePoint). It also includes new features like the OneLayer kernel, LearnableChangePoint operator, and WeightedSum operator for "soft" structure discovery.

BNNs bring advantages over GPs: faster training (approximately linear vs. cubic scaling), better GPU/TPU acceleration, and the ability to combine with traditional deep BNNs for feature discovery.

AutoBNN can be installed via pip: pip install autobnn