Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Amortized Control of Continuous State Space Feynman-Kac Model for Irregular Time Series

Byoungwoo Park, Hyungi Lee, Juho Lee

TL;DR

ACSSM addresses irregularly sampled time series by combining a multi-marginal Doob's $h$-transform with a stochastic optimal control framework to approximate the posterior path measure $\mathbb{P}^{\star}$. It introduces amortized inference via an auxiliary latent variable and a simulation-free latent-dynamics formulation guided by a transformer-based data assimilation scheme, enabling parallel inference and efficient ELBO computation. Theoretical contributions include extending the Doob's $h$-transform to multiple marginals and deriving a tight variational bound that links SOC to the Doob transform. Empirically, ACSSM achieves strong performance on per-time-point classification/regression and sequence interpolation/extrapolation across real-world irregular time series, with substantial gains in training efficiency due to parallel computation.

Abstract

Many real-world datasets, such as healthcare, climate, and economics, are often collected as irregular time series, which poses challenges for accurate modeling. In this paper, we propose the Amortized Control of continuous State Space Model (ACSSM) for continuous dynamical modeling of time series for irregular and discrete observations. We first present a multi-marginal Doob's $h$-transform to construct a continuous dynamical system conditioned on these irregular observations. Following this, we introduce a variational inference algorithm with a tight evidence lower bound (ELBO), leveraging stochastic optimal control (SOC) theory to approximate the intractable Doob's $h$-transform and simulate the conditioned dynamics. To improve efficiency and scalability during both training and inference, ACSSM leverages auxiliary variable to flexibly parameterize the latent dynamics and amortized control. Additionally, it incorporates a simulation-free latent dynamics framework and a transformer-based data assimilation scheme, facilitating parallel inference of the latent states and ELBO computation. Through empirical evaluations across a variety of real-world datasets, ACSSM demonstrates superior performance in tasks such as classification, regression, interpolation, and extrapolation, while maintaining computational efficiency.

Amortized Control of Continuous State Space Feynman-Kac Model for Irregular Time Series

TL;DR

ACSSM addresses irregularly sampled time series by combining a multi-marginal Doob's -transform with a stochastic optimal control framework to approximate the posterior path measure . It introduces amortized inference via an auxiliary latent variable and a simulation-free latent-dynamics formulation guided by a transformer-based data assimilation scheme, enabling parallel inference and efficient ELBO computation. Theoretical contributions include extending the Doob's -transform to multiple marginals and deriving a tight variational bound that links SOC to the Doob transform. Empirically, ACSSM achieves strong performance on per-time-point classification/regression and sequence interpolation/extrapolation across real-world irregular time series, with substantial gains in training efficiency due to parallel computation.

Abstract

Many real-world datasets, such as healthcare, climate, and economics, are often collected as irregular time series, which poses challenges for accurate modeling. In this paper, we propose the Amortized Control of continuous State Space Model (ACSSM) for continuous dynamical modeling of time series for irregular and discrete observations. We first present a multi-marginal Doob's -transform to construct a continuous dynamical system conditioned on these irregular observations. Following this, we introduce a variational inference algorithm with a tight evidence lower bound (ELBO), leveraging stochastic optimal control (SOC) theory to approximate the intractable Doob's -transform and simulate the conditioned dynamics. To improve efficiency and scalability during both training and inference, ACSSM leverages auxiliary variable to flexibly parameterize the latent dynamics and amortized control. Additionally, it incorporates a simulation-free latent dynamics framework and a transformer-based data assimilation scheme, facilitating parallel inference of the latent states and ELBO computation. Through empirical evaluations across a variety of real-world datasets, ACSSM demonstrates superior performance in tasks such as classification, regression, interpolation, and extrapolation, while maintaining computational efficiency.
Paper Structure (49 sections, 11 theorems, 120 equations, 4 figures, 4 tables, 4 algorithms)

This paper contains 49 sections, 11 theorems, 120 equations, 4 figures, 4 tables, 4 algorithms.

Table of Contents

  1. Introduction
  2. Notation
  3. Preliminaries
  4. Continuous-Discrete State Space Model
  5. Controlled Continuous-Discrete State Space Model
  6. Multi Marginal Doob's h-transform
  7. Stochastic Optimal Control
  8. Efficient Modeling of the Latent System
  9. Locally Linear Approximation
  10. Parallel Computation
  11. Auxiliary Variable
  12. Training and Inference
  13. Training Objective Function
  14. Experiment
  15. Per time point classification and regression
  16. ...and 34 more sections

