Technical Summary: Diffusion-Driven State Space Models (DDSSM)

Problem Statement

Practitioners in time-series modeling often face a trade-off between model expressivity and the ability to capture latent system dynamics. Existing approaches typically sacrifice one goal for the other:

• Deep State Space Models (DSSMs) (e.g., Deep Kalman Filters) offer principled inference and end-to-end training but typically assume Gaussian transition densities. This assumption limits their expressivity, often failing to capture multimodal dynamics, and induces a Gaussian prior that can over-regularize latent representations, discarding necessary structural information.
• Diffusion Models are highly expressive and capable of modeling complex, multimodal distributions. However, when applied to time series, they are often used as generative samplers conditioned on past observations or rely on two-stage training (first learning an autoencoder, then a diffusion model). This two-stage approach can produce over-regularized latent spaces that discard information required for downstream generative dynamics.

There remains an open problem regarding how to jointly train an autoencoder and a diffusion model on sequential data in a principled manner to combine the strengths of both frameworks.

Methodology

The authors propose the Diffusion-Driven State Space Model (DDSSM), which integrates diffusion models directly into the transition density of a State Space Model (SSM).

Generative Model

The model assumes observations $x_{1:T}$ are generated from latent states $z_{1:T}$ evolving via a Markov process, potentially conditioned on covariates $u_{1:T}$. The joint likelihood is factorized 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^{(t)}_\psi(z_t | z_{t-j:t-1}, u_{t-j:t}) p^{(t)}_\theta(x_t | z_{t-j+1:t}, u_{t-j+1:t})$$

Key innovations include:

1. Diffusion Transitions: Instead of a parametric Gaussian transition $p^{(t)}_\psi$, the model uses a diffusion process. For each time step $t$, the transition density is modeled by introducing intermediate noisy latent states $z^1_t, \dots, z^K_t$. The transition is defined as:
   $$p^{(t)}_\psi(z^0_t, z^{1:K}_t | z_{t-j:t-1}, u_{t-j:t}) = p(z^K_t) \prod_{k=1}^K p^{(t)}_\psi(z^{k-1}_t | z^k_t, z_{t-j:t-1}, u_{t-j:t})$$
   where $z^0_t = z_t$ and $z^K_t \sim \mathcal{N}(0, I)$. This allows the model to learn arbitrary, potentially multimodal transition densities.
2. Emissions: The emission density $p^{(t)}_\theta$ remains a standard Gaussian parameterized by neural networks, handling residual uncertainty while the diffusion model captures complex non-Gaussian structures.
3. Initialization: The initial state distribution $p_\eta(z_{1:j}|u_{1:j})$ utilizes a Variational Hierarchical Prior (VHP) to avoid restrictive assumptions on the initial latent states.

Inference and Training

Since the true posterior is intractable, the authors employ a variational inference framework. They derive an Evidence Lower Bound (ELBO) that mirrors the conditional independence structure of the true posterior.

• Variational Posterior: The posterior is factorized to include a smoothed transition term (dependent on past latents and future observations) and a diffusion noising process term.
• Joint Training: Unlike two-stage methods, the DDSSM is trained end-to-end. The objective maximizes the ELBO, which includes:
  • A reconstruction loss (negative log-likelihood of emissions).
  • A KL divergence term for initialization.
  • A diffusion chain loss (derived from the reverse process of the diffusion model) that encourages the encoder to produce latent states predictable under the learned transition dynamics.
• Preconditioning: To ensure gradient stability across noise levels during joint training, the authors employ preconditioning (Karras et al., 2022) to re-scale inputs and outputs of the diffusion network.

Key Contributions

1. DDSSM Framework: The introduction of a State Space Model where the transition density is replaced by a diffusion model, enabling the learning of flexible, multimodal latent dynamics without restrictive Gaussian assumptions.
2. End-to-End Training: A principled solution to the open problem of jointly training an autoencoder and a diffusion model on sequential data. This avoids the over-regularization issues associated with two-stage training approaches common in latent time-series diffusion.
3. Theoretical Derivation: A rigorous derivation of the ELBO for this hybrid model, demonstrating how the optimal variational posterior factorizes to match the true posterior's structure, thereby narrowing the approximation gap.

Experimental Results

The authors evaluated the DDSSM on a simulated time-series dataset featuring bimodal noise (a latent random walk with bimodal transitions).

• Setup: They compared the DDSSM against a Deep Kalman Filter (DKF) with a Variational Hierarchical Prior and a Last Observation Carried Forward (LOCF) baseline.
• Metrics: Performance was measured using Jensen-Shannon Divergence (JSD) between the model's 1-step predictive distribution and the true distribution.
• Findings:
  • The DKF, constrained by Gaussian transitions, placed probability mass between the two true modes, failing to capture the bimodal structure.
  • The DDSSM successfully captured both modes of the true distribution.
  • Quantitatively, the DDSSM achieved a significantly lower JSD than the DKF.
  • Qualitatively, the DDSSM produced reconstructions and forecasts that were representative of the data, whereas the DKF forecasts lacked consistent patterns and failed to reconstruct observed data well.

Significance and Claims

The paper claims that by casting the problem of latent time-series diffusion as a State Space Model with a diffusion transition distribution, it is possible to achieve end-to-end training that allows for flexible reconstructions without penalizing the learning of latent dynamics.

The authors position the DDSSM as a method that resolves the tension between the expressivity of diffusion models and the structural rigor of SSMs. They assert that the Gaussian prior assumption in traditional Deep SSMs is harmful for learning latent representations in multimodal settings, and that their approach provides a more expressive prior over latent trajectories. The work extends the literature on latent diffusion models for time series by moving beyond two-stage training paradigms.

The authors note that future work will involve applying the DDSSM to large-scale datasets and incorporating further ideas from deep SSM literature to fully utilize the power of the diffusion transition distribution.