Scalable Gaussian Processes for Integrated and Overlapping Measurements Via Augmented State Space Models

TL;DR

This paper introduces a scalable method for Gaussian Process modeling of integrated astronomical measurements using augmented state space models, enabling efficient analysis of overlapping and variable exposure data.

Contribution

It extends the GP-SSM equivalence to handle integrated measurements with scalability, and provides an open-source package compatible with existing tools.

Findings

Achieves O(N) runtime for integrated GP posteriors on CPU.

Parallelizable to O(N/T + log T) on GPU with T workers.

Supports complex kernels like quasiperiodic in astronomy.

Abstract

Astronomical measurements are often integrated over finite exposures, which can obscure latent variability on comparable timescales. Correctly accounting for exposure integration with Gaussian Processes (GPs) in such scenarios is essential but computationally challenging: once exposure times vary or overlap across measurements, the covariance matrix forfeits any quasiseparability, forcing O($N^2$) memory and O($N^3$) runtime costs. Linear Gaussian state space models (SSMs) are equivalent to GPs and have well-known O($N$) solutions via the Kalman filter and RTS smoother. In this work, we extend the GP-SSM equivalence to handle integrated measurements while maintaining scalability by augmenting the SSM with an integral state that resets at exposure start times and is observed at exposure end times. This construction yields exactly the same posterior as a fully integrated GP but in O($N$) time on a CPU, and is parallelizable down to O($N/T + \log T$) on a GPU with $T$ parallel workers. We present smolgp (State space Model for O(Linear/log) GPs), an open-source Python/JAX package offering drop-in compatibiltiy with tinygp while supporting both standard and exposure-aware GP modeling. As SSMs provide a framework for representing general GP kernels via their series expansion, smolgp also brings scalable performance to many commonly used covariance kernels in astronomy that lack quasiseparability, such as the quasiperiodic kernel. The substantial performance boosts at large $N$ will enable massive multi-instrument cross-comparisons where exposure overlap is ubiquitous, and unlocks the potential for analyses with more complex models and/or higher dimensional datasets.