Latent Matters: Learning Deep State-Space Models

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

cs.LGFeb 26, 2026

Latent Matters: Learning Deep State-Space Models

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

Organizations: Technical University of Munich · AWS AI Labs · Machine Learning Research Lab, Volkswagen Group, Munich · Eötvös Loránd University

Abstract

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.

Figures & tables

Panel of Figure 1 — Figure 1: VHP-DKS (CO) on pendulum (image data). Learning the state-space representation with the CO framework. Distortion and rate are balanced by the Lagrange multiplier $\lambda$ , which is updated (cf. Alg. 1 ) such that the model first improves the reconstruction quality/constraint by learning the rotation angle (see epoch 70). As soon as the constraint is satisfied, $\lambda$ decreases and the model starts learning the underlying dynamics, i.e. to represent the angular velocity. Robustness w.r.t. the hyperparameter $\mathcal{D}_{0}$ is demonstrated in App. A.7.1 (Fig. 8 ).

Panel of Figure 2 — Figure 2: Pendulum (image data). In contrast to annealing (bottom), CO (middle & top) enables the model to learn the underlying dynamic system, as we verify in Table 1a . Furthermore, the VHP (top) significantly improves the quality of generated sequences. This is because the VHP learns a prior $p(\mathbf{z}_{1})=\mathop{\mathbb{E}_{p(\bm{\zeta})}}\big[p(\mathbf{z}_{1}|\,\bm{\zeta})\big]$ that matches the manifold of $\mathop{\mathbb{E}_{p_{\mathcal{D}}}}\big[q(\mathbf{z}_{1}|\,\mathbf{x}_{1:T},\mathbf{u}_{1:T})\big]$ (cf. first two columns of the latent-space visualisations).

Panel of Figure 3 — Figure 3: VHP-EKVAE (CO) on reacher (RGB image data). Five-dimensional state-space representation (disentangled): the first three dimensions ( $z_{1},z_{2},z_{3}$ ) represent the two joint angles, the last two dimensions ( $z_{4},z_{5}$ ) represent the respective angular velocities. The barrel shape indicates the model has learned that the first joint can do a 360 degree turn; whereas the second joint is restricted to avoid self-collisions (cf. Fig. 6 ).

Table 1: The CO framework significantly improves system identification, as indicated by the high correlation (

R^{2}

of an OLS regression) between inferred and ground-truth states. This leads to an increased prediction accuracy of the model, measured by the MSE of 500 predicted sequences conditioned on

\big\{\mathbf{z}_{1}^{(n)}\sim q(\mathbf{z}_{1}|\,\mathbf{x}_{1:5}^{(n)},\mathbf{u}_{1:4}^{(n)})\big\}_{n=1}^{500}

. Note that high ELBO values do not imply the model can accurately predict the observed system.

model	$R^{2~(\text{ang. }\phi)}_{\text{ OLS reg.}}$	$R^{2~(\text{vel. }\dot{\phi})}_{\text{ OLS reg.}}$	$\text{MSE}_{\text{ predict}}^{~(\text{smoothed})}$	$\text{MSE}_{\text{ predict}}^{~(\text{filtered})}$
VHP-KVAE (CO)	0.989	0.043	2.87E-3	4.24E-4
KVAE (annealing)	0.652	0.134	3.16E-3	6.67E-4
VHP-RSSM (CO)	0.915	0.086	3.10E-3	4.80E-4
RSSM (annealing)	0.158	0.060	3.23E-3	6.94E-4

Table 2: Pendulum (image data), cf. Tab. 1a . As a consequence of the RNN-based transition models, the angular velocity is not encoded in

\mathbf{z}_{t}

but in the RNN. This is indicated by the low correlation (

R^{2}

) with the ground-truth and verified by the low accuracy of the smoothing-based predictions, as we explain in Fig. 2 on the example of the KVAE.

Figure 6: VHP-EKVAE (CO). Policy for reaching an encoded goal position (red dot) that we learned based on the disentangled (position–velocity) state-space representation in Fig. 3 . This verifies that the EKVAE learns an accurate model of the reacher including self-collision avoidance.

