A novel stability-enhanced Gaussian process variational autoencoder (SEGP-VAE) is proposed for indirectly training a low-dimensional linear time invariant (LTI) system, using high-dimensional video data. The mean and covariance function of the novel SEGP prior are derived from the definition of an LTI system, enabling the SEGP to capture the indirectly observed latent process using a combined probabilistic and interpretable physical model. The search space of LTI parameters is restricted to the set of semi-contracting systems via a complete and unconstrained parametrisation. As a result, the SEGP-VAE can be trained using unconstrained optimisation algorithms. Furthermore, this parametrisation prevents numerical issues caused by the presence of a non-Hurwitz state matrix. A case study applies SEGP-VAE to a dataset containing videos of spiralling particles. This highlights the benefits of the approach and the application-specific design choices that enabled accurate latent state predictions.
Figures & tables
Fig. 1: SEGP-VAE architecture with semi-contracting kernel.
Fig. 2: KL divergence (unweighted) and L1 norm during training.
SEGP
PEGP
MLP2
NREG
2 -norm
0.02
0.63
0.63
2.53
TABLE I: ∥A^−Ai∥2 where i∈{SEGP,PEGP,MLP2,NREG} .
Fig. 3: Learnt covariance matrices of SEGP-VAE prior (left) & NREG-VAE prior (right). Top left block is the covariance of the radius and bottom right block is the covariance of the angular position. Top right and bottom left are the intra-dimensional covariances.
Fig. 4: Absolute error between the learnt SEGP mean and the angular position trajectory data, averaged over the test set (left). Variance of the learnt SEGP for the angular position, averaged over the test set (right).
Fig. 5: Absolute error between the mean of learnt models and the angular position trajectory data, averaged over the test set (left). The variance of learnt models for the angular position, averaged over the test set (right).
Fig. 6: Randomly sampled latent trajectory from test set, plotted against the SEGP-VAE prior (top) and posterior (bottom). Radius plotted on left and angular position on right.
Fig. 7: Randomly sampled latent trajectory from test set, plotted against the MLP2-VAE prior (top) and posterior (bottom). Radius plotted on left and angular position on right.
Fig. 8: Randomly sampled latent trajectory from test set, plotted against the learnt PEGP-VAE prior (top) and posterior (bottom). Radius plotted on left and angular position on right.
Fig. 9: Randomly sampled latent trajectory from test set, plotted against the learnt NREG-VAE prior (top) and posterior (bottom). Radius plotted on left and angular position on right.
Fig. 10: Randomly sampled videos from test set (top); corresponding latent trajectories (radius along x-axis, angular position along y-axis) and the SEGP posterior mean (middle); SEGP-VAE reconstructed video (bottom).
We evaluate JEPA-style predictive representation learning versus reconstruction-based autoencoders on a controlled "TV-series" linear dynamical system with known latent state and a single noise parameter. While an initial comparison suggests JEPA is markedly more robust to noise, further diagnostics show that autoencoder failures are strongly influenced by asymmetries in objectives and by bottleneck/component-selection effects (confirmed by PCA baselines). Motivated by these findings, we introduce gated predictive autoencoders that learn to select predictable components, mimicking the beneficial feature-selection behavior observed in over-parameterized PCA. On this toy testbed, the proposed gated model is stable across noise levels and matches or outperforms JEPA.
Gaussian processes (GPs) offer appealing properties but are costly to train at scale. Sparse variational GP (SVGP) approximations reduce cost yet still rely on Cholesky decompositions of kernel matrices, ill-suited to low-precision, massively parallel hardware. While one can construct valid variational bounds that rely only on matrix multiplications (matmuls) via an auxiliary matrix parameter, optimising them with off-the-shelf first-order methods is challenging. We make the inverse-free approach practical by proposing a better-conditioned bound and deriving a matmul-only natural-gradient update for the auxiliary parameter, markedly improving stability and convergence. We further provide simple heuristics, such as step-size schedules and stopping criteria, that make the overall optimisation routine fit seamlessly into existing workflows. Across regression and classification benchmarks, we demonstrate that our method 1) serves as a drop-in replacement in SVGP-based models (e.g., deep GPs), 2) recovers similar performance to traditional methods, and 3) can be faster than baselines when well tuned.
Deep state-space models (DSSMs) enable temporal predictions by learning the underlying dynamics of observed sequence data. They are often trained by maximising the evidence lower bound. However, as we show, this does not ensure the model actually learns the underlying dynamics. We therefore propose a constrained optimisation framework as a general approach for training DSSMs. Building upon this, we introduce the extended Kalman VAE (EKVAE), which combines amortised variational inference with classic Bayesian filtering/smoothing to model dynamics more accurately than RNN-based DSSMs. Our results show that the constrained optimisation framework significantly improves system identification and prediction accuracy on the example of established state-of-the-art DSSMs. The EKVAE outperforms previous models w.r.t. prediction accuracy, achieves remarkable results in identifying dynamical systems, and can furthermore successfully learn state-space representations where static and dynamic features are disentangled.
Alexej Klushyn, Richard Kurle, Maximilian Soelch +2
Technical University of Munich · AWS AI Labs · Machine Learning Research Lab, Volkswagen Group, Munich +1