Technisches Resümee: Diffusion-Driven State Space Models (DDSSM)

Problemstellung

Praktiker in der Zeitreihenmodellierung stehen oft vor einem Zielkonflikt zwischen der Ausdrucksstärke eines Modells und der Fähigkeit, latente Systemdynamiken zu erfassen. Bestehende Ansätze opfern typischerweise eines der beiden Ziele:

• Deep State Space Models (DSSMs) (z. B. Deep Kalman Filter) bieten eine fundierte Inferenz und ein End-to-End-Training, setzen jedoch meist Gaußsche Übergangsdichten voraus. Diese Annahme schränkt ihre Ausdrucksstärke ein, da sie oft nicht in der Lage ist, multimodale Dynamiken zu erfassen, und eine Gaußsche Prior-Verteilung induziert, die latente Repräsentationen übermäßig regularisieren und notwendige strukturelle Informationen verwerfen kann.
• Diffusionsmodelle sind hochgradig expressiv und fähig, komplexe, multimodale Verteilungen zu modellieren. Wenn sie jedoch auf Zeitreihen angewendet werden, werden sie oft als generative Sampler eingesetzt, die auf vergangenen Beobachtungen konditioniert sind, oder sie beruhen auf einem zweistufigen Training (erstes Erlernen eines Autoencoders, dann eines Diffusionsmodells). Dieser zweistufige Ansatz kann übermäßig regularisierte latente Räume erzeugen, die Informationen verwirfen, die für nachgelagerte generative Dynamiken erforderlich sind.

Es besteht ein offenes Problem hinsichtlich der Frage, wie man einen Autoencoder und ein Diffusionsmodell auf sequenziellen Daten auf eine fundierte Weise gemeinsam trainieren kann, um die Stärken beider Frameworks zu kombinieren.

Methodik

Die Autoren schlagen das Diffusion-Driven State Space Model (DDSSM) vor, welches Diffusionsmodelle direkt in die Übergangsdichte eines State Space Models (SSM) integriert.

Generatives Modell

Das Modell nimmt an, dass Beobachtungen $x_{1:T}$ aus latenten Zuständen $z_{1:T}$ generiert werden, die einem Markov-Prozess folgen, potenziell konditioniert auf Kovariaten $u_{1:T}$. Die gemeinsame Likelihood wird wie folgt faktorisiert:
$$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})$$

Zu den zentralen Innovationen gehören:

1. Diffusions-Übergänge: Anstelle einer parametrischen Gaußschen Übergangsdichte $p^{(t)}_\psi$ verwendet das Modell einen Diffusionsprozess. Für jeden Zeitschritt $t$ wird die Übergangsdichte durch die Einführung intermediärer verrauschter latenter Zustände $z^1_t, \dots, z^K_t$ modelliert. Der Übergang ist definiert als:
   $$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})$$
   wobei $z^0_t = z_t$ und $z^K_t \sim \mathcal{N}(0, I)$. Dies ermöglicht es dem Modell, beliebige, potenziell multimodale Übergangsdichten zu lernen.
2. Emissionen: Die Emissionsdichte $p^{(t)}_\theta$ bleibt eine Standard-Gauß-Verteilung, die durch neuronale Netze parametrisiert wird, um residuelle Unsicherheit zu handhaben, während das Diffusionsmodell die komplexen nicht-gaußschen Strukturen erfasst.
3. Initialisierung: Die Verteilung des Anfangszustands $p_\eta(z_{1:j}|u_{1:j})$ nutzt einen Variational Hierarchical Prior (VHP), um restriktive Annahmen über die Anfangszustände zu vermeiden.

Inferenz und Training

Da die wahre Posterior-Verteilung unintraktabel ist, verwenden die Autoren ein variatonal-inferentielles Framework. Sie leiten eine Evidence Lower Bound (ELBO) ab, die der bedingten Unabhängigkeitsstruktur der wahren Posterior-Verteilung entspricht.

• Variational Posterior: Die Posterior-Verteilung wird so faktorisiert, dass sie einen geglätteten Übergangsterm (abhängig von vergangenen Latenzen und zukünftigen Beobachtungen) sowie einen Diffusions-Rauschprozess-Term enthält.
• Gemeinsames Training: Im Gegensatz zu zweistufigen Methoden wird das DDSSM End-to-End trainiert. Das Ziel ist die Maximierung der ELBO, welche Folgendes umfasst:
  • Einen Rekonstruktionsverlust (negative Log-Likelihood der Emissionen).
  • Einen KL-Divergenz-Term für die Initialisierung.
  • Einen Diffusionsketten-Verlust (abgeleitet aus dem reversen Prozess des Diffusionsmodells), der den Encoder dazu anregt, latente Zustände zu erzeugen, die unter den gelernten Übergangsdynamiken vorhersagbar sind.
