Stability Enhanced Gaussian Process Variational Autoencoders

1. Problem Statement

The paper addresses the challenge of learning compact, low-dimensional representations of high-dimensional, time-varying data (specifically video) where the underlying physical process is unobserved (latent).

• Context: In many scientific domains (e.g., material science, fluid dynamics), the true state of a system cannot be directly measured, but video recordings of related phenomena are available.
• Limitations of Existing Methods:
  • Standard VAEs: While effective at compression, they often lack interpretability and do not enforce physical laws (like stability), leading to brittle models.
  • Physics-Informed Learning: Existing approaches often require strict prior knowledge of system rules or rely on post-hoc stability checks, which can lead to numerical instability during training (e.g., non-Hurwitz state matrices causing unbounded growth).
• Goal: To develop a framework that learns a latent Linear Time-Invariant (LTI) system from video data, ensuring the learned dynamics are globally stable (specifically semi-contracting) and interpretable, without requiring constrained optimization algorithms.

2. Methodology

The proposed solution is the SEGP-VAE, which integrates a novel Stability-Enhanced Gaussian Process (SEGP) prior into a Variational Autoencoder framework.

A. Theoretical Foundation: SEGP

The core innovation is the derivation of the mean and covariance functions for a GP prior based strictly on the definition of an LTI system:
$$\dot{x}(t) = Ax(t) + Bu(t)$$
$$y(t) = Cx(t) + Du(t)$$

• Derivation: The authors derive analytical expressions for the mean $\mu_y(t)$ and covariance $K_y(t, t')$ of the latent output $y(t)$, accounting for initial conditions $x(0)$ and stochastic inputs $u(t)$.
• Stability Constraint (Semi-contraction): To ensure stability, the system is constrained to be semi-contracting. This means the distance between any two trajectories never increases, regardless of initial conditions. Mathematically, this requires the existence of a positive definite matrix $P$ such that $PA + A^\top P \preceq 0$.
• Unconstrained Parametrization (Theorem 3): A major contribution is the re-parametrization of the state matrix $A$. Instead of optimizing $A$ directly (which risks violating stability constraints), the authors express $A$ in terms of unconstrained variables ($V_1, V_2, V_3$):
  $$A = -\frac{1}{2}P^{-1}V_2 V_2^\top + P^{-1}V_3$$
  where $P = V_1 V_1^\top$. This guarantees that any $A$ generated during training is inherently semi-contracting, preventing numerical explosions (non-Hurwitz issues) and allowing the use of standard unconstrained optimizers (e.g., Adam).

B. SEGP-VAE Architecture

The SEGP serves as the prior distribution $P_\psi(y)$ for the latent states within a VAE:

1. Encoder: A Convolutional Neural Network (CNN) maps video frames to a variational posterior distribution $Q_\phi(y|v)$.
   • Specific Design: For the case study involving rotational motion, the encoder includes a "spatial-to-polar" mapping using Spatial Softmax and a "temporal unwrapping" trick to handle the $2\pi$ discontinuity of angular data, ensuring continuous latent trajectories.
2. Decoder: A Multi-Layer Perceptron (MLP) maps the latent state back to pixel probabilities (Bernoulli distribution for binary video).
3. Training Objective: The model is trained using an augmented Evidence Lower Bound (ELBO):
   $$\mathcal{L} = \mathbb{E}[\log P(v|y)] - \beta \cdot D_{KL}(Q || P_\psi) + \lambda \cdot r(\psi)$$
   • Regularization: An $L_1$ penalty is applied to the state matrix $A$ to encourage sparsity and interpretability.
   • Feature Scaling: The KL divergence is computed in a normalized space to prevent features with larger magnitudes (e.g., angle) from dominating the loss.

3. Key Contributions

1. Novel SEGP Prior: A Gaussian Process prior derived directly from LTI system dynamics, providing a combined probabilistic and interpretable physical model.
2. Stability via Parametrization: A complete, unconstrained parametrization of the LTI state matrix $A$ that guarantees semi-contracting stability. This eliminates the need for constrained optimization and prevents numerical instability caused by non-Hurwitz matrices.
3. Generalization of LFM: Extends Latent Force Models to handle non-zero initial conditions, non-zero mean inputs, and direct feedthrough terms ($D \neq 0$), broadening applicability beyond previous works.
4. Application-Specific Architecture: Demonstrates how to handle rotational dynamics and coordinate discontinuities within the encoder to ensure smooth latent state evolution.

4. Case Study & Results

The method was validated on a dataset of spiralling particles (simulated video of a particle moving in a plane with polar coordinates $r, \theta$).

• Performance:
  • Reconstruction: The model achieved high-fidelity video reconstruction with low per-pixel error.
  • Latent Prediction: The posterior mean of the latent state closely matched the ground truth trajectories with narrow uncertainty bands.
  • Stability: The learned system remained stable throughout training, avoiding the divergence seen in standard GPs.
• Comparison with Standard GP:
  • A standard Squared Exponential (SE) GP failed to capture the correct uncertainty dynamics. It modeled uncertainty as an exponentially decaying function of time separation, whereas the SEGP correctly modeled the integration of input uncertainty over time (for the angle) and the decay of initial condition uncertainty (for the radius).
  • The SEGP learned the uncorrelated nature of radius and angle dynamics naturally, whereas the standard GP required this to be enforced.
• Parameter Recovery: The learned state matrix $A$ had a spectral norm error of only $0.016$ compared to the ground truth, indicating the model successfully recovered the underlying physical dynamics.

5. Significance

• Robust Scientific ML: The paper bridges the gap between data-driven deep learning and physics-based modeling. By embedding stability directly into the model architecture via parametrization, it ensures that learned models are physically plausible and numerically robust.
• Control Design: Since the latent model is a stable LTI system, the output can be directly used for control design, a critical requirement for robotics and autonomous systems.
• Scalability: The use of unconstrained optimization makes the approach scalable and compatible with standard deep learning toolkits, removing the barrier of complex constrained solvers often required for physics-informed learning.

In conclusion, the SEGP-VAE offers a rigorous framework for learning interpretable, stable latent dynamics from high-dimensional video data, successfully overcoming the numerical and theoretical limitations of previous GP-VAE approaches.