Appendix figures & tables22 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 7 — Figure 7: Graphical model of the EKVAE. Red arrows indicate the variational inference networks used for the auxiliary variables $\mathbf{a}_{t}$ and the VHP variable $\bm{\zeta}$ . Given samples $\mathbf{a}_{1:T}$ , the states $\mathbf{z}_{1:T}$ are inferred through (closed-form) Bayesian filtering/smoothing using the locally-linear transition model in Eq. ( 9 ) and the time-invariant auxiliary-variable model in Eq. ( 12 ).

Figure 8: Pendulum (image data). To evaluate the robustness of the CO framework w.r.t. $\mathcal{D}_{0}$ , we measure the correlation between inferred and ground-truth angular velocity by $R^{2}$ (OLS regression). Here, $\mathcal{D}_{0}^{\ast}$ corresponds to the value determined by our heuristic (see App. A.1 ). Note that the abrupt drop at $1.15\,\mathcal{D}_{0}^{\ast}$ is due to the fact that the constraint cannot be satisfied (cf. heuristic), leading to an ill-posed constrained optimisation problem and thus to a poorly trained model.

Panel of Figure 9 — Figure 9: Pendulum (image data). Statistic evaluation of different annealing schedules. For this purpose, we measure the correlation between inferred and ground-truth angular velocity by $R^{2}$ (OLS regression) and compare the best schedule with CO (right), cf. Tab. 1a in Sec. 6.2 . The statistics are based on 25 experimental runs each using a random seed and indicate that CO facilitates system identification.

Panel of Figure 10 — Figure 10: VHP-EKVAE (CO) trained on pendulum image data (supplementary to Tab. 1a in Sec. 6.2 ). In combination with the constrained optimisation framework, the EKVAE identifies the dynamical system of the pendulum and learns to predict it accurately.

Panel of Figure 11 — Figure 11: EKVAE (CO) trained on pendulum image data (supplementary to Tab. 1a in Sec. 6.2 ). Without the VHP, the EKVAE identifies the dynamical system of the pendulum but does not learn to process samples from the prior, which results in a broken generative model.

Panel of Figure 12 — Figure 12: EKVAE (annealing) trained on pendulum image data (supplementary to Tab. 1a in Sec. 6.2 ). Without the constrained optimisation framework, the EKVAE does not learn to accurately predict the observed dynamical system.

Panel of Figure 13 — Figure 13: VHP-DVBS (CO) trained on pendulum image data (supplementary to Tab. 1a in Sec. 6.2 ). In combination with the constrained optimisation framework, DVBS identifies the dynamical system of the pendulum and learns to predict it accurately.

Panel of Figure 14 — Figure 14: DVBS (CO) trained on pendulum image data (supplementary to Tab. 1a in Sec. 6.2 ). The original empirical Bayes prior proposed for DVBS leads to a poorer generative model than the VHP.

Panel of Figure 15 — Figure 15: DVBS (annealing) trained on pendulum image data (supplementary to Tab. 1a in Sec. 6.2 ). Without the constrained optimisation framework, DVBS does not learn to accurately predict the observed dynamical system.

Panel of Figure 16 — Figure 16: VHP-DKS (CO) trained on pendulum image data (supplementary to Tab. 1a in Sec. 6.2 ). In combination with the constrained optimisation framework, DKS identifies the dynamical system of the pendulum and learns to predict it accurately.

Panel of Figure 17 — Figure 17: DKS (CO) trained on pendulum image data (supplementary to Tab. 1a in Sec. 6.2 ). Without the VHP, DKS identifies the dynamical system of the pendulum but does not learn to process samples from the prior, which results in a broken generative model.

Panel of Figure 18 — Figure 18: DKS (annealing) trained on pendulum image data (supplementary to Tab. 1a in Sec. 6.2 ). Without the constrained optimisation framework, DKS does not learn to accurately predict the observed dynamical system.

