Technische Samenvatting: Diffusion-Driven State Space Models (DDSSM)

Probleemstelling

Praktijkgebruikers in tijdreeksmodellering worden vaak geconfronteerd met een afweging tussen de expressiviteit van een model en het vermogen om latente systeemdynamiek te vangen. Bestaande benaderingen offeren doorgaans één van beide doelen op:

• Deep State Space Models (DSSMs) (bijv. Deep Kalman Filters) bieden principiële inferentie en end-to-end training, maar gaan doorgaans uit van Gaussische transitiedichtheden. Deze aanname beperkt hun expressiviteit, waardoor ze vaak niet in staat zijn om multimodale dynamiek te vangen, en induceert een Gaussische prior die latente representaties kan overreguleren, waardoor noodzakelijke structurele informatie verloren gaat.
• Diffusiemodellen zijn zeer expressief en in staat om complexe, multimodale distributies te modelleren. Wanneer ze echter worden toegepast op tijdreeksen, worden ze vaak gebruikt als generatieve samplers geconditioneerd op het verleden, of vertrouwen ze op een twee-fasen training (eerst een autoencoder leren, en dan een diffusiemodel). Deze twee-fasen aanpak kan overgeregulariseerde latente ruimtes produceren die informatie weggooien die nodig is voor downstream generatieve dynamiek.

Er bestaat een openstaand probleem over hoe men een autoencoder en een diffusiemodel op een principiële manier gezamenlijk kan trainen op sequentiële data om de sterke punten van beide kaders te combineren.

Methodologie

De auteurs stellen de Diffusion-Driven State Space Model (DDSSM) voor, die diffusiemodellen direct integreert in de transitiedichtheid van een State Space Model (SSM).

Generatief Model

Het model gaat ervan uit dat observaties $x_{1:T}$ worden gegenereerd vanuit latente toestanden $z_{1:T}$ die evolueren via een Markov-proces, potentieel geconditioneerd op covariaten $u_{1:T}$. De gezamenlijke waarschijnlijkheid wordt gefactoriseerd als:
$$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})$$

Kerninnovaties zijn onder meer:

1. Diffusie-transities: In plaats van een parametrische Gaussische transitie $p^{(t)}_\psi$, gebruikt het model een diffusieproces. Voor elk tijdstap $t$ wordt de transitiedichtheid gemodelleerd door tussenliggende ruisgevoelige latente toestanden $z^1_t, \dots, z^K_t$ te introduceren. De transitie wordt gedefinieerd 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})$$
   waarbij $z^0_t = z_t$ en $z^K_t \sim \mathcal{N}(0, I)$. Dit stelt het model in staat om willekeurige, potentieel multimodale transitiedichtheden te leren.
2. Emissies: De emissiedichtheid $p^{(t)}_\theta$ blijft een standaard Gaussische distributie geparametriseerd door neurale netwerken, wat residuele onzekerheid afhandelt terwijl het diffusiemodel complexe niet-Gaussische structuren vangt.
3. Initialisatie: De initiële toestandsdistributie $p_\eta(z_{1:j}|u_{1:j})$ maakt gebruik van een Variational Hierarchical Prior (VHP) om restrictieve aannames over de initiële latente toestanden te vermijden.

Inferentie en Training

Omdat de ware posterior intractabel is, gebruiken de auteurs een variationele inferentiekader. Ze leiden een Evidence Lower Bound (ELBO) af die de voorwaardelijke onafhankelijkheidsstructuur van de ware posterior weerspiegelt.

• Variationele Posterior: De posterior wordt gefactoriseerd om een gesmoothde transitieterm (afhankelijk van het verleden van latente toestanden en toekomstige observaties) en een diffusie-ruisproces term te bevatten.
• Gezamenlijke Training: In tegen tegenstelling tot twee-fasen methoden, wordt de DDSSM end-to-end getraind. Het doel is het maximaliseren van de ELBO, die bestaat uit:
  • Een reconstructieverlies (negatieve log-likelihood van emissies).
  • Een KL-divergentieterm voor initialisatie.
  • Een diffusieketenverlies (afgeleid van het reverse proces van het diffusiemodel) dat de encoder stimuleert om latente toestanden te produceren die voorspelbaar zijn onder de geleerde transitiedynamiek.
• Preconditioning: Om gradiëntstabiliteit te garanderen over verschillende ruisniveaus tijdens de gezamenlijke training, gebruiken de auteurs preconditioning (Karras et al., 2022) om de inputs en outputs van het diffusienetwerk te herschalen.

Belangrijkste Bijdragen

1. DDSSM Framework: De introductie van een State Space Model waarbij de transitiedichtheid wordt vervangen door een diffusiemodel, wat het leren van flexibele, multimodale latente dynamiek mogelijk maakt zonder restrictieve Gaussische aannames.
2. End-to-End Training: Een principiële oplossing voor het openstaande probleem van het gezamenlijk trainen van een autoencoder en een diffusiemodel op sequentiële data. Dit voorkomt de overreguleratieproblemen die geassocieerd worden met de twee-fasen training van veel latente tijdreeks-diffusiemodellen.
3. Theoretische Afleiding: Een rigoureuze afleiding van de ELBO voor dit hybride model, die aantoont hoe de optimale variationele posterior factoriseert om overeen te komen met de structuur van de ware posterior, waardoor de benaderingskloof wordt verkleind.

Experimentele Resultaten

De auteurs evalueerden de DDSSM op een gesimuleerde tijdreeksdataset met bimodaal ruis (een latente random walk met bimodale transities).

• Setup: Ze vergeleken de DDSSM met een Deep Kalman Filter (DKF) met een Variational Hierarchical Prior en een Last Observation Carried Forward (LOCF) baseline.
• Metrieken: Prestaties werden gemeten met behulp van de Jensen-Shannon Divergentie (JSD) tussen de 1-stap voorspellende distributie van het model en de ware distributie.
• Bevindingen:
  • De DKF, beperkt door Gaussische transities, plaatste de waarschijnlijkheidsmassa tussen de twee ware modi, waardoor de bimodale structuur niet werd gevangen.
  • De DDSSM slaagde erin beide modi van de ware distributie te vangen.
  • Kwantitatief behaalde de DDSSM een significant lagere JSD dan de DKF.
  • Kwalitatief produceerde de DDSSM reconstructies en voorspellingen die representatief waren voor de data, terwijl de DKF-voorspellingen geen consistente patronen vertoonden en de geobserveerde data niet goed konden reconstrueren.

Betekenis en Claims

Het artikel claimt dat door het probleem van latente tijdreeksdiffusie te formuleren als een State Space Model met een diffusie-transitiedistributie, het mogelijk is om end-to-end training te bereiken die flexibele reconstructies toestaat zonder het leren van latente dynamiek te benadelen.

De auteurs positioneren de DDSSM als een methode die de spanning oplost tussen de expressiviteit van diffusiemodellen en de structurele strengheid van SSM's. Zij stellen dat de Gaussische prior-aanname in traditionele Deep SSM's schadelijk is voor het leren van latente representaties in multimodale settings, en dat hun aanpak een expressievere prior biedt over latente trajecten. Het werk breidt de literatuur over latente diffusiemodellen voor tijdreeksen uit door verder te gaan dan de twee-fasen trainingsparadigma's.

De auteurs merken op dat toekomstig werk zal bestaan uit het toepassen van de DDSSM op grootschalige datasets en het integreren van verdere ideeën uit de diepe SSM-literatuur om de kracht van de diffusie-transitiedistributie volledig te benutten.