PaperlayerView exact revision on arXiv

Source-linked AI summary

Latent Matters: Learning Deep State-Space Models

Alexej Klushyn, Richard Kurle, Maximilian Soelch, Botond Cseke, Patrick van der Smagt

arXiv:2602.23050v2cs.LG

TL;DR

DSSMs are intended to learn latent dynamics from partially observed sequences, but maximising the ELBO alone does not ensure correct system dynamics. The paper proposes constrained optimisation with an empirical Bayes prior and introduces EKVAE, combining amortised inference with extended Kalman filtering and smoothing. Experiments report improved system identification and prediction accuracy, with EKVAE outperforming RNN-based models and learning disentangled position–velocity representations.

  • Problem

    DSSMs often fail to learn the correct dynamics underlying observed data, even though accurate dynamics are needed for prediction and downstream tasks.

  • Method

    The paper introduces constrained optimisation for DSSMs, an empirical Bayes prior, and EKVAE, which combines amortised variational inference with extended Kalman filtering and smoothing.

  • Results

    The proposed framework improves system identification and prediction accuracy, while EKVAE achieves significantly higher prediction accuracy than state-of-the-art RNN-based models.

  • Takeaways & Limitations

    EKVAE can learn disentangled position–velocity state representations that support model-based reinforcement learning reward functions and policy learning.

  • Takeaways & Limitations

    RNN-based transition models can produce non-Markovian state spaces, restricting Bayesian filtering and smoothing and lowering prediction accuracy; the constrained lower-bound guarantee requires 1 ≥ λ ≥ 0.

Abstract

from arXiv · show

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.

1 Introduction

The paper argues that DSSMs often fail to learn the true dynamics despite their intended latent-state representation, and proposes constrained optimisation, an empirical Bayes prior, and EKVAE to address this. Experiments on pendulum and reacher data show improved model learning and prediction accuracy.

  • Motivation: DSSMs model sequence data with nonlinear transition and observation models, but often fail to learn the correct system dynamics needed for accurate prediction and downstream tasks.The observation model maps latent variables to data and accounts for measurement noise, while the transition model is intended to capture the underlying dynamics.
  • Motivation: High sequential ELBO values do not guarantee that a DSSM has learned the underlying system dynamics.The paper identifies this mismatch as one of three main reasons for inaccurate learned models.
  • Contributions: The paper proposes constrained optimisation as a general DSSM training framework, complemented by an empirical Bayes prior.The framework is designed to ensure good reconstruction as a basis for learning system dynamics.
  • Contributions: EKVAE combines extended Kalman filtering and smoothing with amortised variational inference and neural linearisation, dispensing with RNNs.It is also designed to learn state-space representations that disentangle static and dynamic features.
  • Evaluation: Experiments on moving-pendulum image data and the reacher environment show that the proposed approaches improve learned model accuracy and prediction accuracy.The evaluation uses angle and high-dimensional RGB observations, and applying constrained optimisation to established DSSMs produces a substantial prediction-accuracy increase.

2 Background: A Rate–Distortion Perspective on Deep State-Space Models

DSSMs represent observed sequences through lower-dimensional latent states under a Markov assumption, with amortised variational inference learning the latent dynamics. A rate–distortion view separates reconstruction quality from the compression term associated with learning dynamics.

  • DSSM formulation: DSSMs model observed sequence data with typically lower-dimensional latent variables representing the system’s underlying state.The latent sequence is used to represent the dynamics of the observed data.
  • DSSM formulation: The Markov assumption makes the future state and current observation depend solely on the current latent state.This imposes the state-space structure used by the model.
  • DSSM formulation: The model can include optional control signals, while separate parameters govern the observation and transition components.The paper denotes these parameter sets by θ and ψ.
  • Variational inference: Amortised variational inference introduces a recognition model that learns the dynamics jointly with the transition model.The resulting objective is the sequential evidence lower bound.
  • Rate–distortion perspective: The ELBO is decomposed into distortion, which measures reconstruction ability, and rate, which enables learning the underlying dynamics.Balancing these terms can improve DSSM learning.

3 Constrained Optimisation Framework for Improved System Identification