Figure 19: Summary of all models. Predicted sequences of a moving pendulum conditioned on $\mathbf{z}_{1}\sim q(\mathbf{z}_{1}|\,\mathbf{x}_{1:5},\mathbf{u}_{1:4})$ or, in case of the EKVAE, on $\mathbf{z}_{1}\sim p(\mathbf{z}_{1}|\,\mathbf{a}_{1:5},\mathbf{u}_{1:4})$ , where the auxiliary variables are obtained through $\mathbf{a}_{1:5}\sim q(\mathbf{a}_{1:5}|\,\mathbf{x}_{1:5})$ . The average prediction accuracy, measured by the MSE, can be found in Tab. 1a (Sec. 6.2 ).

Panel of Figure 20 — Figure 20: VHP-EKVAE (CO). Predicted sequence of a moving reacher conditioned on the smoothed distribution $\mathbf{z}_{1}\sim p(\mathbf{z}_{1}|\,\mathbf{a}_{1:5},\mathbf{u}_{1:4})$ , where $\mathbf{a}_{1:5}\sim q(\mathbf{a}_{1:5}|\,\mathbf{x}_{1:5})$ . The average prediction accuracy is depicted in Tab. 1c (Sec. 6.2 ).

Figure 21: VHP-KVAE (CO). The predictions demonstrate that the KVAE encodes the angular velocity of the pendulum in $\mathbf{h}_{t}=\text{LSTM}(\mathbf{a}_{1:t})$ and not in $\mathbf{z}_{t}$ . This causes the poor smoothing-based predictions, as $\mathbf{h}_{1}=\text{LSTM}(\mathbf{a}_{1})$ does not have access to sequence data and therefore cannot infer the angular velocity. See Tab. 2 in Sec. 6.3 .

Figure 22: VHP-RSSM (CO). The predictions demonstrate that the RSSM encodes the angular velocity of the pendulum in $\mathbf{h}_{t}=\text{LSTM}(\mathbf{z}_{1:t},\mathbf{u}_{1:t})$ and not in $\mathbf{z}_{t}$ . This causes the poor smoothing-based predictions, as $\mathbf{h}_{1}=\text{LSTM}(\mathbf{z}_{1},\mathbf{u}_{1})$ does not have access to sequence data and therefore cannot infer the angular velocity. See Tab. 2 in Sec. 6.3 .

Panel of Figure 23 — Figure 23: VHP-EKVAE (CO). Visualisation of the policy learned based on the disentangled (position–velocity) state-space representation in Fig. 3 . For this purpose, the goal position (red dot) is encoded in the latent space. The results show that the EKVAE has learned an accurate model of the observed system. (see Sec. 6.4 )

Dataset	Optimiser	Implementation Details
Pendulum	Adam	Observations	256 (flattened 16 $\times$ 16)
	1 $e$ -3	Time Steps	15
		Actions	1
		Auxiliary Variables	2
		Latents	3
		$q_{\phi}(\mathbf{a}_{t}\|\mathbf{x}_{t})$	FC 128, 128, 128. ReLU activation.

Appendix

Table 3: Model architectures of EKVAE. FC refers to fully-connected layers.

Dataset	Optimiser	Implementation Details
Pendulum	Adam	Observations	256 (flattened 16 $\times$ 16)
	1 $e$ -3	Time Steps	15
		Actions	1
		Latents	3
		$q_{\phi}(\mathbf{z}_{t}\|\mathbf{x}_{1:T},\mathbf{u}_{1:T})$	BiLSTM 128. sigmoid activation. FC 64. ReLU activation.
		$p_{\theta}(\mathbf{x}_{t}\|\mathbf{z}_{t})$	FC 128, 128, 128. ReLU activation. Gaussian.

Appendix

Table 4: Model architectures of DKS. FC refers to fully-connected layers.

