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

TL;DR

The authors address the computational bottleneck of integrating measurements over exposure intervals in time-series GP analyses, especially when exposures overlap across instruments. They show that linear Gaussian state-space models provide an exact GP posterior and extend this equivalence by augmenting the state with an integral variable that is reset at exposure starts and observed at exposure ends, achieving exact integrated GP results in O($N$) time on CPU and scalable GPU parallelization to O($N/T$) with additional log factors. The resulting smolgp framework delivers exposure-aware GP modeling with drop-in compatibility to tinygp, and extends efficiently to commonly used covariance kernels including non-quasiseparable ones like the quasiperiodic kernel. Validation demonstrates numerical equivalence to existing integrated GP approaches and dramatic improvements in runtime and memory for large, multi-instrument datasets, enabling robust analyses of massive time-series data 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.

Figures (5)

• Figure 1: Equivalence of linear Gaussian SSMs to GPs. Top: The true signal (in black) is the sum of a SHO and Matérn-5/2 kernel. Synthetic measurements (black points) are noisy samples of this process. The purple '/' band shows the predicted mean and 1$\sigma$ variance given the data from the SSM method; the green '\\' band shows the same for the full GP method. Middle: The same as the top panel but decomposed into each component kernel's prediction at the test points (filled bands) and at the data (points). Bottom: The bottom two panels show the residuals for the predicted mean and variance from the overall model prediction (black line) and component predictions (blue/orange lines) at the test points, and well as at the data (points), between the SSM method and the full GP method. The horizontal dashed lines denote machine precision for 64-bit floating point numbers in the dynamic range of the data; In all cases, the difference between smolgp and tinygp is within this level.
• Figure 2: Numerical validation that the integrated SSM (purple curve) presented in this work produces the same result as the full integrated GP (green curve) from LuhnIntGP. Top: We draw an example stochastic signal (black curve) from a SHO kernel with a $\sim$300 sec timescale (Appendix \ref{['appendix:sho']}). The data points are mock measurements of this curve with $180$ s exposures and uncertainties of 0.3 (arb. units). Middle: the difference between the conditioned (points) and predicted (black curve) mean of the SSM and GP approach. Bottom: The same for the variances. The horizontal dashed lines show machine epsilon for a single 64-bit floating point numbers in the dynamic range of the simulated data; this also corresponds to the linear regime of the y-axis symmetric log scale. The two methods generally agree to within an order of magnitude of this level, consistent with the accumulated floating point error over all the flops in the calculation. Periodicity in the residuals likely stem from the same periodicity in the kernel propagating through the various computations that yield the final mean and variance.
• Figure 3: The same as Figure \ref{['fig:ss_vs_gp_integrated']} but for a mock dataset with two instruments; one with 55 s exposures and one with variable $180 \pm 30$ s. Again, the two methods generally agree to within the expectation from numerical precision.
• Figure 4: Runtime and memory benchmarking comparing the performance on instantaneous data of the full/dense GP solution (i.e. Eq. \ref{['eq:full GP solution']}, as in tinygp, green diamonds), using QSM algebra (also via tinygp, orange squares), and our implementation of the sequential SSM solver (smolgp, blue circles); these three cases were all tested on a CPU, as their performance on a GPU is degraded. Dashed lines trace the theoretical scaling from the largest value tested. In all cases, the top panel shows the wall-clock timed average of five runs, while the bottom shows the peak memory usage during the function execution. The purple stars show the parallel SSM solver (Sarkka2021 as implemented in smolgp) as tested on a NVIDIA RTX 6000 Ada GPU running CUDA v12.8. Left: Results for the log-likelihood as a function of $N$ data points. Middle: Results for conditioning at the $N$ data points, including initialization. Right: Results for conditioning on $N$ data points and then predicting at $M = 100N$ test points, to simulate a typical high-resolution prediction scenario. Takeaway: The SSM shares the linear runtime scaling as QSMs, but is typically more memory efficient (especially for predictions). Computing the likelihood is faster in the QSM framework, although for conditioning the best runtime performance is achieved by the parallel SSM (with high memory overhead on a GPU).
• Figure 5: Same as Figure \ref{['fig:benchmark']} but with integrated measurements. The full/dense GP solution (green diamonds) uses the LuhnIntGP method, which we implemented in tinygp. The SSM solution (blue circles) uses our augmented SSM approach (Section \ref{['sec:integrated_ssm']}), as implemented in smolgp. Likewise, the parallel version (purple circles) implements Section \ref{['sec:parallel']} in smolgp. Takeaway: As there is no quasiseparable framework compatible with integrated measurements, the SSM method provides a transformative speedup and memory savings over traditional GP methods which are forced to construct the full covariance matrix.