The constrained optimisation framework trains DSSMs by enforcing a reconstruction-quality constraint while optimising the rate, and extends this with a variational hierarchical prior and a two-phase algorithm. The method’s guarantees and optimisation procedure include explicit parameter restrictions and staged training.

  • Motivation: High ELBO values can correspond to different rate–distortion combinations and therefore do not necessarily indicate correct system dynamics.Predefined annealing schedules for balancing reconstruction and compression often fail to achieve the desired result.
  • Constrained optimisation: The framework transfers VAE constrained optimisation to DSSMs by optimising rate subject to the reconstruction constraint D(θ, ϕ) ≤ D0.D0 is a hyperparameter specifying the desired reconstruction-quality baseline.
  • Constrained optimisation: The Lagrange multiplier λ weights distortion, linking the constrained formulation to the sequential ELBO.The paper formulates the sequential ELBO as the Lagrangian of the constrained optimisation problem.
  • Constrained optimisation: The lower-bound guarantee holds only when 1 ≥ λ ≥ 0.This condition is stated for the Lagrangian objective in relation to the marginal log likelihood.
  • Initial distribution: A standard-normal initial prior can over-regularise the approximate posterior or produce a broken generative model when it poorly matches the posterior.The proposed empirical Bayes prior is intended to address this mismatch.
  • Optimisation algorithm: The optimisation algorithm begins by reducing reconstruction error, then learns the prior and transition dynamics after D(θ, ϕ) ≤ D0 is satisfied.This divides training into initial and main phases.

4 Extended Kalman VAE

The EKVAE combines amortised variational inference with Bayesian filtering and smoothing, replacing RNN-based approximations with neural linearisation and auxiliary variables. Its architecture supports analytic inference, sample-based transition optimisation, and disentangled static and dynamic latent features.

  • 4 Extended Kalman VAE: The EKVAE combines amortised variational inference with extended Kalman filtering or smoothing while dispensing with RNNs.It focuses on the smoother version because smoothing can learn a more precise model.
  • 4.1 Neural Linearisation of the Dynamic Model: Neural linearisation learns the dynamic model’s Jacobian as a function of the current state and action instead of computing a Taylor expansion.The transition model is designed to find the best linearisation at each time step globally as a function of state and action.
  • 4.2 Linear Auxiliary-Variable Model: A linear auxiliary-variable model gives high-dimensional observations a low-dimensional representation while preserving analytic posterior computation.The model learns the nonlinear mapping from auxiliary variables to observations with a VAE encoder–decoder pair, while their dependence on latent states remains linear.
  • Disentangling Static and Dynamic Features: The time-invariant auxiliary-variable model uses globally learned or predefined H and R, enabling disentangled state-space representations.With rectangular identity H, the first Da latent dimensions correspond to static features and the remaining Dz − Da dimensions represent dynamic features.
  • 4 Extended Kalman VAE: The EKVAE–CO integration defines distortion and rate from the transition, observation, recognition, and empirical Bayes prior models.The rate uses closed-form computation over zt and Monte Carlo marginalisation over zt−1 to enable sample-based optimisation of transition parameters.

5 Related Work

Related DSSM work commonly uses ELBO-based amortised variational inference and RNNs to approximate Bayesian filtering or smoothing. Alternative empirical-Bayes and auxiliary-variable approaches provide context for the paper’s constrained optimisation framework and EKVAE.

  • Related Work: Annealing modifies the ELBO rate weight β from 0 to 1, but predefined schedules can be suboptimal.The paper extends a VAE-based approach by deriving a Lagrangian formulation of the sequential ELBO from distortion and rate.
  • Related Work: Simple Gaussian prior or initial distributions can produce less accurate transition models, whereas empirical-Bayes methods learn an initial pseudo state from the full sequence.The VHP method can directly substitute a Gaussian p(z1) without the additional restrictions described for another empirical-Bayes approach.
  • Related Work: Amortised variational inference is widely used for DSSMs, with deterministic RNNs commonly replacing or approximating classic Bayesian filtering and smoothing.The paper evaluates and integrates deep Kalman filters/smoothers and deep variational Bayes filters/smoothers in its constrained optimisation framework.
  • Related Work: Earlier models also use hidden states or switch variables to predict parameters of time-inhomogeneous linear state-space models for closed-form Bayesian inference.The Kalman VAE combines auxiliary variables with linear Gaussian observation and transition models whose parameters are conditioned on additional variables.

6 Experiments

Experiments show that constrained optimisation improves DSSM system identification and prediction, while EKVAE learns accurate, disentangled state representations for partially observed dynamical systems.

  • Experimental setup: The evaluation uses pendulum and reacher environments with angle, partially observed state, and image observations, including 64 × 64 pixels RGB images.The pendulum dataset contains 500 sequences of 15 images; the reacher dataset contains 2000 sequences of 30 time steps.
  • CO and DKS: VHP-DKS trained with CO learns pendulum angular velocity after satisfying the reconstruction constraint, unlike annealing.CO first improves reconstruction by learning rotation angle; after the constraint is satisfied, optimisation targets the underlying dynamics.
  • CO and DKS: The CO framework improves system identification and prediction accuracy, evaluated using state correlation and MSE over 500 predicted sequences.High ELBO values alone do not imply accurate prediction, whereas CO substantially increases prediction accuracy, particularly with the VHP.
  • EKVAE: EKVAE outperforms DKS and DVBS in prediction accuracy and identifies the reacher dynamics from 64 × 64 pixels RGB images.Its reacher representation uses five dimensions, separating joint angles from angular velocities.
  • RNN-based transition models: RNN-based transition models can store angular velocity in the recurrent state rather than z_t, reducing smoothing-based prediction accuracy.For KVAE, the initial recurrent state lacks sequence information needed to infer angular velocity; filtered predictions are more accurate than smoothed initial-state predictions.
  • Disentangled representations and policy learning: EKVAE disentangles static and dynamic features, enabling reward-free policy learning for pendulum swing-up and steady rotation.Policies target encoded position or velocity representations without external rewards.