• Preconditioning: Um die Gradientenstabilität über verschiedene Rauschpegel während des gemeinsamen Trainings zu gewährleisten, verwenden die Autoren Preconditioning (Karras et al., 2022), um die Inputs und Outputs des Diffusionsnetzwerks neu zu skalieren.

Zentrale Beiträge

1. DDSSM-Framework: Die Einführung eines State Space Models, bei dem die Übergangsdichte durch ein Diffusionsmodell ersetzt wird, was das Lernen flexibler, multimodaler latenter Dynamiken ohne restriktive Gaußsche Annahmen ermöglicht.
2. End-to-End-Training: Eine fundierte Lösung für das Problem des gemeinsamen Trainings eines Autoencoders und eines Diffusionsmodells auf sequenziellen Daten. Dies vermeidet die Probleme der Überregulierung, die mit den im Bereich der latenten Zeitreihen-Diffusion üblichen zweistufigen Trainingsansätzen assoziiert sind.
3. Theoretische Herleitung: Eine rigorose Ableitung der ELBO für dieses hybride Modell, die zeigt, wie die optimale Variational Posterior so faktorisiert, dass sie der Struktur der wahren Posterior-Verteilung entspricht, wodurch die Approximationslücke verringert wird.

Experimentelle Ergebnisse

Die Autoren evaluierten das DDSSM auf einem simulierten Zeitreihen-Datensatz, der bimodales Rauschen (einen latenten Random Walk mit bimodalen Übergängen) aufweist.

• Setup: Sie verglichen das DDSSM mit einem Deep Kalman Filter (DKF) mit einem Variational Hierarchical Prior und einer Last Observation Carried Forward (LOCF) Baseline.
• Metriken: Die Leistung wurde mittels der Jensen-Shannon-Divergenz (JSD) zwischen der 1-Schritt-Prädiktionsverteilung des Modells und der wahren Verteilung gemessen.
• Ergebnisse:
  • Der DKF, der durch Gaußsche Übergänge beschränkt ist, platzierte die Wahrscheinlichkeitsmasse zwischen den beiden wahren Modi und konnte die bimodale Struktur somit nicht erfassen.
  • Das DDSSM konnte beide Modi der wahren Verteilung erfolgreich erfassen.
  • Quantitativ erreichte das DDSSM eine signifikant geringere JSD als der DKF.
  • Qualitativ erzeugte das DDSSM Rekonstruktionen und Vorhersagen, die repräsentativ für die Daten waren, während die Vorhersagen des DKF keine konsistenten Muster aufwiesen und die beobachteten Daten nicht gut rekonstruieren konnten.

Bedeutung und Behauptungen

Die Autoren behaupten, dass durch die Darstellung der latenten Zeitreihen-Diffusion als ein State Space Model mit einer Diffusions-Übergangswahrscheinlichkeit ein End-to-End-Training möglich ist, das flexible Rekonstruktionen ermöglicht, ohne das Lernen latenter Dynamiken zu behindern.

Die Autoren positionieren das DDSSM als eine Methode, die den Spannungsgrad zwischen der Ausdrucksstärke von Diffusionsmodellen und der strukturellen Stringenz von SSMs auflöst. Sie behaupten, dass die Gaußsche Prior-Annahme in traditionellen Deep SSMs schädlich für das Lernen latenter Repräsentationen in multimodalen Settings ist und dass ihr Ansatz einen expressiveren Prior über latente Trajektorien bietet. Die Arbeit erweitert die Literatur zu latenten Diffusionsmodellen für Zeitreihen, indem sie über zweistufige Trainingsparadigmen hinausgeht.

Die Autoren merken an, dass zukünftige Arbeiten darin bestehen werden, das DDSSM auf groß angelegte Datensätze anzuwenden und weitere Ideen aus der Deep-SSM-Literatur zu integrieren, um die Leistungsfähigkeit der Diffusions-Übergangswahrscheinlichkeit voll auszuschöpfen.