Dataset	Optimiser	Implementation Details
Pendulum	Adam	Observations	256 (flattened 16 $\times$ 16)
	1 $e$ -3	Time Steps	15
		Actions	1
		Latents	3
		$q_{\phi}(\mathbf{z}_{t}\|\mathbf{x}_{t:T},\mathbf{u}_{t:T})$	LSTM 128. sigmoid activation. FC 64. ReLU activation.
		$p_{\theta}(\mathbf{x}_{t}\|\mathbf{z}_{t})$	FC 128, 128, 128. ReLU activation. Gaussian.

Appendix

Table 5: Model architectures of DVBS. FC refers to fully-connected layers.

Dataset	Optimiser	Implementation Details
Pendulum	Adam	Observations	256 (flattened 16 $\times$ 16)
	1 $e$ -3	Time Steps	15
		Actions	1
		Auxiliary Variables	2
		Latents	3
		$q_{\phi}(\mathbf{a}_{t}\|\mathbf{x}_{t})$	FC 128, 128, 128. ReLU activation.

Appendix

Table 6: Model architectures of KVAE. FC refers to fully-connected layers.

Dataset	Optimiser	Implementation Details
Pendulum	Adam	Observations	256 (flattened 16 $\times$ 16)
	1 $e$ -3	Time Steps	15
		Actions	1
		Latents	3
		$q_{\phi}(\mathbf{z}_{t}\|h_{t-1},\mathbf{x}_{t:T},\mathbf{u}_{t:T})$	LSTM 128. sigmoid activation. FC 64. ReLU activation.
		$p_{\theta}(\mathbf{x}_{t}\|\mathbf{z}_{t})$	FC 128, 128, 128. ReLU activation. Gaussian.

Appendix

Table 7: Model architectures of RSSM. FC refers to fully-connected layers.

State Space Models

State space models represent sequences through evolving latent states, enabling efficient temporal modeling, system identification, and forecasting.

See all 290 papers

Efficient Learning of Deep State Space Models via Importance Smoothing
May 20, 2026John-Joseph Brady, Nikolas Nusken, Yunpeng LiVariational Inference State Space Models
Centre for Oral, Clinical and Translational Sciences, King’s College London, London, United Kingdom · Department of Mathematics, King’s College London, London, United Kingdom
Diffusion-Driven State Space Models
Jun 19, 2026Jack Ruder, Michael WojnowiczLatent Dynamics Modeling Latent Diffusion Models
1Montana State University
End-to-End Identifiable and Consistent Recurrent Switching Dynamical Systems
May 7, 2026Carles Balsells-Rodas, Zhengrui Xiang, Xavier Sumba +1Dynamical Systems Latent Variable Models
Imperial College London

May 20, 2026J/K move · Enter open · S save

Efficient Learning of Deep State Space Models via Importance Smoothing

John-Joseph Brady, Nikolas Nusken, Yunpeng Li

Centre for Oral, Clinical and Translational Sciences, King’s College London, London, United Kingdom · Department of Mathematics, King’s College London, London, United Kingdom

Variational Inference State Space Models

Open paper AskPDF arXiv

Latent state space systems are ubiquitous in statistical modelling, arising naturally when time series are observed through noisy measurements. However, training deep state space models (DSSMs) at scale remains difficult. Two largely distinct strategies have emerged for training DSSMs. The first, auto-encoding DSSMs, trains generative models by optimising a variational lower bound. The second backpropagates through the outputs of classical sequential Monte Carlo (SMC) algorithms. Such approaches can train DSSMs for both discriminative and generative tasks, but their inherently sequential forward passes scale poorly on modern hardware. We propose parallel variational Monte Carlo (PVMC), a new training method that bridges these paradigms and robustly trains DSSMs for both discriminative and generative tasks. Across a set of benchmark experiments, PVMC matches or exceeds state-of-the-art performance while training

10\times

faster than the fastest competing SMC-based approach.

Abstract

State Space Models

Explore similar work

Efficient Learning of Deep State Space Models via Importance Smoothing

Diffusion-Driven State Space Models

End-to-End Identifiable and Consistent Recurrent Switching Dynamical Systems