Koopman-Equivariant Gaussian Processes

TL;DR

This work introduces Koopman-Equivariant Gaussian Processes (KE-GPs), a probabilistic framework that blends Gaussian process regression with Koopman operator structure to model dynamical systems whose evolution is linear in time but nonlinear in initial conditions. By encoding Koopman equivariance through a trajectory-based symmetrization and a spectral kernel construction, KE-GPs achieve closed-form multi-step trajectory posteriors with quantified uncertainty and improved generalization. The authors provide a theoretical analysis showing reduced sample complexity via information-gain rates, and develop a scalable variational inference scheme with inducing trajectories that leverages the KE structure to avoid time-context inducing points. Empirical results on predator-prey dynamics, robotic demonstrations, and weather data demonstrate competitive or superior forecasting and robust uncertainty quantification compared to kernel-based GP methods and Koopman regression baselines. This framework offers a principled, scalable approach for learning dynamical representations with tractable probabilistic forecasts and representations.

Abstract

Credible forecasting and representation learning of dynamical systems are of ever-increasing importance for reliable decision-making. To that end, we propose a family of Gaussian processes (GP) for dynamical systems with linear time-invariant responses, which are nonlinear only in initial conditions. This linearity allows us to tractably quantify forecasting and representational uncertainty, simultaneously alleviating the challenge of computing the distribution of trajectories from a GP-based dynamical system and enabling a new probabilistic treatment of learning Koopman operator representations. Using a trajectory-based equivariance -- which we refer to as \textit{Koopman equivariance} -- we obtain a GP model with enhanced generalization capabilities. To allow for large-scale regression, we equip our framework with variational inference based on suitable inducing points. Experiments demonstrate on-par and often better forecasting performance compared to kernel-based methods for learning dynamical systems.

Figures (5)

• Figure 1: Backward time equivariance interval (red) and the simulation-induced prediction horizon (green).
• Figure 2: Multi-step mean and 2-sigma interval of the prediction for predator population from the predator-prey dynamics for our proposed Koopman-equivariant GP (KE-GP), a generic contextual kernel (C-GP), and a Koopman operator regression approach (KOR) for noise-free (top) and noisy (bottom) training data.
• Figure 3: Empirical information gain $\hat{\gamma}$ for a 2D linear system scaled to remove effects of constants. The improved rates confirm our theoretical results for Koopman-equivariant GPs, leading to a lower information gain compared to their non-equivariant counterpart (\ref{['eq:SDK']}), even when a randomly sampled eigenvalue spectrum $\{\lambda_j\}_{j=1}^D$ is used instead of the true spectrum.
• Figure 4: Visualization of the GP covariances in space and time. The spatial, KE-GP prior already strongly indicates the shape of the NLL-optimized covariance.
• Figure 5: Ablation Studies: test RMSE with varying data and spectral parameters; we report mean and interquartile range over 10 runs.

Theorems & Definitions (24)

• Definition 3.1: Koopman-equivariance
• Theorem 3.2
• Remark 4.1: Strict complexity reduction
• Theorem 4.2
• Remark B.1: Operator boundedness
• Definition B.2: Non-recurrent domain
• Lemma B.3: Universality of (\ref{['eq:KoopObs']})
• proof
• Remark B.4
• Lemma C.1