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

Technical Summary: Scalable Gaussian Processes for Integrated and Overlapping Measurements via Augmented State-Space Models

Problem Statement
Astronomical measurements are frequently obtained by integrating signals over finite exposure times (e.g., accumulating photoelectrons in a CCD). When the timescale of the underlying astrophysical variability is comparable to the exposure length, standard instantaneous Gaussian Process (GP) models become biased estimators. Correctly accounting for this integration requires modeling the measurement as an average of the latent process over the exposure interval.

While GP regression is a standard tool for modeling stochastic phenomena (e.g., stellar oscillations in radial velocity data), handling integrated measurements poses significant computational challenges. For non-overlapping, constant exposures, specialized kernels (like those in celerite) can maintain quasi-separability, allowing for $O(N)$ solutions. However, in the general case involving variable exposure times and, crucially, overlapping exposures (common when combining data from multiple instruments with different cadences), the covariance matrix loses its quasi-separable structure. This forces the use of dense matrix operations, resulting in $O(N^2)$ memory and $O(N^3)$ runtime costs. Existing methods for integrated kernels (e.g., Luhn et al. 2026) rely on constructing the full dense covariance matrix, which becomes computationally prohibitive for large datasets ($N \sim 10^5$) typical of modern multi-instrument surveys.

Methodology
The authors propose extending the equivalence between Gaussian Processes and Linear Gaussian State-Space Models (SSMs) to handle integrated and overlapping observations while preserving scalability.

1. Augmented State-Space Formulation:
   The core innovation is the augmentation of the standard SSM state vector. For a latent process $x(t)$ governed by a linear stochastic differential equation (SDE), the authors introduce an integral state $z(t)$ such that $dz/dt = x(t)$.

   • The augmented state vector becomes $\tilde{x} = [x, z]^T$.
   • The system dynamics are updated to include the integration of the latent state.
   • Reset Mechanism: At the start of each exposure $t_s$, the integral state $z$ is reset to zero. At the end of the exposure $t_e$, the observation $y$ is modeled as the normalized integral of the state over the interval: $y = z(t_e)/\delta + \epsilon$, where $\delta$ is the exposure duration.

2. Handling Overlaps:
   To accommodate overlapping exposures from multiple instruments, the framework assigns a distinct integral state $z_i$ to each instrument $i$. The augmented state vector expands to include all instrument-specific integrals. The reset matrices are applied selectively: $z_i$ is reset only at the start of an exposure belonging to instrument $i$, allowing simultaneous accumulation of integrals for overlapping instruments. This construction naturally encodes the cross-correlations required for overlapping data without requiring dense matrix operations.

3. Scalable Inference:
   The authors derive modified Kalman Filter and Rauch–Tung–Striebel (RTS) Smoother algorithms for this augmented system.

   • Sequential Complexity: The filtering and smoothing steps operate on the augmented state dimension ($d + N_{inst}$), maintaining a computational complexity of $O(N)$ on a CPU.
   • Parallel Complexity: Leveraging the associativity of the filtering and smoothing distributions, the method is reframed as an all-prefix-sums problem. This allows for parallel implementation on GPUs, reducing wall-clock time to $O(N/T + \log T)$, where $T$ is the number of parallel workers.

4. Software Implementation:
   The authors introduce smolgp, an open-source Python/JAX package. It offers drop-in compatibility with the existing tinygp library, supporting both standard and exposure-aware GP modeling. It implements the sequential and parallel solvers described above.

Key Results

• Numerical Equivalence: The authors validate the augmented SSM approach against the full dense GP implementation (Luhn et al. 2026) for both single-instrument and multi-instrument (overlapping) scenarios. The results for conditioned means, variances, and log-likelihoods agree to within machine precision (differences $\sim 10^{-14}$ to $10^{-15}$), confirming the theoretical equivalence.
• Performance Benchmarks:
  • Memory: The SSM approach is significantly more memory-efficient than the dense GP method. While dense GPs require $O(N^2)$ memory (exceeding 100 GB for $N \sim 10^5$), the SSM method scales linearly, making large datasets feasible on standard hardware.
  • Runtime: For integrated data, where no quasi-separable (QSM) alternative exists, the SSM method provides a transformative speedup. On a CPU, the sequential SSM solver is the fastest for prediction and competitive for likelihood computation. On a GPU, the parallel SSM solver achieves the fastest conditioning times, rivaling QSM methods for likelihoods.
  • Prediction: The method supports efficient prediction at arbitrary test points using a specific algorithm (Algorithm 1) that is embarrassingly parallel and memory-light compared to dense matrix multiplications.

Significance and Claims
The paper claims that this work enables massive multi-instrument cross-comparisons where exposure overlap is ubiquitous, a scenario previously limited by computational cost. By maintaining $O(N)$ (or $O(\log N)$ on GPUs) complexity, the method unlocks the potential for:

• Analyzing datasets with $N > 10^5$ observations (e.g., daily solar monitoring from multiple instruments like KPF, NEID, EXPRES, and HARPS-N).
• Employing more complex models and higher-dimensional datasets that were previously intractable due to the $O(N^3)$ bottleneck of integrated GPs.
• Extending scalable GP modeling to kernels that lack quasi-separability (e.g., quasiperiodic kernels) when integrated measurements are involved.

The authors emphasize that the framework is not limited to time series but applies to any binned data with a sortable independent coordinate, and it naturally supports derivative observations (e.g., FF' models) due to the inclusion of state derivatives in the SSM formulation. The work positions smolgp as a scalable, flexible tool for next-generation astronomical time-domain analysis.