Key Result

Theorem 3.1

Let us define a sequence of functions $\{h_i\}_{i \in [1:k]}$, where each $h_i : [t_{i-1}, t_i) \times \mathbb{R}^d \to \mathbb{R}_{+}$, for all $i \in [1:k]$, is a conditional expectation $h_i(t, \mathbf{x}_t) := \mathbb{E}_{\mathbb{P}}\left[ \prod_{j \geq i}^k f_{j}(\mathbf{y}_{t_j}|\mathbf{X}_{t_ Then, with the initial condition $\mu^{\star}_0(d\mathbf{x}_0) = h_1(t_0, \mathbf{x}_0)\mu_0(d\math

Figures (4)

  • Figure 1: Conceptual illustration. Given the observed time stamps $\mathcal{T} = \{t_i\}_{i \in [1:4]}$ and the unseen time stamps $\textcolor{red}{\mathcal{T}_u}$ ($\color{red}{\times}$ in figure), the encoder maps the input time series $\{\mathbf{o}_{t}\}_{t \in \mathcal{T}}$ into auxiliary variables $\{\mathbf{y}_t\}_{t \in \mathcal{T}}$. These auxiliary variables are then utilized to compute the control policies $\{\alpha\}_{i \in [1:5]}$ through a masked attention mechanism that relies on two different assimilation schemes. The computed policies $\{\alpha\}_{i \in [1:5]}$ control the prior dynamics $\mathbb{P}$ over the interval $[0, T]$ to approximate the posterior $\mathbb{P}^{\star}$ in the latent space. Finally, the sample path $\textcolor{rgb(46,117,182)}{\mathbf{X}^{\alpha}_{0:T}} \sim \mathbb{P}^{\alpha}$ are decoded to generate predictions across the complete time stamps $\textcolor{rgb(124,159,102)}{\mathcal{T}'} = \mathcal{T} \cup \mathcal{T}_u$ ($\textcolor{rgb(124,159,102)}{\circ}$ in figure), over the entire interval $[0, T]$.
  • Figure 2: Two types of information assimilation schemes.
  • Figure 3: Example of the pendulum sequence. (Up) The input image sequences $\{\mathbf{o}\}_{t \in \mathcal{T}}$ observed at irregular time stamps. (Down) The angular values of $\textcolor{rgb(106,153,195)}{sin(\theta_t)}$ and $\textcolor{rgb(255,96,38)}{cos(\theta_t)}$ where $\theta_t$ represents the angle of the pendulum at time $t \in [0, 100]$, are used as regression targets.
  • Figure 4: Illustraion of Masked Attention. Given the observed time stamps $\mathcal{T} = \{t_i\}_{i=1}^3$ ($\circ$) and the unseen time stamps $\textcolor{red}{\mathcal{T}_u}$ ($\times$), attention scores are computed based on the latent observations $\{\mathbf{y}_t\}_{t \in \mathcal{T}}$. (Up) The mask of the History assimilation scheme consists of two components: one masks attention scores related to future time stamps, and the other masks those related to unseen time stamps. (Down) The mask of the Full assimilation scheme only blocks the attention scores corresponding to unseen time stamps. Using these masks, masked attention is calculated. For History assimilation scheme, the latent variables $\mathbf{z}_{t_i}$ include information up to time $t_i$, while for Full assimilation scheme, $\mathbf{z}_{t_i}$ incorporates all available information. Finally, latent variables corresponding to unseen time stamps are filled with the nearest past latent variable value.

Theorems & Definitions (24)

  • Theorem 3.1: Multi-marginal Doob's $h$-transform
  • Theorem 3.2: Dynamic Programming Principle
  • Theorem 3.3: Verification Theorem
  • Lemma 3.4: Hopf-Cole Transformation
  • Corollary 3.5: Optimal Control
  • Theorem 3.6: Tight Variational Bound
  • Remark 3.7: Diagonalization
  • Theorem 3.8: Simulation-free estimation
  • Remark 3.9: Non-Markov Control
  • Definition B.1: Infinitesimal Generator
  • ...and 14 more