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
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 λ , 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, λ decreases and the model starts learning the underlying dynamics, i.e. to represent the angular velocity. Robustness w.r.t. the hyperparameter D0 is demonstrated in App. A.7.1 (Fig. 8 ).
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(z1)=Ep(ζ)[p(z1∣ζ)] that matches the manifold of EpD[q(z1∣x1:T,u1:T)] (cf. first two columns of the latent-space visualisations).
Figure 3: VHP-EKVAE (CO) on reacher (RGB image data). Five-dimensional state-space representation (disentangled): the first three dimensions ( z1,z2,z3 ) represent the two joint angles, the last two dimensions ( z4,z5 ) 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 ( R2 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 {z1(n)∼q(z1∣x1:5(n),u1:4(n))}n=1500 . Note that high ELBO values do not imply the model can accurately predict the observed system.
model
R OLS reg.2(ang. ϕ)
R OLS reg.2(vel. ϕ˙)
MSE predict(smoothed)
MSE predict(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 zt but in the RNN. This is indicated by the low correlation ( R2 ) 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
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
Figure 7: Graphical model of the EKVAE. Red arrows indicate the variational inference networks used for the auxiliary variables at and the VHP variable ζ . Given samples a1:T , the states z1: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. D0 , we measure the correlation between inferred and ground-truth angular velocity by R2 (OLS regression). Here, D0∗ corresponds to the value determined by our heuristic (see App. A.1 ). Note that the abrupt drop at 1.15D0∗ 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.
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 R2 (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.
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.
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.
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.
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.
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.
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.
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.
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.
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 z1∼q(z1∣x1:5,u1:4) or, in case of the EKVAE, on z1∼p(z1∣a1:5,u1:4) , where the auxiliary variables are obtained through a1:5∼q(a1:5∣x1:5) . The average prediction accuracy, measured by the MSE, can be found in Tab. 1a (Sec. 6.2 ).
Figure 20: VHP-EKVAE (CO). Predicted sequence of a moving reacher conditioned on the smoothed distribution z1∼p(z1∣a1:5,u1:4) , where a1:5∼q(a1:5∣x1: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 ht=LSTM(a1:t) and not in zt . This causes the poor smoothing-based predictions, as h1=LSTM(a1) 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 ht=LSTM(z1:t,u1:t) and not in zt . This causes the poor smoothing-based predictions, as h1=LSTM(z1,u1) does not have access to sequence data and therefore cannot infer the angular velocity. See Tab. 2 in Sec. 6.3 .
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 × 16)
1 e -3
Time Steps
15
Actions
1
Auxiliary Variables
2
Latents
3
qϕ(at∣xt)
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 × 16)
1 e -3
Time Steps
15
Actions
1
Latents
3
qϕ(zt∣x1:T,u1:T)
BiLSTM 128. sigmoid activation. FC 64. ReLU activation.
pθ(xt∣zt)
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 × 16)
1 e -3
Time Steps
15
Actions
1
Latents
3
qϕ(zt∣xt:T,ut:T)
LSTM 128. sigmoid activation. FC 64. ReLU activation.
pθ(xt∣zt)
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 × 16)
1 e -3
Time Steps
15
Actions
1
Auxiliary Variables
2
Latents
3
qϕ(at∣xt)
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 × 16)
1 e -3
Time Steps
15
Actions
1
Latents
3
qϕ(zt∣ht−1,xt:T,ut:T)
LSTM 128. sigmoid activation. FC 64. ReLU activation.
pθ(xt∣zt)
FC 128, 128, 128. ReLU activation. Gaussian.
Appendix
Table 7: Model architectures of RSSM. FC refers to fully-connected layers.
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× faster than the fastest competing SMC-based approach.
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
In many domains, practitioners seek models that produce accurate forecasts while faithfully capturing latent system dynamics. Existing approaches typically sacrifice one of these goals: deep state space models often assume Gaussian latent transitions, limiting fit and forecasting, while diffusion models are highly expressive but lack principled inference for the underlying dynamics. To combine the strengths of both, we introduce the Diffusion-Driven State Space Model (DDSSM), which replaces the conventional Gaussian transition distribution with a diffusion model. Our DDSSM resolves the open problem of how to jointly train an autoencoder and a diffusion model on sequential data, thereby extending the literature on latent diffusion models for time series. Moreover, we find that the DDSSM empirically outperforms a state-of-the-art deep SSM at fitting and forecasting a simulated time series with multimodal transitions.
Learning identifiable representations in deep generative models remains a fundamental challenge, particularly for sequential data with regime-switching dynamics. Existing approaches establish identifiability under restrictive assumptions, such as stationarity or limited emission models, and typically rely on variational autoencoder (VAE) estimators, which introduce approximation gaps that limit the recovery of the latent structure. In this work, we address both the theoretical and practical limitations of this setting. First, we establish identifiability of a broad class of recurrent nonlinear switching dynamical systems under flexible assumptions, significantly extending prior results. Second, we introduce ΩSDS, a flow-based estimator that enables exact likelihood optimization using expectation-maximisation. Through empirical validation on both synthetic and real-world data, our results demonstrate that ΩSDS achieves improved disentanglement compared to VAE-based estimators and more accurate forecasting of underlying dynamics.