Diffusion-Driven State Space Models

Abstract: 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.

• The paper introduces DDSSM by replacing Gaussian transitions with a diffusion-based transition model to capture non-Gaussian, multimodal dynamics.
• The paper integrates end-to-end training for both the autoencoder and diffusion processes, enhancing latent inference and forecasting performance.
• The paper demonstrates superior predictive accuracy and reconstruction compared to deep Kalman filters, validated by lower Jensen-Shannon Divergence.

Diffusion-Driven State Space Models: Integrating Diffusion and Deep State Space Approaches

Motivation and Formulation

The paper "Diffusion-Driven State Space Models" (2606.21036) proposes a new latent dynamical generative model that unifies diffusion models and deep state-space models (DSSMs) for sequential data. Conventional DSSMs, such as deep Kalman filters, parameterize latent transitions with Gaussian distributions, which limits the expressivity of the latent dynamics and frequently leads to over-regularized latent spaces. Meanwhile, diffusion models exhibit high expressivity in generative modeling but lack principled inference with respect to underlying system dynamics, especially for sequential data. Existing latent diffusion frameworks for time series typically train autoencoders and diffusion models in a two-stage manner, resulting in suboptimal latent representations that hinder downstream generative performance.

The Diffusion-Driven State Space Model (DDSSM) addresses these limitations by:

1. Replacing the Gaussian transition in DSSMs with a diffusion-based transition model, inducing a nonparametric transition prior in latent space.
2. Enabling joint, end-to-end training of the autoencoder and diffusion components for sequential data—resolving the challenge of integrating latent diffusion and principled inference in time series modeling. This formulation results in a highly flexible latent process capable of modeling multi-modal and non-Gaussian transitions, leveraging the inductive biases of state space models together with the expressive power of diffusion-based generation.

Model Architecture and Variational Inference

Generative Process

The generative process of DDSSM posits observed sequences $x_{1:T}$ produced from latent trajectories $z_{1:T}$, modulated by covariates $u_{1:T}$. The joint probability is factorized using state-space assumptions—including Markov order $j$—as:

$$p(x_{1:T}, z_{1:T}|u_{1:T}) = p_{\eta, \theta}(z_{1:j}, x_{1:j}|u_{1:j}) \prod_{t=j+1}^T p_\psi(z_t|z_{t-j:t-1}, u_{t-j:t}) p_\theta(x_t|z_{t-j+1:t}, u_{t-j+1:t})$$

Critically, the transition density $p_\psi(z_t|z_{t-j:t-1}, u_{t-j:t})$ is parameterized as a diffusion model—introducing intermediate latent states at each transition step per time $t$, which allows the model to capture complex, multimodal transition dynamics not possible with Gaussian distributions. The emission model remains Gaussian, with arbitrary expressivity via neural networks.

Inference and ELBO Derivation

Exact posterior inference is intractable, so the model resorts to variational inference. The variational posterior is factorized to mirror the true posterior's conditional independence structure—a practice that provides strong inductive bias and narrows the ELBO approximation gap. The variational posterior includes Kalman smoothing factors, with future observations influencing the latent state at each timestep, and leverages convolutional or GRU-based neural architectures for summarizing history and future context. Emphasis is given to joint training of both the encoder (inference network) and the diffusion transition model, resulting in latent states that are both reconstructable and consistent with learned dynamics.

The ELBO objective incorporates:

• Initialization KL (hierarchical prior)
• Reconstruction error (from emissions)
• Diffusion chain loss (KL between forward and reverse diffusion processes)
• Transition KL (between variational posterior and transition prior)

Preconditioning strategies for gradient normalization across noise levels, from Karras et al. (2022), are adopted to stabilize training.

Comparison to Prior Work

DDSSM improves on deep Kalman filters (DKF), DVAEs, and recent latent diffusion approaches for time series. Unlike previous latent diffusion methods—which rely on two-stage training and regularize the latent space towards a Gaussian distribution—DDSSM enables end-to-end training, jointly regularizing the latent space according to the learned diffusion transition model. This principal innovation directly tackles the observed phenomenon that two-stage approaches discard information necessary for predictive and generative tasks, confirmed both in this paper and prior literature (Liu et al., 2024; Vahdat et al., 2021; Shmakov et al., 2023).

Normalizing flow-based SSMs offer exact likelihoods but are constrained by invertibility requirements, making DDSSM preferable for many dynamical domains. Models that apply diffusion directly in observation space (Rasul et al., 2021; Tashiro et al., 2021) suffer from prohibitive computational cost and limited temporal flexibility, amplifying the incentive for latent space diffusion.

Empirical Evaluation and Numerical Results

A synthetic simulation study demonstrates the practical advantages of flexible transition modeling with DDSSM. On a latent random walk with bimodal noise, DDSSM produces predictive distributions that closely match the analytic ground truth, whereas DKF (Gaussian transitions) places mass between modes and fails to capture multi-modality. Quantitatively, DDSSM achieves significantly lower Jensen-Shannon Divergence (JSD) to the true predictive distribution compared to DKF and simple last-observation baselines.

DDSSM also produces more accurate reconstructions and forecasts for observed data, demonstrating trajectories that align with ground truth, unlike DKF, whose forecasts do not exhibit coherent patterns. The evidence strongly supports the contention that Gaussian priors in SSMs over-regularize latent trajectories, degrading both reconstruction and prediction performance.

Practical and Theoretical Implications

The integration of diffusion processes into state-space models fundamentally expands the class of latent transition distributions available to researchers—enabling precise modeling of complex, multimodal, and non-Gaussian dynamical phenomena. This is particularly relevant in physics, finance, and neural modeling domains, where latent dynamics rarely conform to simple parametric forms.

From a theoretical perspective, DDSSM advances the literature on sequential latent diffusion by establishing a viable, principled pathway for end-to-end training. Furthermore, the approach creates a bridge for future enhancements—such as incorporation of extended Kalman filter techniques, normalizing flows, alternative emission models, and application to large-scale real-world datasets.

Speculation on Future Directions

DDSSM's demonstrated empirical superiority and architectural flexibility suggest several future research avenues:

• Scaling to high-dimensional and long-horizon sequential data, including video, weather, and spatiotemporal forecasting.
• Incorporating meta-learning, online adaptation, and interactive inference mechanisms, exploiting the model's non-Gaussian transition capacity.
• Exploring alternative diffusion schedules, time-homogeneous transition parameterizations, and structured state-space models for specialized domains.
• Integrating DDSSM with differentiable particle filtering and normalizing flow-based inference for further expressivity.
• Investigating interpretability and uncertainty quantification in latent trajectories inferred via diffusion transitions.

Conclusion

The Diffusion-Driven State Space Model (DDSSM) introduces a unified framework for flexible, expressive latent dynamical modeling in sequential data. By replacing Gaussian transitions with diffusion models and enabling joint, end-to-end training, DDSSM avoids the pitfalls of over-regularized latent spaces and limited expressivity that have hindered previous deep state-space and latent diffusion models. Empirical results validate the approach's ability to recover multi-modal transition distributions and provide superior forecasting performance. The theoretical framework, architecture, and training methodology open new possibilities for future developments in sequential generative modeling and probabilistic forecasting.