7 Conclusion

The paper proposes constrained optimisation for learning DSSMs and introduces EKVAE, combining extended Kalman filtering/smoothing with amortised variational inference. Experiments show improved system identification and prediction accuracy, while EKVAE supports disentangled representations for reinforcement learning.

  • The constrained optimisation framework is applied to generic DSSMs, extending the VHP empirical Bayes prior and its optimisation algorithm.
  • EKVAE combines extended Kalman filtering/smoothing with amortised variational inference and neural linearisation.
  • Applying constrained optimisation to established DSSMs facilitates system identification and substantially increases prediction accuracy.The evaluation includes models such as DKF/DKS and DVBF/DVBS, with VHP avoiding over-regularisation and broken generative models.
  • EKVAE achieves significantly higher prediction accuracy than state-of-the-art RNN-based models.
  • EKVAE learns disentangled position–velocity representations that can define reward functions and support model-based reinforcement-learning policies.

Funding Transparency Statement

The authors report no related third-party funding or support during the 36 months before submission and no potentially influential financial relationships.

  • No authors received related third-party funding or support during the 36 months before submission.
  • No authors had financial relationships with entities that could potentially be perceived to influence the submitted work during that period.

A.1 Heuristic for Determining D0

The experiments determine the desired reconstruction-quality threshold D0 from the model’s best distortion under classical amortised variational inference.

  • D0 is set heuristically to 0.9 Dmax, where Dmax is the best distortion achieved by the model under classical amortised variational inference.

A.2 Learning the Initial Distribution

The VHP introduces a learnable empirical Bayes prior and an upper bound on the rate, yielding the constrained optimisation Lagrangian used for DSSM training.

  • The VHP defines a VAE-like lower bound on the optimal empirical Bayes prior p*(z1).
  • The VHP formulation introduces an upper bound on the rate.
  • The resulting Lagrangian is L(θ, ψ, ϕ, ψ0, ϕ0; λ) = R(ψ, ϕ, ψ0, ϕ0) + λ(D(θ, ϕ) − D0).
  • This Lagrangian corresponds to the constrained optimisation problem defined for the DSSM.

A.3 Integrating the Deep Kalman Filter and Smoother With the Constrained Optimisation Framework

The constrained-optimisation framework is integrated with deep Kalman and deep variational Bayes filters and smoothers by explicitly specifying their ELBO distortion and rate terms. The integrations replace standard initial-prior constructions with an empirical Bayes prior and retain filtering or smoothing variants through restricted posterior conditioning.

  • Deep Kalman Filter and Smoother: DKF/DKS integration formulates the constrained-optimisation Lagrangian by specifying the distortion and rate underlying the sequential ELBO.The objective is expressed as the ELBO’s negative distortion and rate, with the framework built from Eq. (7).
  • VHP Integration: The VHP-DKS replaces the standard-normal initial prior with an empirical Bayes prior while leaving the DKS distortion unchanged.The filter variant uses posterior conditioning on x1:t and u1:t instead of the full sequence.
  • Deep Variational Bayes Filter and Smoother: DVBF/DVBS integration likewise defines distortion and rate terms for the constrained-optimisation Lagrangian.The DVBS objective is written as the negative of its distortion and rate, and the filter version changes the posterior conditioning to filtering information.
  • VHP-DVBS Integration: For VHP-DVBS, replacing the deterministic initial transformation with the VHP simplifies the marginal posterior and permits a common posterior factorisation across all time steps.The conditional posterior remains implemented as in DVBS, while VHP-DVBF uses filtering rather than smoothing observations.
  • Extended Kalman VAE: The EKVAE combines auxiliary variables, locally linearised dynamics, and Kalman filtering or smoothing to evaluate its variational objective.Its derivation uses closed-form integration for the current latent state and Monte Carlo integration for preceding states and the empirical Bayes variable.
  • Extended Kalman VAE: The EKVAE filter differs from the smoother by conditioning on data only through the current time, enabling sample-based optimisation of transition parameters.The smoother uses pψ(zt|a1:T,u1:T−1), whereas the filter uses pψ(zt|a1:t,u1:t−1).
Loading 2602.23050v2…