Amortized Control of Continuous State Space Feynman-Kac Model for Irregular Time Series
Authors: Byoungwoo Park
Hyungi Lee
Juho Lee
Venue: ICLR 2025 Oral (2025)
Keywords: stochastic optimal control
variational inference
state space model
irregular time series
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
TL;DR: We propose a multi-marginal Doob's $h$-transform for irregular time series and variational inference with stochastic optimal control to approximate it.

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.

Published as a conference paper at ICLR 2025
AMORTIZED CONTROL OF CONTINUOUS STATE SPACE FEYNMAN-KAC MODEL FOR IRREGULAR TIME SERIES
Byoungwoo Park, Hyungi Lee, Juho Lee
KAIST
{bw, park, 1hk2708, juholee)ekaist.com
ABSTRACT
Many real-world datasets, such as healthcare, climate, and economics, are often collected 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 obser- vations. Following this, we introduce variational inference algorithm with a tight evidence lower bound (ELBO), leveraging stochastic optimal control (SOC) theory approximate the intractable Doob's 's fi-transform and simulate the condi- tioned 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, incorporates a simulation- latent dynamics framework and wansformer-basee 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. Code available https: ptttps://gee //gtthbtttt///////////////t com foom/wwpaate com/bw--paz</ACSMM
INTRODUCTION
State space models (SSMs) are widely used for modeling time series data (Fraccaro et al., 2017;
Rangapuram et al., 2018; de Bézenac et al., 2020). However, SSMs typically assume evenly spaced ob- servations, whereas many real-world datasets, such as healthcare (Silva et al., 2012), climate (Menne al., 2015), are often collected as irregular time series. This poses challenges for accurate modeling. provokes the neural differential equations (Chen et al., 2018; Rubanova et al., 2019; Li et al., 2020; Kidger et al., 2020; Zeng et al., 2023) for continuous-tinn time dynamical modeling of time- series. Moreover, by incorporating stochastic dynamical models with SSM models, the application of continuous-dyaamical state space models (CD-SSM) (Jazwinski, 2007) with neural networks have emerged which allows for more flcxible modeling (Schirmer et al., 2022; Ansari et al., 2023). However, they typically rely on the Bayesian recursion approximate the desired latent posterior distribution (filtering/smoottin
distribution) involves sequential prediction and updating, resulting in computational costs that scale with the length of observations (Särkkä GaaríiaFerránnee 2020).
Instead of relying on Bayesian recursion, we propose an Amortized Control of continuous State Space Model (ACSSM), a variational inference (VI) algorithm that directly approximates the posterior distribution using measure-theoretio formulation, enable simulation of continuous trajectories from such distributions. To achieve this, we introduce Feynman- Kac model (Del Moral, 201 Chopin et al., 2020), which facilitates thorough sequential analysis. Within this formulation, propose a multi multi-marginal Doob's s h-transform to provide conditional latent dynamics described by stochastic differential equations (SDEs) which conditioned on irregularly sampled, set of discrete observations. extends the previous Doob's 's h--transform (Chopin et al., 2023) by incorporating multi-marginal constraints instead of single terminal constraint.
Considering that the estimation of the Doob's h- trannfform is infeasible in general, we utilize the theory of stochastic optimal control (SOC) (Fleming & Soner, 2006; Carmona, 2016) to simulate the

---

Published as a conference paper at ICLR 2025
conditioned SDEs by approximating the Doob's s h-transform. It allow to propose a tight evidence lower bound (ELBO) for the aforementioned VI algorithm by establishing a fundamental connection between the partial differential equations (PDEs) associated with Doob's h--rranformm and SOC. The Doob's ss h--transfomm often referred to as the twist-function in Sequential Monte Carlo (SMC) literature (Guarniero al., 2017) to approximate the smoothing distributions. Building on this, (Heng et al., 2020) introduced an algorithm to approximate the twisted transition kernel directly, while a recent concurrent study (Lu & Wang, 2024) extended this approach to continuous-time settings. However, both studies primarily emphasize approximation methods rather than practical applications.
In practical situations, the computation of ELBO for VI algorithm might impractical due to the instability and high memory demands associated with gradient computation of the approximated stochastic dynamics over the entire sequence interval (Liu al., 2024; Park et al., 2024). To address this issue, we propose two efficient modeling approaches: 1) We establish amortized inference by introducing an auxiliary variable to the latent space, generated by neural network encoder-decoder. It maps the high-dimensional time-series into suitable low-dimensional space, allowing more flexible parameterization of the latent dynamics. Moreover, amortization allows the inference of the posterior distribution for novel time-series sequence without relying Bayesian recursion by incorporating the learned control function. 2) We leverage the simulation-free property, which enables closed-form sampling from intermediate latent marginal distributions that can be computed in a temporally parallel way. Additionally, we explore a more flexible linear approximation of the drift function in controlled SDEs to enhance the efficiency of the proposed controlled dynamics.
We evaluated ACSSM on several time-series tasks across various real-world datasets, Our experiments show that ACSSM consistently outperforms existing baseline models in each tasks, demonstrating its effectiveness in capturing the underlying dynamics of irregular time-series. Additionally, ACSSM achieves significant computational efficiency, enabling faster training times compared to dynamics- based models that rely on numerical simulations. A summary of the key concepts of ACSSM, along with related works, is provided in Appendix A. We summarize our contributions as follows:
We extend the theory of Doob's ' h h-trrasffrrm to multi-marginal cases. This indicates the existence of class of conditioned SDEs that depend on future observations, where the solutions of these SDEs lead to the true posterior path measure within the framework of CD-SSM.
We reformulate the simulation of conditioned SDEs as a SOC problem to approximating an impractical Doob's 's h--transfomm. By leveraging the connection between SOC theory and Doob's h- h-transform, we propose a variational inference algorithm with a tight ELBO.
For practical real-world applications, we introduce an efficient and scalable modeling approach that enables parallelization of latent dynamic simulation and ELBO computation.
We demonstrate its superior performance across various real-world irregularly sampled time- series tasks, including per-time point classification, regression, and sequence interpolation and extrapolation, all with computational efficiency.
Notation Throughout this paper, we denote path measure by p(+), defined on the space of continuous functions Ω C( TJ, Rd). We sometimes denote with TP the expectation as B [-] Er[-|X,
x]. where the stochastic processes corresponding PP()) are represented as x and their time-
marginal distribution at time t E [0, T' is given by the fosh-forward measure p4,² = (x ) 1() with marginal density p(? This marginal density represents the Radon-Nikodym derivative duy ( (x)) (x) = p ( (x)d((x)) where L denotes the Lebesgue measure. Additionally, for a function V J0, 77 R + R. we define the first and second derivatives with respect to x E R2 as V V and Vxx respectively, and the derivative with respect to time E [0, T] as V.V. For a sequence of functions (V) (eeeE()L)) we will denote (t, x) := Vi, and [1 k] (1, k). Finally, the Bullback-Leible (KL.) divergence du between two probability measures 4 and is defined D (ulp) Sge log de (x)dp(x) when μ is absolutely continuous with respect to v, and (p)( = too otherwise.
2 PRELIMINARIES
Contimuous-Discen State Space Model Consider for a set of time-stamps (regular or irregular) (1.)/2 over an interval T [0, 7). i... = to ... S le T. The CD-SSM assumes
2

---

Published as a conference paper at ICLR 2025
continuous-time Markov state trajectory Xo:r in latent space R is given solution of the SDE:
(Prior State) dX, = b(t, X + dWe.
(1)
where Xo di and (Wa) tejoo is a -V Wiener process that is independent the μoo Since
X₂ is Markov process, the time-evolution of marginal distribution is governed by a transition density, which is the solution the Pokker-Plance equation assocaited with X,. This allows to define a path measure F that represent the weak solutions of the SDE in (1) over an interval [0, T)
For a measurement model (y, ). we consider the case that we have only access to the re- alization of the (latent) observation process at each discrete-time stamps [E]))E( ie., Yt, ~ ge (yr,,,Xa,) I,,), vi E [] kj. In this paper, our goal is to infer the classes of SDEs which inducing the filteringsmmoohingg path measure F P" (-IHe,), the conditional distribution over the interval [0, I'I for given IP and a set of observations up to time Eg. Hi. = {y. S k):
(Posterior Dist.) dlM (Xo:r
dI (Ht,))
II 1=1
9: (y X, ddP )dP(Xo.r)
(2)
Z(He,
where the normalizing constant Z(H,,)) Ep IT == 9: (y x serve observations likelihood.
The path measure formulation of the posterior distribution described in (2) referred to as Feynman-Kan
models. See (Del Moral, 2011; Chopin al., 2020) for more comprehensive understanding.
CONTROLLED CONTINUOUS-DISCREN STATE SPACE MODEL
In this section, we introduce our proposed model, ACSSM. First, we present the Multi-marginal Doob's ' h-transform, outlining the continuous dynamics for I in Section 3.1. Then, in Section 3.2, we frame the VI for approximating IP* using SOC. To support scalable real-world applications, we discuss efficient modeling and amortized inference in Section 3.3?.
3.1 MULTI MARGINAL DOOB'S h-TRANSFORM
Before applying to approximate the posterior distribution Fx in SOC, we first show that a class of SDEs exists whose solutions induce path measure equivalent to P in (2). This formulation provides a valuable insight for defining an appropriate objective function for the SOC problem in the next section. To do so, we first define sequence of normalized potential functions {fi (filee(aa)) where
each _2d R,, for all E [1 k).
fly (x..) g:ly,, (X,,) Lg(g))
where Lg(g.) = Sgg 9t. (y,, [x,, )dP (xu:r), for all i [1 k] is the normalization constant. Then, we can observe that the potential functions (fiJke(1:N) defined in (3) satisfying the normalizing = and di² (Xo:T(He,,) II I f((y,, (x dXP()) from (2). define the
property i.e., I f( Now, with the choice of reference measure P induced by Markov process in (1), conditional SDEs conditioned on He, which inducing the desired path measure PP. Note that this an extension of the original Doob's 's h-trransform (Doob, 1957), incorporating multiple marginal constraints. Below, we summarize the relevant result.
Theorem 3.1 (Multi- marginal Doob's h- t-transform). Let us define a sequence of functions {h: veeaaa where each 14 [4-1,44) 4:) x Rd + for all i E [1: : k), is conditional expectation hú(t, xt) == Ep I1 f5 (x,,) XX, = x where {felve(uan is defined in (3). Now, we define a function h: [0, x Rdd + R,, by integrating the functions (hiveepa))
h(t, Σ p.
See details on Appendix B.2, Proofs and detailed derivations are provided in Appendix B.
3

---

Published as a conference paper at ICLR 2025
Then, with the initial condition pū((xx)) = h1 (to. (d)))) the solution of the following condi-
tional SDE inducing the posterior path measure 1 in (2):
(Conditioned State)
dX; = [b( x;) + Vo h(t, X,)] dt + dWW
(5)
Theorem 3.1 demonstrates that we can obtain sample trajectories from P in (2) by simulating the dynamics in (5). However, estimating the functions {h,) [e([::] requires both the estimation of the sequence of normalization constants (E))e and the computation of conditional expectations, which infeasible in general. For these reasons, we propose VI algorithm to approximate the functions {h (E(L:A) and derive the variational bound for the VI by exploiting the theory of SOC.
3.2 STOCHASTIC OPTIMAL CONTROL
The SOC (Fleming & Soner, 2006; Carmona, 2016) a mathematical framework that deals with the problem of finding control policies in order to achieve certain object. In this paper, we define following
control-affie SDE, adjusting the prior dynamics in (1) with Markov control [0. T] x R R
(Controlled State) dX; [b(t. X + a(t, X dt + dWW.
(6)
where XØ ~ .4- We refer to the SDE in (6) as controlled SDE. We can expect that for a well- defined function set a E A. the class of controlled SDE (6) encompass the SDE in (5). This implies that the desired path measure F* can be achieved through the SOC formulation. In general, the goal of SOC is to find the optimal control policy a" that minimizes a given arbitrary cost function [1, Xt, (r) L.e.aa( (1, = arg min (1, x), and determine the value function V((,,,)) = min [[1, Xt, O), where V(t, x,) == [[, Xx, a*,aa)) S J(, a) holds for any a E A. Below, we demonstrate how, with carefully chosen cost function, the theory of SOC establishes connection between two classes of SDEs (5) and (6). This connection enables the development of a variational inference algorithm with a tight evidence lower bound (ELBO). To this end, we consider the following cost function:
J(t,x,, a.) = Ep [- |lals, x;)i
|a(ssx;)1|'s
Σ log fiy,, XX); x,
(7)
Then, the value function V for (7) satisfies the dynamic programming principle (Carmona, 2016):
Theorem 3.2 (Dynamic Programming Principle). Let us consider sequence of left continuous
functions [ViJee(us+1), where each V; (1b,(( tz) x IR))
1 |a,ll ds log f. (y,, + V4 XX
(t,) = min Ep aEk
Xx
(8)
for all i E [1 [6] and V4+1 = O. Then, for any 0- < t < u S T, the value function V for the cost function in (7) satisfying the recursion defined as follows:
J.
11(0) 1
Vtt, x1) = min Epa aeA
||a.11 ds Σ log f. + Vronz+ Xin X Xx : (9)
2
((IStiStr)))
with the indexing function [()) = maxfi E [1 kl]tt S u} and f. = flyy. X
The value functions presented in Theorem 3.2 suggest that the objective of the optimal control policy for the interval it,,1,t,) is not just to minimize the negative log-potential llo f.y.,..) for the immediate observation Vt,- Instead, it also involves considering future costs (V) and the corresponding future observations {y,, be(+++1s]] since V, follows recursive structure. Since our goal is to approximate P* in (5), it natural that the optimal control policy a should reflect the future observations {y,, [i as the h-function inherently does. Next, we will derive the explicit form of the optimal control policy.
Theorem 3.3 (Verification Theorem). Suppose there exist a sequence of left continuous functions (,), Rd)), for all E [1 k), satisfying the following Hamittonian-lacoii Bellman (HJB) equation:
xt,,)) E ((t-1,,)
(VxVeee as. + |lane 12] =0, t11St<
0,V,, + A.Vi, + min
(10)
aE
Vi(t,, x) -loog f(((,)))))) (i,) t= Vi [1 k))
(11)
4

---

Published as a conference paper at ICLR 2025
where minimum is attained by az (t. VxVi x). Now, define function [0, T) x R R by integrating the optimal controls (ar))e(1:)).
h a' at,,,) = Σ xt
(12)
Then, V(t,x,) - X1, a²) S I
[t, (t, Xx, a) holds for any (t, xx) [0, R& and a E A.
Note that the optimal control (12) for the cost function (7) share a similar structure as in (4). Since the theory of PDEs establishes a fundamental link between various classes of PDEs and SDEs (Richter Berner, 2022: Berner et al., 2024), it allows us to reveal the inherent connection between (12) and (4), thereby enable us to simulate the conditional SDE (5) an alternative way.
Lemma 3.4 (Hopf-Cole Transformation). The h function satisfying the following linear PDE:
Ocha, Athi, 0, f:-1 St t; h,(4; x) = fi(y,, x1),++1(1,,)) t=
(13)
Vi [1:
(14)
Moreover, for logarithm transformation V = loghh h. V satisfying the HJB equation in (10-11),
According Lemma 3.4, the solution of linear PDE in (13-14) is negative exponential the solution of the HJB equation in (10-11). 1). Therefore, it leads to the following corollary:
Corollary 3.5 (Optimal Control). For optimal control a* induced by the cost function (7) with dynamics (6), it satisfies ax Vx log h. In other words, we can simulate the conditional SDEs in (5) by simulating the controlled SDE (6) with optimal control a².
Corollary 3.5 states that the Markov process induced by a² and Vx log h is equivalent under same initial condition. However, comparing the pa induced by the controlled dynamics (6) with an initial condition μoo the conditioned dynamics (5) has an intial condition μ to inducing the desired path measure IP* In other words, although we find the optimal control a", the constant discrepancy between μá and y still remain, thereby keeping P* and pa misaligned. Fortunately, surrogate cost function can be derived from the variational representation under the KL-divergence, allowing us to find the optimal control while minimizing the discrepancy.
Theorem 3.6 (Tight Variational Bound). Let us assume that the path measure Ie induced by for (PPP
any a E A satisfies DKL(IP" [P+)) < oc. Then, for a cost function J in (7) and μā in (5), it holds:
D (P*1IT> = DxL (polm))) + Exo~po [T (0, Xo, at) = L((x) + log Z(H ) 2 0,
(15)
where the objective funnctooo L(a) (a negative ELBO) is given by:
[(a)= = Epp I.
k
log 9. g: x 2 [la(s, x;)) ds Σ
2 - log Z(He.))
(16)
Moreover. assume that L(a) has global minimum a* arg min L(a). Then the equality holds in (16) i..,, L(a*)((()))))))) = 1--llog log Z(He,) and μÓ = Pn almost everywhere with respect to Pa-
Theorem 3.6 suggests that the optimal control a² arg min, L(a) provides the tight variational bound for the likelihood functions [gi and the prior path measure P induced by (1). Further- more, a* ensures that μ Po. indicating that simulating the controlled SDE in (6) with a* and initial condition oll generates trajectories from the posterior path measure IP* in (2). In practice, the optimal control a* can be approximated using highly flexible neural network, which serves as function approximator (i.e., a == a1=; 0)), optimized through gradient descent- based optimization (Li al., 2020; Zhang& & Chen, 2022: Vargas et al., 2023).
However, applying gradient descent- based optimization necessitates computing gradients through the simulated diffusion process over the interval [0, 7] to estimate the objective function (16) for neural network training. This approach can become slow, unstable, and memory-intensive as the time horizon or dimension of latent space increases (Iakovlev et al., 2023; Park et al., 2024). It contrasts with the philosophy of many recent generative models (Ho et al., 2020; Song et al., 2020), which aim to decompose the generative process and solve the sub-problems jointly. Additionally, for inference, it requires numerical simulations such Buler-Maruyama solvers (Kloeden & Platen, 2013), which can also be mons-mnnnnnng for a long time series. It motivated us to propose an efficient and scalable modeling approach for real-world applications described in the next section.
5

---

Published as a conference paper at ICLR 2025
Amortized Inference
Latent Obs.
Control Policy
Latent Dynamics
Output
Masked Attention
A/m,
O0:T
Y
01
065
015
Encoder
X
Decoder
Time
Time
ytu
Pull
Consalee
Figure I: Conceptual illustration. Given the observed time stamps T (e))e( and the unseen time stamps Tu (x in figure), the encoder maps the input time series {or HET into auxiliary variables (yeT These auxiliary variables are then utilized to compute the control policies {a) through masked attention mechanism that relies on two different assimilation schemes. The computed policies [Ej:ss) control the prior dynamics P over the interval [0, 7'J to approximate the posterior 1* in the latent space. Finally, the sample path X Ipa are decoded generate predictions across the complete time stamps T' = =TUT, (o in figure), over the entire interval [0, T
3.3 EFFICIENT MODELING OF THE LATENT SYSTEM
The linear approximation of the general drift function provides a simulation-free property for the dynamics and significantly enhances scalability while ensuring high-quality performance (Deng et al., 2024). Motivate by this property, we investigate a class of linear SDEs to improve the efficiency of the proposed controlled dynamics, ensuring superior performance compared to other baselines. We introduce the following affine linear SDEs:
dx, - + a) dt + dWW, where X ~ N(mo: En);
(17)
and a matrix A E Rdxd and a vector a Rdd The solutions for X, (17) has a closed-form Gaussian distribution for any E [0, T). where the mean m, and covariance E, can be explicitly computed by solving the ODEs (Särkkä & Solin, 2019):
m. e Ar moo - A ( At I)a
(18)
Σ, = - - ( c-AA--e) A
(19)
However, calculating the moments m, and E, in (18) involves computing matrix exponentials, inversions, and performing numerical integration. These operations can be computationally intensive, especially for large matrices or when high precision is required. These computations can be simplifies by restricting the matrix A to be a diagonal or semi-position definite (SPD).
Remark 3.7 (Diagonalization). Since SPD matrix admits the decompoccooooiiionn A EDE with E E idxx and D E diag(R.)) the process X, expressed in a standard basis can be transformed to X, which have diagonalized drift function. In the space spanned by the eigen-basis E, the dynamics in (1) can be rewritten into:
dx, = [ -DX, + a dt + dWW, where Xo ~ N(moo Eo).
(20)
X, - B. E.W, m = E m, and = E Note that W, E TW for any t [0,T] due to the orthonormality of E, so We can be regarded as a standard Wiener process. Because of D diag(A) where X = (.I.... >aa and each >; 2 for all i E [1 : d,, we can obtain the state distributions of X, for any E [o by solving ODEs in (18-19) analytically, without the need for numerical computation. The results then transforming back to the standard basis ie.. m, Em,, Σ, E.,ET.
E,

---

Published as a conference paper at ICLR 2025
Locally Linear Approximation To leverage the advantages of linear SDEs in (17) which offer simulation-free property, we aim to linearize the drift function in (6). However, the naive formulation described in (17) may limit the expressiveness of the latent dynamics for real-world applications. Hence, we introduce a parameterization strategy inspired by (Becker et al., 2019; 99 Klushyn et al., 2021a), which leverage neural networks enhance the flexibility by fully incorporating an attentive structure with a given observations Yo.r while maintaining a linear formulation:
xx; A- + oi) dt + dW..
(21)
where the approximated drift function is constructed as affine formulation with following components:
A = Σ w (2.)A a₁ = Boze.
(22)
The matrix A, is given by convex combination of L trainable base (fa produced by the neural network fo. Additionally, B, Raxd is a trainable matrix. The latent variable z, is produced by the trans- former Te, which encodes the given observations 7o.7, depending on the task-specific information assimilation scheme.
Figure 2: Two types of infor-
matrices {A (A hej--L) where the weights wa
mation assimilation schemes.
O 211
Ztz
Ztg
FFN
FFN
FFN
N
Attn
Attn
Attn
Figure 2 illustrates two assimilation schemes using masked atten- tion mechanism': the history assimilation scheme. the transformer To encodes information up the current time and outputs 2,, Le., 2, = (Ht,). and the full assimilation scheme, the transformer To encodes information over the entire interval [0, 7] and outputs z₁, i.e,, 2, = To (HT). Note that this general formulation brought from the control formulation enables more flexible use of information encoded from observations. In contrast, previous based CD-SSM method (Schirmer et al., 2022) relies on recurrent updates, which limits them to typically using historical information only.
Y
History
Z
Z
Z
FFN
FFN
FFN
N
Attn
Attn
Attn
-ilttrniig
O'L,
Oytz
Oy
Full
Furthermore, since observations are updated only at discrete time steps (t, (ejss] the latent variables 2: remain constant within any interval E [t,-1,6,) (.) for all i E [1 : k]. making A; and a constant as well. As a result, the dynamics (21) remain linear over local intervals. This structure enables us to
derive a closed-form solution for the intermediate latent states.
Theorem 3.8 (Simulation frrel estimation). Let us consider sequences of SPD matrices {A, Aiiiel
that admit the eigen-decoopossitoon A, EDD.E with E E R and D; diag(IR.) 2 ffr all i E [1 k). control vectors (ass [e[[:K] and following control-afiine SDEs for all i E [1 k]:
dx,- = - A,X, + oa] dt + odW,, 1 LE (4)-1,44)-
(23)
Then, with Xo ~ (m, E,), ), the solution (23) is a Gaussian process N N(m,,. En,) with:
(I clta-ta 1)Ds) ( ân
m = E e
mo
Σ, EE c
- (-EEjr1(0)-1) (1)--1))) DI0
1
Σ
E/ -e (I elt (( E
2 A=
where iy
E m ES = Mg E and à E au all E [1: k..
Parallel Computation Given an associative operator é and a sequence of elements [st. ... s8x the parallel scan algorithm computes the all all-prefix-sum which returns the sequence
[8t, (St, ( St,),... , (8t, 6 St, 0
stx)
(24)
in Ollog K) time. Leveraging the linear formulation described in Theorem 3.8 and the inherent parallel nature of the transformer architecture for sequential structure, our method can be integrated with the parallel scan algorithm (Blelloch, 1990) resulting efficient computation of the marginal
Gaussian distributions by computing both moments {m:)eeforr and { (Echeloo,ry in a parallel sense
"See Appendix D.2 for more details.
*See Appendix C. for more details.
7

---

Published as a conference paper at ICLR 2025
Remark 3,9 (Non-Marko) Control). Note that (21-22) involves approximating the Markov control by a non-Markov control a(H)) == a², parameterized by neural network 8. However, Theorem 3.3 establishes that the optimal control should be Markov, as it is verified by the HJB equation (Van Han- del, 2007). In our case, we expect that with a high-capacity neural network 0, the local minimum fad obtained after M gradient descent steps g 8t VoL(a" 9 yields L(a*) = L(a -
Auxiliary Variable Moreover, to enhance flexibility, we treat Yo:r as an auxiliary variable in the latent space which is produced by neural network encoder qs applied to the given time series data 00:T i.e,, Yo.r ~ (Yo. (0:7), where it is factorized as
4o (yorr /
II go (y,,,,oo)) (o..) = III (os,,,, Eny
(25)
with fixed variance Eg- Additionally, it enables the modeling of nonlinear emission distributions where it is factorized as
through a neural network decoder
Po(or Yo:r)
II po (or, yt,);
(26)
with the likelihood function Pe depending the task at hand. This formulation first maps the time series 0o:T into suitable low-dimensional space Ya:7, allowing more efficient modeling of the latent dynamics Xirr- The information capturing the underlying physical dynamics resides in much lower- dimensional space compared to the original sequence (Fraccaro et al., 2017). Therefore, performing generative modeling in this reduced latent space (rather than directly in the high-dimensional domain, e.g., pixel values in an image sequence) offers greater flexibility in parameterization.
3.4 TRAINING AND INFERENCE
Training Objective Function We jointly train, in an end-to-end manner, the amortization pa- rameters {ő) 8/) for the encoder-decoder pair, along with the parameters of the latent dynamics 0 (fo. Ba, Tø, mp, Eu, (A [e(t-r)) which include the parameters required for controlled la- tent dynamics. The training is achieved by maximizing the evidence lower bound (ELBO) of the observation log-likelihood for given the time series Op.r:
I Pe (or, yr, (V(O.)) log =
log Pvi (Do:T) 2 E Ho (Vo.T)ourr)
(27)
I U6 (o lo..)
Hy
log (o - L()) EL,O 2,
(28)
Since Z(Hr,) = (yor)) the prior over auxiliary variable g(Yorr) can be computed using the ELBO L(0) proposed in (16) as part of our variational inference procedure for the latent posterior P in proposed Sec 2, including all latent parameters 0 containing the control a Note that our modeling is computationally favorable for both training and inference, since estimating marginal distributions in a latent space can be parallelized. This allows for the generation of latent trajectory over the entire interval without the need for numerical simulations. The overall training and inference processes are summarized in the Algorithm 3 and Algorithm 4 in Appendix, respectively.
4 EXPERIMENT
In this section, we present empirical results demonstrating the effectiveness of ACSSM in model- ing real-world irregular time-series data. The primary objective was to evaluate its capability to capture the underlying dynamics across various datasets. To demonstrate the applicability of the ACSSM, we conducted experiments on four tasks: per-time regression/classifcattio and sequence interpolation//ettapllttoo using four datasets: Human Activity, USHCN (Menne et al., 2015), and Physionet (Silva et al., 2012). We compare our approach against various baselines including RNN architecture (RKN-A, (Becker et al., 2019), GRU-A, (Chung al., 2014), GRU-D (Che
8

---

Published as a conference paper at ICLR 2025
et al., 2018)) well dynamics- based models (Latent-ODE (Chen et al., 2018; Rubanova et al., 2019), ODE- EE-N (Rubanova et al., 2019), GRU-ODE-B (De Brouwer et al., 2019), CRU (Schirmer et al., 2022), Latent-SDEn (Zeng et al., 2023)) and attention-based models (mTAND (Shukla & Marlin, 2021)), which have been developed for modeling irregular time series data. We reported the averaged results over five runs with different seed. The best results are highlighted in bold, while the second- best results are shown in blue. Additional experimental details can be found in Appendix D.
4.1 PER TIME POINT CLASSIFICATION & REGRESSION
Human Activity Classification We investigated the classification performance of our proposed model. For this purpose, we trained the model on the Human Activity dataset, which contains time- series data from five individuals performing various activities such as walking, sitting, lying, and standing etc. The dataset includes 12 features in total, representing 3D positions captured by four sensors attached to the belt, chest, and both ankles. Following the pre- pe-ppoceessng approach proposed by (Rubanova et al., 2019), the dataset comprises 6,554 sequences, each with 211 time points. The task is to classify each time point into one of seven activities.
Table I: Test Accuracy (%).
Model
Acc
Latent ODE!
87.0 1. 2.8
Latent-SDEis ? mTAND'
90.6 + 0.1
91.1 + 0.2
ACSSM (Ours)
91.4 0.4
f result from (Shukla & Marlin, 2)21). 1 resalt from (Zeng et al., 2023).
Table reports the test accuracy, showing that ACSSM outperforms all baseline models. We employed the full assimilation scheme maintain consistency with the other baselines, which infer the latent state using full observation. It is important to note that the two dynamical models, Latent-ODE and Latent-SDE,, incorporate a parameterized vector field in their latent dynamics, thereby rely on numerical solvers to infer intermediate states. Besides, mTAAND is an attention- based method that does not depend on dynamical or state-space models, thus avoiding numerical simulation. We believe the significant performance improvement of our model comes from its integration of an attention mechanism into dynamical models. Simulation-frrr dynamics avoid numerical simulations while maintaining temporal structure, leading to more stable learning.
Pendulum Regression Next, we explored the problem of sequence regression using pendulum experiment (Becker et al., 2019), where the goal is to infer the sine and cosine of the pendulum angle from irregularly observed, noisy pendulum images (Schirmer et al., 2022). To assess our performance, we compared it with previous dynamics- based models, reporting the regression MSE on a held-out test set as shown in Table 2. We employed the full assimilation scheme. The experimental results demonstrated that our proposed method out- performed existing models, delivering superior performance. These findings highlight that, even when linearizing the drift function, the amortization and proposed neural network based locally linear dynamics in See 3.3 preserves the expressivity of our approach, enabling more accurate inference of non-linear systems.
Table 2: Test MSE (x10-3) )).
Model
MSE
Latent Latent-ODet
15.70 1. 0.29
CRU
4.63 +
Latent-SDEa
3.84 0.35
S5"
3.41 1: 0.27
ITTAND
3.20 = 0.60
ACSSM (Ours)
2.98
f result from (Schimner et al., 2022).
I resalt from (Zeng al., 2023).
* result from (Smith et al., 2023).
Moreover, particularly in comparison to CRU, we believe that the significant performance improve- ment stems from the fundamental differences in how information is leveraged. To infer intermediate angular values, utilizing not only past information but also future positions of the pendulum can enhance the accuracy of these predictions. this regard, while CRU relies solely on the past positions of the pendulum for its predictions, our model is capable of utilizing both past and future positions.
4.2 SEQUENCE INTERPOLATION & EXTRAPOLATION
Datasets We benchmark the models on two real-world datasets, USHCN and Physionet. The USHCN dataset (Menne et al., 2015) contains 1,218 daily measurements from weather stations across the US with five variables over a four-year period. The Physionet dataset (Silva et al., 2012) contains 8,000 multivariate clinical time-serics of 41 features recorded over 48 hours.
We begin by evaluating the effectiveness of ACSSM on the interpolation task. Following the approach of Schirmer al. (2022) and Rubanova al. (2019), each model is required
Interpolation
See Appendix 1 for details the datasets.

---

Published as a conference paper at ICLR 2025
Table Test MSE (x10-2) for inter/extra-polation on USHCN and Physionet.
Interpolation
Extrapolation
Runtime (scc//ppoch)
Model
USHON
Physionet
USHCN
Physionet
USHCN
Physionet
mTAAN
1,766 0.009
0.208 0.025
2,360 + 0.038
0.340 0.020
7
10
RKN-A
0,009 0.002
0.186 0.030
1.491 + 0.272
0,703 0.050
94
39
GRU-A!
0,090 0.089
0.271 0.057
2.081 = 0.054
0,870 = 0.077
3
GRU-D
0,944 0.011
0.338 0.027
1.718 + 0.015
0,873 0.071
292
5736
1.798 0.009
0.212 a 0.027
2.034 +: 0.005
0,725 0.072
110
791
ODE-RNNT
0,831 = 0.008
0,236 & 0.009
1.955 + 0,466
0.467 +..000000 0.006
81
299
GRU MILL-ODEBB B
0.841 0.142
0.521 + 0.038
5.437 #: 1,020
0.798 0.071
389
90
CRU
0,016 = 0.006
0,182 1 0.091
1,273 4 0.066
0.629 0.093
122 (57.8)"
114 (63.5)
ACSSM (Ours)
0.006 0.001
0.116 0.011
0.941: 0.014
0.627 0.019
2.3
3.8
result from (Schirmer al., 2022).
to infer all time points & E T based on a subset of observations OreT where T C T', For the interpolation task, the encoded observations YieT were assimilated by using the full assimilation scheme for the construction of the accurate smoothing distribution. The interpolation results presented in Table 3, where report the test MSE evaluated the entire time points T'. For all datasets, ACSSM outperforms other baselines in terms of test MSE. It clearly indicate the expressiveness of the ACSSM, trajectories TT sampled from approximated path measure over the entire interval T are contain sufficient information for generating accurate predictions.
Extrapolation We evaluated ACSSM's performance on the extrapolation task following the ex- perimental setup of Schirmer al. (2022). Each model infer values for all time stamps E T', where T denotes the union of observed time stamps T = (E()k())))) and unseen time stamps To = (4+hejk+1::)) i.e,, T' = T U Tu. For the Physionet dataset, input time stamps T covered the first 24 hours, while target time stamps T' spanned the rest hours. In the USHCN dataset, the timeline was split evenly, with 4 = 7 We report the test MSE for unseen time stamps Tn T -T based on the observations on time stamps T. For modeling an accurate filtering distribution, we employed the history assimilation scheme for this task. As illustrated in Table 3, ACSSM consistently outperformed all baseline models in terms of MSE on the USHCN dataset, achieving a significant performance gain over the second-best best model. For the Physionet dataset, ACSSM exhibited comparable performance.
Computational Efficiency To evaluate the training costs in comparison to dynamics- based models that depend on numerical simulations, we re-ran the CRU model on the same hardware used for training our model, indicated by in Table 3. Specifically, we utilized a single NVIDIA RTX A6000 GPU. As illustrated in Table 3, ACSSM significantly lowers training costs compared to dynamics- based models. Notable, ACSSM demonstrated a runtime that was more than 16 66-25x faster than CRU, even though both models aim to approximate I well. It highlight that the latent modelling approach discussed in Sec 3.3 improves efficiency while achieving an accurate approximation of P*.
5 CONCLUSION AND LIMITATION
In this work, we proposed the method for modeling time series with irregular and discrete observations, which we called ACSSM. By using a multi-marginal Doob's h-transform and a variational inference algorithm by exploiting the theory of SOC, ACSSM efficiently simulates conditioned dynamics. leverages amortized inference, simulation-fre latent dynamics framework, and a transformer-based data assimilation scheme for scalable and parallel inference. Empirical results show that ACSSM outperforms in various tasks such as classification, regression, and extrapolation while maintaining computational efficiency across real-world datasets.
Although we present the theoretical basis of our method, thorough analysis in followings remain an open challenge. The variational gap due to the linear approximation may lead to cumulative errors over time, which requires further examination specially for a long time-series such as LLM. Unlike SMC variants (Heng et al., 2020; Lu & Wang, 2024) that use particle-bat based importance weighting to mitigate approximation errors, ACSSM depends on high-capacis neural networks to accurately approximate the optimal control. Incorporating multiple controls, akin to a multi- dynamics approach (Han & Hu, 2020), may alleviate these challenges by enhancing flexibility and robustness.
10

---

Published as a conference paper at ICLR 2025
REPRODUCIBILITY STATEMENT
On the theoretical part, all proofs and assumptions are left to Appendix B due the space constraint. The training and inference algorithms detailed in Algorithm 3 and Algorithm 4, respectively. Additional implementation details, such as data preprocessing are included in Appendix D. We believe these details are sufficient for interested readers to reproduce the results.
ETHICS STATEMENT
In this work, proposed a method for modeling irregular time series for practical applications, suggesting that ACSSM does not directly influence ethical or societal issues in a positive or negative way. However, because ACSSM can be applied health-care datasets, we believe it has the potential to benefit society by improving health and well-being of people.
ACKNOWLEDGMENTS
This work was partly supported by Institute of Information & communications Technology Planning Evaluation(IITP) grant funded by the Korea gMV)) (No.RS-2019-I19007; Artificial Intelligence Graduate School Program(KAIST), No.RS-2022-1120711 Meta-learning Applicable to Real-world Problems, 20.RS-2024-005009277 Global Frontier Lab).
REFERENCES
Abdul Fatir Ansari, Alvin Heng, Andre Lim, and Harold Soh. Neural dontinuous-discrete state space models for eregularly-samplee time series. In International Conference on Machine Learning. 2023.
Elizabeth Louise Baker, Gefan Yang, Michael L Severinsen, Christy Anna Hipsley, and Stefan Sommer. Conditioning non-linear and infinitedimmensional diffusion processes. arXiv preprint arXiv:2402.014344 2024.
P. Baldi. Stochastic Calculus: An Introduction Through Theory and Exercises. Universitext. Springer International Publishing, 2017. ISBN 9783319622262. URL https: /books google.
Philipp Becker, Harit Pandya, Gregor Gebhardt, Cheng Zhao, James Taylor, and Gerhard Neumann. Recurrent kalman networks: Factorized inference high-dimensional deep feature spaces. In International conference machine learning, pp. 554. PMLR, 2019.
Julius Berner, Lorenz Richter, and Karen Ullrich. An optimal control perspective on diffusion-based generative modeling. Transactions on Machine Learning Research, 2024. ISSN 2835-8856. URL https: https:://peeette /ttt:s:///penreview .nt forum?ideeooooo o
Guy E Blelloch. Prefix sums and their applications. School of Computer Science, Carnegie Mellon University, 1990.
Michelle Boué and Paul Dupuis. A variational representation for certain functionals of brownian motion. The Annals of Probability, 26(4):1641-169) 1998.
René Carmona. Lectures on BSDEs, stochastic control, and stochastic differential games with financial applications. SIAM, 2016.
Zhengping Che, Sanjay Purushotham, Kyunghyun Cho, David Sontag, and Yan Liu. Recurrent neural networks for multivariate time series with missing values. Scientific reports, 8(1):6085, 2018.
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
Raphaël Chetrite and Hugo Touchette. Nonequilibrium markov processes conditioned on large deviations. In Annales Henri Poincaré, volume 16, pp. 2005-2057. Springer, 2015.
11

---

Published as a conference paper at ICLR 2025
Nicolas Chopin, Omiros Papaspiliopoulos, et al. An introduction sequential Monte Carlo, volume 4. Springer, 2020.
Nicolas Chopin, Andras Fulop, Jeremy Heng, and Alexandre H. Thiery, Computational doob's h-transforms for online filtering of discretely observed diffusions, 2023.
Junyoung Chung, Caglar Gulcehre, KyungHyun Cho, and Yoshua Bengio. Empirical evaluation of gated recurrent neural networks on sequence modeling. arXiv preprint arXiv:1412.3555, 2014.
Emmanuel de Bézenac, Syama Sundar Rangapuram, Konstantinos Benidis, Michael Bohlke- Schneider, Richard Kurle, Lorenzo Stella, Hilaf Hasson, Patrick Gallinari, and Tim Januschowski. Normalizing kalman filters for multivariate time series analysis. Advances in Neural Information Processing Systems, 33:2995-3007, 2020.
Edward De Brouwer, Jaak Simm, Adam Arany, and Yves Moreau. Gru-ode-bayy bayes: Continuous modeling of sporadically-observe time series. Advances in neural information processing systems, 32, 2019.
P. Del Moral. Feynman- Kac Formulae: Genealogical and Interacting Particle Systems with Applica- tions. Probability and Its Applications. Springer New York, 2011. ISBN 9781441919021. URL https: /books....////
google co. kr/books 21d=IN9
Wei Deng, Weijian Luo, Yixin Tan, Marin Biloš, Yu Chen, Yuriy Nevmyvaka, and Ricky TQ Chen. Variational schr\." odinger diffusion models. arXiv preprint urXiv:2405.04777 2024.
Alexander Denker, Francisco Vargas, Shreyas Padhy, Kieran Didi, Simon Mathis, Vincent Dutordoir, Riccardo Barbano, Emile Mathieu, Urszula Julia Komorowska, and Pietro Lio. Deft: Efficient finetuning of conditional diffusion models by learning the generalised h-transform. arXiv preprint vrXiv:2406.0178/1 2024.
Kingma Diederik. Adam: method for stochastic optimization. (No Title), 2014.
Andreas Doerr Christian Daniel, Martin Schiegg, Souyen-Tuong Duy. Stefan Schaal, Marc Toussaint, and Trimpe Sebastian. Probabilistic recurrent state-space models. In International conference on machine learning, pp. 1280-1289. PMLR, 2018.
Joseph Doob. Conditional brownian motion and the boundary limits of harmonic functions. Bulletin de la Société mathématique de France, 85:431-455 1957.
Wendell H Fleming and Halil Mete Soner. Controlled Markov processes and viscosity solutions, volume 25, Springer Science & Business Media, 2006.
Marco Fraccaro, Simon Kamronn, Ulrich Paquet, and Ole Winther. A disentangled recognition and nonlinear dynamics model for unsupervised learning. Advances neural information processing systems, 30, 2017.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces, arXiv preprint arXiv:2312.00752. 2023.
Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. arXiv preprint 1IX.v:2111.00366, 2021.
Picralberto Guarniero, Adam M Johansen, and Anthony Lee. The iterated auxiliary particle filter. Journal of the American Statistical Association, 1666520) 2017.
Jiequn Han and Ruimeng Hu. Deep fictitious play for finding markovian nash equilibrium in multi-agent games. In Mathematical and scientific machine learning, pp. 221-245. PMIR, 2020.
Carsten Hartmann, Lorenz Richter, Christof Schütte, and Wei Zhang. Variational characterization of free energy: Theory and algorithms. Entropy, 19616)626, 2017.
Jeremy Heng, Adrian N Bishop, George Deligiannidis, and Amaud Doucet. Controlled sequential monte carlo. The Annals of Statistics, 48(5):2004-292) 2020.
12

---

Published as a conference paper at ICLR 2025
Jeremy Heng, Valentin De Bortoli, Arnaud Doucet, and James Thornton. Simulating diffusion bridges with score matching. arXiv preprint arXiv:2111.07243, 2021.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabillsstc models. Advances in neural information processing systems, 33:6840-660 6851, 2020.
Valerii lakovlev, Cagatay Yildiz, Markus Heinonen, and Harri Lähdesmäki. Latent neural ODEs with sparse bayesian multiple shooting. In The Eleventh International Conference Learning Representations, 2023, URL https: openreview net/
Andrew H Jazwinski. Stochastic processes and filtering theory. Courier Corporation, 2007.
Patrick Kidger, James Morrill, James Foster, and Terry Lyons. Neural controlled differential equations for irregular time series. Advances in Neural Information Processing Systems, 33:6696-6707, 2020.
P.E. Kloeden and E. Platen. Numerical Solution of Stochastic Differential Equations. Stochastic Modelling and Applied Probability. Springer Berlin Heidelberg, 2013. ISBN 9783662126165. URL https:tpp//.......t.. //toks/boosk.//oooo... google. co kr hww./boosss..eoo dtt.okksssierr999AAAAeA.
Alexej Klushyn, Richard Kurle, Maximilian Soelch, Botond Cseke, and Patrick van der Smagt. Latent matters: Learning deep state-space models. In M. Ranzato, A. Beygelz- imer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural In- formation Processing Systems, volume 34, pp. 10234-10245. Curran Associates, Inc., wtwpeeiiiinnnnnuuiineeeee neurips cc/paper_ lapper_fille wwlesspppper//000 Pap-0o-Pppppp...eeeee
2021a. URL https: http///poceeeensceeeeeeepces.ese
file/ /alee/544 4etp:22aaa4110188222222000020800
Alexej Klushyn, Richard Kurle, Maximilian Soelch, Botond Cseke, and Patrick van der Smagt. Latent matters: Learning deep state-space models. Advances in Neural Information Processing Systems, 24:10234-10244 2021b.
Dieterich Lawson, Allan Raventós, Andrew Warrington, and Scott Linderman. Sixo: Smoothing inference with twisted objectives. Advances in Neural Information Processing Systems, 35: 38844-38858, 2022.
Christian Léonard.
Stochastic derivatives and generalized h-transforms of markov processes. arXiv preprint arXiv:1102.3172 2011.
Christian Léonard. A survey of the schr\" odinger problem and some of its connections with optimal transport. arXiv preprint arXiv:1308.022 2013.
Alexander K Lew, Tan Zhi-X Chi-Xuan, Gabriel Grand, and Vikash K Mansinghka. Sequential monte carlo steering of large language models using probabilistic programs. arXiv preprint arXiv:2306.03081, 2023.
Mario Lezcano Casado. Trivializations for gradient-basee optimization on manifolds. Advances in Neural Information Processing Systems, 32, 2019.
Xuechen Li, Ting-Kam Leonard Wong, Ricky T. Q. Chen, and David Duvenaud. Scalable gradients for stochastic differential equations. International Conference on Artificial Intelligence and Statistics, 2020.
Guan-Horng Liu, Yaron Lipman, Maximilian Nickel, Brian Karrer, Evangelos Theodorou, and Ricky T. Q Chen. Generalized schrödinger bridge matching. In The Twelfth International Conference on Learning Representations, 2024. URL https: /ttps/////p/eeeeeeee net forum? id-SoismgeX72.
Xingchao Liu, Lemeng Wu. Mao Ye, and qiang liu. Learning diffusion bridges on constrained domains. In The Eleventh International Conference on Learning Representations, 2023, URL
https:/ wtwpeeeeeevveeeeeeeeeeeeeeeeeeee net forum? feemmmid=-iii.. 3=WHlyCaOTbs
Jianfeng Lu and Yuliang Wang. Guidance for twisted particle filter: a continuous-tinm time perspective. arXiv preprint 029399900.0 2024.
13

---

Published as a conference paper at ICLR 2025
MJ Menne, CN Williams Jr. RS Vose, and Data Files. Long-term daily climate records from stations across the contiguous united states, 2015.
Bernt Oksendal. Stochastic Differential Equations An Introduction with Applications. Springer- Verlag, Berlin, Heidelberg, 1992.
Byoungwoo Park, Jungwon Choi, Sungbin Lim, and Juho Lee. Stochastic optimal control for diffusion bridges in function spaces. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. ht wpw.reveevvveeeeeeeeeeeeeeeeeeee ew. net/ forum? ? d=M
Stefano Peluchetti. Diffusion bridge mixture transports, schrödinger bridge problems and generative modeling. Journal of Machine Learning Research, 24(374):1-51, 2023.
Angus Phillips, Hai-Dang Dau, Michael John Hutchinson, Valentin De Bortoli, George Deligiannidis, and Arnaud Doucet. Particle denoising diffusion sampler. arXiv preprint arXiv:2402.06320, 2024.
Syama Sundar Rangapuram, Matthias W Seeger, Jan Gasthaus, Lorenzo Stella, Yuyang Wang, and Tim Januschowski. Deep state space models for time series forecasting. Advances neural information processing systems, 31, 2018.
Lorenz Richter and Julius Berner. Robust sde-based variational formulations for solving linear pdes via deep learning. In International Conference on Machine Learning, pp. 18649-18666. PMLR, 2022.
H. Risken and T. Frank. The Fokker-Planck Equation: Methods of Solution and Applications. Springer Series in Synergetics. Springer Berlin Heidelberg, 2012. ISBN 9783642615443. URL https: //books www.go. kr/books ?id Pid-GSSuCQAAQEAA
Chris G Rogers and David Williams. Diffusiioss, Markov processes and martingales: Volume 2, Itô calculus, volume 2. Cambridge university press, 2000.
Yulia Rubanova, Ricky TQ Chen, and David K Duvenaud. Latent ordinary differential equations for irreullrly-sampled time series. Advances in neural information processing systems, 32, 2019.
S. Särkkä. Bayesian Filtering and Smoothing. Bayesian Filtering and Smoothing. Cambridge University Press, 2013. ISBN 9781107030657. URL https: //boks/bok...oo gwwgle. co kr/
cooks?id=5V1sAAAAOO
S. Särkkä and A. Solin. Applied Stochastic Differential Equations. Institute of Mathematical Statistics Textbooks. Cambridge University Press, 2019. ISBN 9781108693448, URL https:
//books.... google..oooooooooooooooooooo.. co. kr/booksb htt.b/ookkksiidd-uuuuAAAAAAA.
Simo Särkkä and Ángel F Garefa-Fermández. Temporal parallelization of bayesian smoothers. IEEE Transactions on Automatic Control, 66(1):299 306, 2020.
Mona Schirmer, Mazin Eltayeb, Stefan Lessmann, and Maja Rudolph. Modeling irregular time series with continuous recurrent units. International conference on machine learning, pp- 19388-19405. PMLR, 2022.
Yuyang Shi, Valentin De Bortoli, Andrew Campbell, and Arnaud Doucet. Diffusion schrödinger bridge matching. Advances in Neural Information Processing Systems, 36, 2024.
Satya Narayan Shukla and Benjamin Marlin. Multi-time attention networks for irregularly sampled time series. arXiv preprint 100. 2021.
Ikaro Silva, George Moody, Daniel Scott, Leo Celi, and Roger G Mark. Predicting in-hospital mortality of icu patients: The physionet/comppttig in cardiology challenge 2012. In 2012 computing in cardiology, pp. 245-248, IEEE, 2012.
Jimmy T.H. Smith, Andrew Warrington, and Scott Linderman. Simplified state space layers for
sequence modeling. In The Eleventh International Conference on Learning Representations, 2023. URL https: httpp:://pperreeeeett net forum? d 8Hw3Axqks.
14

---

Published as a conference paper at ICLR 2025
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13356. 2020.
Belinda Tzen and Maxim Raginsky. Neural stochastic differential equations: Deep latent gaussian models in the diffusion limit. arXiv preprint arXiv:1905.000 0099009883, 2019.
Ramon Van Handel. Stochastic calculus, filtering, and stochastic control. Course notes., URL http://ww: princeton. cclu/vaa/accm217/AACM117 pdf. 14, 2007.
Francisco Vargas, Andrius Ovsianas, David Fernandes, Mark Girolami, Neil D Lawrence, and Nikolas Nüsken. Bayesian learning via neural schrödinger-föllmer flows. Statistics and Computing, 33(1): 3, 2023.
Mao Ye, Lemeng Wu, and qiang liu. First hitting diffusion models for generating manifold, graph and categorical data. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022. URL https: //ppnneeveww net/forum? 1d=
Sebastian Zeng, Florian Graf, and Roland Kwitt. Latent SDEs on homogeneous spaces. In Thirty-seventh Conference Neural Information Processing Systems, 2023. URL https: / h//ppenreveewnee//eeeeeeeeeeeeee
Qinsheng Zhang and Yongxin Chen. Path integral sampler: stochastic control approach for sampling. In International Conference on Learning Representations, 2022. URL https://
openreview net/ forum? id= _uC22ynuuY.....
Stephen Zhao, Rob Brekelmans, Alireza Makhzani, and Roger Grosse. Probabilistic inference in language models via twisted sequential monte carlo. arXiv preprint arXiv:2404.15544 2024.
15

---

Published as a conference paper at ICLR 2025
A BRIEF REVIEWS ON THE KEY CONCEPTS AND RELATED WORKS
In this section, we provide brief overview of key concepts to help clarify the foundation of our proposed method. Additionally, we review related works to offer deeper understanding.
Bayesian filtering and smoothing (Särkkä, 2013) serve as fundamental state estimation techniques for probabilistic SSMs. Formally, SSMs can be defined as follows:
Probabilistic SSMs
(Latent transition) X1, "x dx,, Xo ~ (Xo) (Observation) Yt. ~ 9t, (y,, X...
(29)
The SSMs consist of the Rd valued latent variable (22) which are assumed to follow a time- inhomogeneous Markov process with Markov transition densities (ei)vett:))- and the observations fy,, ui] assumed to be generated from an observation (emission) density goal is to obtain the filtering/mmoothing distribution, for given observations (yx, Ne(l: given by:
Then, the
pX Hax))
Z(H p HH, ((oo))(X.T))
(30)
where p(Xu.r) is the prior distribution and Z(H..) is normalization constant, defined as:
Z(Hr.) = / p(H.,, Xor II Ix)
(31)
To obtain the filtering/smoothing distribution in (30), it typically relies on recursive Bayesian updates, which scale proportionally with the number of observations, resulting in computational complexity of O(k) in this context. Previous works (Doerr et al., 2018; Becker et al., 2019; Klushyn et al., 2021b) have proposed RNN-based approximate Bayesian inference methods. However, these models generally assume evenly spaced observations, making it challenging to accurately model irregular time series. On the other hand, deterministic linear SSMs such as S4 (Gu et al., 2021), S5 (Smith et al., 2023), and Mamba (Gu & Dao, 2023) have been introduced, demonstrating improved inference efficiency through parallel computing algorithms.
In contrast, we propose probabilistic SSM that enables efficient and powerful modeling of latent systems. Our approach supports parallel computation, leading to significant gains in both training and inference efficiency. Specifically, we incorporate a parallel scan algorithm (Blelloch, 1990) into probabilistic inference, effectively reducing the computational complexity from Okk) Olog k).
Twist function for Conditioned SSMs (30), i.e,, pX,, (H.,) over all / E [0, T] (referred here as the conditioned sampling problem), is necessary to compute the marginal distribution p(X, (He.). However, this involves expec- tation over the marginalized distribution p(X,|H,,) = |H,, ) which is generally intractable. Fortunately, the distribution p(X, (H..) ) can be factorized as (Chopin al., 2020):
To sequentially sample from the smoothing distribution
p(X.(H.,) & ((.,))))) / p(Heeu (x.T)dX
(32)
= =p(H((X))((X) / p(Hes, )
(33)
Thus, by approximating the term (xX,) == problem can be effectively addressed. Here, the intractable term v,, often referred to as twist function, has been the focus of various approximation algorithms proposed in the Sequential Monte Carlo (SMC) literature (Guarniero et al., 2017; Heng et al., 2020; Lawson et al., 2022; Lu Wang, 2024). Recently, these methods have also been adapted to large language models (LLMs) (Zhao et al., 2024), further demonstrating their versatility in solving controlled language generation problems.
NTT, the conditioned sampling
It is worth noting that our h-function defined in (4) serves an instance of such twist function within the Feynman- Kac representation. Essentially, enables us to establish a connection between twisted SSMs and the multi-marginal Doob's h-transform.
16

---

Published as a conference paper at ICLR 2025
Feynman-Kae Models The Peynman-Kac model provides substantial expressive advantages for analyzing SSMs (Chopin et al., 2020). By offering a flexible measure-theoretie framework, allows for efficient representation of conditioned SSMs in (30) through the corresponding Feynman-Ka formulae (Del Moral, 201 Notably, the posterior distribution in (30) can be expressed via a Feynman- Kac model, as described in (2). In the machine learning literature, the Feynman-Kan model has been utilized for abstracting processes such as LLM fine-tuning (Lew et al., 2023) and diffusion based sampler (Phillips et al., 2024). Furthermore, we extend the Feynman-Kae model continuous settings for time-series modeling. For more in-depth understanding, refer to (Chopin et al., 2020).
Continuous Dynamical Models To accurately model irregular time-series data, neural differential equation families have been proposed as an effective paradigm. This is because the latent, continuous dynamics underlying physical time-series can be well-approximated using data-driven methods with parameterized vector fields. Specifically, Neural ODE (Chen et al., 2018) introduced neural network parameterized vector fields for continuous- time dynamical modeling of time-series data. Building on this, Latent-ODE: (Rubanova et al., 2019) proposed latent dynamics by encoding the entire dataset into an initial state, ODE-RNN as continuous encoder alternative to standard RNNs. GRU-ODE-B (De Brouwer et al., 2019) incorporated Bayesian principle into Neural ODEs to enable online updates for new observations. On the stochastic dynamics side, Latent-SDE (Li et al., 2020) introduced a variational bound for posterior inference in SDEs, while Latent-SDEy, SDE (Zeng et al., 2023) focused on stochastic dynamies evolving within homogeneous latent space. These neural differential equations generally rely on numerical simulations with dynamic solvers, which can result in significant computation times for both training and inference.
To model continuous time-series using probabilistic SSMs, CD-SSM (Jazwinski, 2007) extends the discrete transitions of latent variables to stochastic transitions governed by SDEs. This approach has inspired neural network-based CD-SSM methods (Schirmer et al., 2022; Ansari et al., 2023) for han- dling irregularly sampled real-world time-series datasets. Compared to our approach, while Schirmer et al. (2022); Ansari et al. (2023) also leverage locally linear dynamics to improve scalability, they still require numerical approximations to infer Gaussian moments, which limits their ability to fully utilize parallel computation. In contrast, our method successfully leverages parallel computation, as demonstrated by Theorem 3.8, significantly reducing both training and inference costs conditioned state-space modeling. Moreover, our SOC formulation with approximated control a offers ficxibility in sequential modeling, allowing for various information assimilation schemes utilizing powerful neural network architectures, such as transformers.
Doob's ii-transform for Conditioned SDEs In contrast to conditioned SSMs using a twist function for discrete transitions, we propose a conditioned SDEs for continuous transitions to model irregular time-series data. This involves extending the traditional Doob's h-transform to multi-marginal settings. Specifically, the Doob's ' h-transform (Doob, 1957; Rogers & Williams, 2000; Chetrite & Touchette, 2015) is technique in probability theory that modifies the behavior of a Markov process, effectively conditioning it to reach a desired state or outcome. It can be understood as reweighting the paths of a stochastic process to make certain events more probable.
In the machine learning context, the Doob's 's h-transform has been applied to tasks such as diffusion models (Ye et al., 2022; Liu al., 2023; Peluchetti, 2023; Shi et al., 2024; Denker et al., 2024; Park et al., 2024), simulating diffusion bridges (Heng et al., 2021; Baker et al., 2024), posterior approximations (Park et al., 2024), online filtering (Chopin et al., 2023). However, prior works typically focus on single-marging marginal cases, where for finite time horizon [0, T). the conditioning is solely on a terminal constraint P(x, E A) for some set E ((X).
One of our contribution is extends this concept to multi-marginet cases, capturing the continuous dynamies of the posterior distribution conditioned on a collection of observations. Specifically, in Theoren 3.1, we define marginals, XX,,,, E A!...., X., E As) for any A, E ((R)) and E [1 : k]. Since inferring the corresponding lt-function in these multi-marginal settings is more complex, we reformulate this challenge as a SOC problem. This reformulation also requires extensions of existing theoretical results, such as Theorems 3.2 and 3.3. Additionally, establish a tight variational bound in Theorem 3.6 demonstrating that our proposed multi multi-marginal Doob's h-transform can be efficiently approximated using the proposed SOC objective.
conditional SDES where the constraint is given by set of
17

---

Published as a conference paper at ICLR 2025
B PROOFS AND DERIVATIONS
In this section, we present the proofs and derivations for all relevant theorems, lemmas, and corollaries We first restate the core concepts of stochastic caleulus, which will be used without further explanation.
Assumptions. Throughout the paper, we work with probability space (Ω, (Fr Je P), where the filtration {Fa (E).7) supports an -d-valued Fr-adapted Wiener process W, for all & [0, 7). It is important to note that the P-null set is included in Fo. indicating that any event with probability zero at time is measurable in the initial o-algebra.
We assume that b and a satisfy following conditions:
(Lipschitz condition): For any E [0, T). w E Ω and x, x' E Rdd where co > 0 is Lipschitz constant, the functions satisfy the inequality [b(t, w, w. b(t, w, w,x' S co|x x').
(Linear growth condition): For every x R2, the F--progressively measurable processes b(l, x) satisfy E S [6,[=ds] oc and [b(t,x)]
(1 + xx] for L [0, and e > 0.
(Control function): For any / E [0, T. E Q2 x E R. and 0, 0' E O, the control function a is L-- Lipschitz function, |a((, x, 0) a(l, x, 0') < L - 0% Moreover it satisfy E [ lag|ds] < co.
Definition B.1 (Infinitesimal Generator). Let us consider an Itô diffusion process of the form:
dX, = b(t, Xx) + (t) ddW
(34)
Then, an infinitesimal generator of the above diffusion process given by:
E [x)(,,)]
f(x) Vxf b + Trace -ao T vxxf].
A.f = lim
(35)
e10+
t
Definition B.2 (Path Measure). Lef us consider path-sequence of random variables NEeLNN over an interval 0 1 S ... S Ix = T, where each X,, taking values in measurable space (Rd B((R)) with a Borel o-algebra B(Rd). Then, the set of random variables (Xt, Jie(l:N] is given by the probability measure C([O, T), Rd):
N PX E dxto:".. Xax E dxix)= P(dxo) II ,x
(36)
where ( A E B(R"), (x....., = JA pi(x._11 Xt, )xxt, where Pi (x := p(l,, Xt, (i Xt..1) is transition density obtained by solution of the Tokker-Placnk equation (Risken & Frank, 2012);
NN is a sequence of probability kernels from (IRd, B(R2)) to (Rd B(IR)) for any event
(xpr(x:) = Aip: = Vx. + Trace |oo [ TxxP::] 2
(37)
where A is adjoint operator of the generator in (35) and Pt is the Radon-Nikodye density of pt with respect to the Lebesgue measure. By taking N - 0o, we gel the path measure P(X.. E ddxe.r) which describes the weak solutions of the SDE of the form in equation (34).
Lemma B.3 (Itô's formula). Let v(1, be C² in I and C2 in x and let X, be the Itô diffusion process of the form in equation (34). Then, the stochastic process v(1, X,) is also an hi diffusion process satisfying:
do(l, X,) = (0, X,) + Apo(t, (c)( di +
(())aW.
(38)
Theorem B.4 (Girsanov Theorem). Consider the two It diffusion processes of form
dx, b(t, Xt) + (t, dW. E [0, dY, = T(, Y.)dl + Y.) dW., LE [0,T],
(39)
(40)
where both drift functions b, and the diffusion function a assumed to be invertible are adapted to F. and W ( is P-Wiener process. Moreover, consider TP as the path measures induced by
18

---

Published as a conference paper at ICLR 2025
(39). Let us define H₁ o- b) which is assumed to be satisfying the Novikov's 's condition
(i.e., Ep exp 2
1
2 | ds) oo), and the P-martingale process
M, == exp ([',,dw 2 H.lll ds
(41)
satisfies E Mrl 1. Then for the path measure Q given as dQ MydP, the process W, =
W. So H,ds is Q- Q-Wiener process and Y can be represented as
d', = b(t, Y,)dt + (1,, Y.) dWW, LE [O, Therefore Q-law of the process Y, is same P-law of the process X..
(42)
B. PROOF OF THEOREM 3.1
start the section by showing the normalizing property of (fiheeaa in (3), By definition, satisfied that
k
/
(.) II
(),,, (x1,)d((xo.7)) Ep II g (y (v,,,x))) = Z(H+))
(43)
t=1 ad
1=1
where (i) follows from the conditional indenpendency of yt, given x1, for all i [1 R.]. Hence, get the normalizing property:
I (y (xx:) (Xt,)xxe)) II Leg (9t:)
Ep
II Li=1
(xx.))
Er
Ep II Ho(
(44)
Z(Hu))
Theorem 3.1 MMulti-marginal Doob's h- t-transffrrm.. Let us define a sequence functions {he Ne(LE where each 1y [4-1,4) x Rdd - 18..for all i E [1: k), is conditional expectation hi(1, xx)
Ep
II [,(,,,)X XX,, = x.}. where {Jelveuaa is defined (3). Now, we define a function
h: 0, 7 x Rdd + RR by integrating the functions (1eree1::))
k h(t, == E helt, x) (p)
(45)
i=1
Then, with the initial condition pi (dxx) = ha(to, so))oo(dxo). the solution of the following condi- tional SDE inducing the posterior path measure TPA (2):
(Conditioned State) dx; = X;) + Vx log hit, X;) + dWW
(46)
Proof. We start with the interval [t,-1,t, t22) without loss of generality. For all t (t) and for any A,, C B(Rd), the transition kernel of the conditioned process is defined by the transition kernel of the original Markov process P, and h-ffuccion defined in (3.
h, X.,) h, x.)
(X,, E A|X; x) = (x,,,
Pr, (xe, dx;,).
(47)
By the definition of h, the transition kernel P! (xi. A) is probability kernel: [ hr X.,)
[ gg
Phi((e, (x1, dxx.)
(t, P. (x dxc,)
(48)
Sa hi(ti, X., P;)): dxx:))
(49)
ht(t,xe)
Sa f( 1444 (t., X., P.()P, dxi.)
(50)
Ep
I f (y (x,, xx,))XX X,
Ex
I j= fy (ye, X |X+ X2
1,
(51)
Ex
II fs (yr,, X ) X
19

---

Published as a conference paper at ICLR 2025
where (1) follows from the recursion established by the definition of h in Theorem 3.1:
1r,1, xx.) fi(Y,, Xx, (44, Vi V [
(52)
Now, the infinitesimal generator for Ph can be computed for any & E C 7] Rd) and for all F-almost x E IR::
Bpn (t,. X x ((t,x))
A Pe
lim a.0
(53)
phu (x. Pi (x(x,xx.->
Ep
[p(to. Xete) pp(t,
X, =
lim +.0
(54)
8. (1t( X(ts) sp(t, x)] (ft+t)XX,,))))
Ep
X
1(1l(x))
lim
(55)
a.0
Ep [sp( X(+,) sp(1,
Arllts kx((xx)
Xx
lim
(56)
+.0
Be [pp(ts; Xete) (t))xx)) [ha(t + s, X(te)
he(t, x)) X x
Appt + lim
(57)
8J0
st,
where (1) follows from the definition of the infinitesimal generator. Now, shall compute the second term of the RHS of the equation (57) adapted from (Léonard, 2011). By employing basic stochastic calculus, for stochastic process for =p( Xx) and hi, (,((,X)),
Pc+shist+ = spoti, + lo Puuhia fore hindliu + ki, hilets
(58)
Pel4, = polii0 + f Puddin + f Iisuelpp + hiles
(59)
where [p, h,le So dipidhi, is quadratic variation to and h;. By subtracting equation (59) from
equation (58),
Petaliitts -ight = ftte Sudhii,
its
f. ++o hi + (sp, holets - [sp, h.).
(60)
Applying integration by parts leads to the following equation
(pt-. (.)((1,,ts hie) Petshiiets - spehi,t selhitts ha,)) hi,elsets - sp:) (61)
tts (yuu - (p))dhi, +
1. tts (hi haad + Li, hileto [0, hole
where peohiittt - h1,t) gete pidhii Morcover, by applying Itô's formula, get
dep = Aspidl + TdW dhs,e = Adh + (Vzhs,) T.W..
(62)
Therefore, since W, IP-martiggale,
Ep 1(Pt4 (ha,4ts hx = x] (Pu Aul ( (hiju hi.t) Auspudu + (5,, helets (0,, he] X,
(63)
Ep 1/att rtto
(64)
ptss
+ EE (faee (hi,u hr.e))m(PPud
(yu - (x) Au)
E [kp, hsleto - [sp, hale]
(C)
(A)
(B)
20

---

Published as a conference paper at ICLR 2025
For first term (A), the Hölder's inequality with 1/p + 1/9 1 and p. 9 yields,
pte EB [. (yeu -s) 3
l9o -ullduuuu 1)
1/9
Je
(+o
1/p
Bp
E+x
(65)
14 - (Pilld)) [e
pltt
1/0
ghte
1/P
E
|A P2] du
(66)
Given the bounded and continuous function sp the dominated convergence theorem yields (limato 1 [It 1Pu spr /d/ )u/9 (E lim= f lip - --e9 )u/2
0 and since h4 C. ([ts-1,t)), Rd) and boundedness of in Assumptions B, following inequality holds
< <0, 1 Trace [Vx Vxxhaal
a & 10chual + ((V (V = 1
(67)
for any [1-1, (x) and IP almost surely. In other words, supue[t,tta 1
[lakti, 12] oc, Vi
[0, T]. 8 > 0, and > Therefore we get lim E [e (Puu - (pp) Au), = 0 and we can
get the similar result for the second term (B), i.e., lim Hy"" [1, (hi, hi An O. For the last term (C), by the definition of the quadratic variation of yo and h4 yields:
Exx halees - (e, ha]e] EE [.
t+s
-EE EE [ b+s (Vxip)) T Vx Vxllendu
dpuditiu
(68)
Subsequently, by taking the limit from (57), we get the following result:
Ep [lyptt., X,+s) x)) [h,(t + s, X,4s) h,(t, X, =
lim
(69)
+10
st; x)
Ep [ (V Xx)) V Xu, =x
-im
(70)
e10
xt,
= (Vxp(t, Vx log (t, X.),
(71)
Therefore, the infinitesimal generator for ph is defined by:
A p: = Aep: - (Vxyer) Vx log hi,
(72)
which shows that the conditioned SDE for an interval [4:-1 1, 4:) is given by Hence, integrating the generators over the entire interval yields:
xx bb,, x.) + log hy (1, X.) di + dW,
(73)
Ahe App: + Σ [(Vx))) T Vx loghea] 1ts, (t)
(74)
= A + (Vx))t Vx log hr.
(75)
implies that the conditioned SDE for the entire interval [0, 2 is given by:
dx! [b(1, X.) + Vx log h(4, X.)) di + dWW,
(76)
Now, assume that X (x) where (x))) is absolutely continuous with n (x). Then, the path measure induced by X/ be computed
21

---

Published as a conference paper at ICLR 2025
dIP (xx:T) = dup (xo)
II
phy wi (Xtees is dx),,
(77)
1e X(,)
xx)
II
Pit) (Xtmg-n)) dxucil)
(78)
Ri (1i-i,, Xx,
hitt (1i, Xt, (i
dxp((x))
II
Pun (X -x)
(79)
1,
Ntoc did (xo) difp
h441 (1ks Xta))
II foly,, xx, )dP(x)))
(80)
h1 (to,Xo)
where, for all i E [1 k). we define a increasing sequence (1())se(o))) with :(0) = 1,
{(1) = i - 1 + 1 and i(N) i. Hence, for dx)(xx) = h1 (to. x1 (x)) and h&+1 = 1 yields N
di (xo.)) = II (y 1x, )dP(xo:)
(81)
i=1
II gy (xt, (dP..))
(82)
Z(Hr.)
;=1
dF= (Xu:)).
(83)
It concludes the proof.
B.2 PROOF OF THEOREM 3.2
Theorem 3.2 (Dynamic Programming Principle). Let us consider sequence of left continuous functions (e(Lutt) where each Vi E ((: x IR))
2 ||a.11 2 ds log fily,, X + Vttitt X;),, xx 1
Vi x1) min Epir acA
(84)
for all i E [1: k)] and V&+1 = O. Then, for any 0 t u S T, the value function V for the cost
function in (7) satisfying the recursion defined as follows:
V(1,,.)
min Epo aEA
|aaa122s ds Σ log Sily,, (X:) + VI(0)= Xa (
|X; X,
(8.5)
with the indexing function I(u) = maxfi E [1 A:]l,, S w].
Proof. Following the approach used in the proof of the standard dynamic programming principle with the flow property induced by Markov control (Van Handel, 2007), we can apply similar methods to our cost function. We start the proof by establishing the recursion of J. Let define the sequence of left continuous cost functions {Tihe(LAA+))
1/" |l,,1 ds - log (y,,, T:41 X
Jt,, a) == Ex.
(86)
where we denote Bt [-] [ x). for all i E [1 and J&+1 = 0. Since pa is Markov process, satisfying following recursion, for any S 1 S uuu ST.
J(t Xt, = -E
((()))
1
||aa|| ds Σ log fily,, X + (1 X x a) ((rS)))mo))
2
(87)
22

---

Published as a conference paper at ICLR 2025
with the indexing function [(u)) maxfi E [1: ki]t S u). For any E O, there exists a control a' E Att, I] such that (l1,)) + [1, x')
(88)
E E. 3a
((()) 1
l0'l ds E log (iy,,, + Iiw)) Xa Ki
(89)
2
((()) 1
|a'll ds E log fiy,, + VI(w)+1(11(w) X Xirce)
(90)
2
:(1St/5tow
(l(())
min a'EA[t.T]
la'll ds Σ log (x)) + (1(0)+1(41(00) Xircs)
(91)
Since e was arbitrary, limiting + 0 we get:
((()) 1
v(l,x) 2
min e'EA
BLX. po
|aa'l ds
Σ log Jyy x + Vioo+allio X Xirmn)
(92)
For the reverse direction, consider the control á E Alt, T] obtained from integrating:
(ta)) s [tr()) 7].
(93)
a
Then, by following the definition of the value function al, 2
min aAEA(tz(w)
[t,, X,
(94)
(I())
=
Rt pal
| ds Σ log Jy + (( Xe
(95)
2
|a!l
ds Σ
log fily, ( Vi(w) Xi
2
min
"po
2
((ISt:St(())
(96)
(((w))
= min aEA(L.T]
ELxx, pp
|ā,1
Σ log fy (X) + (i Xa X(m..)
(97)
2 vt,,x).
(98)
Combining both inequalities in (92, 97-98), we arrive at the desired result:
Ai(m) J
((t,x) min Ep aελ
2 ||
ds
Elog log f. (y,,, + V1 (tr(w) X |X: x, ((tSt:Str)))
(99)
This concludes the proof.
B.3 PROOF OF THEOREM 3.3
Theorem 3.3 (Verification Theorem). Suppose there exist sequence of left continuous functions x), E (L,2((4-1,, 44), Rd), for all E [1 k), satisfying the following Hamiltonian-lacobi
Bellman (HJB) equation:
0,V,, + A.Vi, + min aEA
(VxVree) as, |lau, =0, 411st<h
(100)
Vi(t, x) log loyr,( (1,) t= Vi [
(101)
23

---

Published as a conference paper at ICLR 2025
where minimum is attained by aj (f.x) (t), Now, define a function [0, x by integrating the optimal controls (ai)ie(),... .ep
R - R
*
(1,,,))= = Σ (4, ((-)))))))
(102)
==1
Then, [t, (e, S [t, X1, a) holds for any (t, E [0, x Bd and a E A. In other words, at optimal control for V in (9).
Without loss of generality, consider & E [f=-1,4] 42). By applying the Itt's formula to the value function V and taking expectation with respect to po, we obtain
Proof.
I 1 (t.. x (t x) + E 1." (o.Vi. A.V., + (VxV T ai,s) ds (103)
where we denote 1 [x x]. By adding the Lagrangian term E/x [ 1 ||a,,1l Il ds
to both sides of (103), we get for the LHS of (103) E
LHS
[1((1,: x(] + E5
2 |lav,ell ds
(104)
xi
Vi(ti, X(,) +
|las ds
(105)
1
|an,al ds log fo (x;) (4, x9:)
(106)
2
[1, x, x))
(107)
where (i) follows from the definition of HIB equation in (11). Now for the RHS of (103), we have:
RHS Vll,,) + E V." (o.V + A.Vi + (V>Vs,)) Joiss) ds + ET
1
|
| ds
2
(108)
=V((t,x] x) + +E J." (.Vi,a + A,Vi. + (V>Vie) ai, + $lag.all ds
(109)
@ v,(t,)) + Epp U." ([VAPia)) 0n. + |10
min aEk
(VxVi.)) 0 +
|l0 ds
(110)
where (i) follows from the definition of HJB equation in (10). Therefore, we get the following result: Ji(t, x, a) = V,(t, x)
+ E5 D. ([VaWin) as + TV
|aa
min agk
(V>Vi,)TTa ass + [lasssl?)) ds
(1i1)
Due to the fact that
1 1.a ([VaVin)) ai, + 2/ |la
1 |la.ool' ds O, (V>De,.) ai, + atk
min
for all & E (() and Ipa almost x, we conclude that [t, x, a) 2 Vtt,, where the equality holds for at, = minaes (V>Ve.) Toa, + 2 ||oi,,11 = Vx,, Additionaly, it implies that x,,,x,,a)) x)). Subsequently, for any E 4)), the recursion (9) yields:
((()))
min E aeA Vtt,x:)
|a ds Σ log fi (X)) + ((14) X )
(112)
[
t2 |aa,, 2 ds log fi(ye, Xi.) V4+ (ti, X.)
min Ebxx ack
(113)
=)((tx,))
(114)
24

---

Published as a conference paper at ICLR 2025
This implies that the optimal control for the value function V in (9) over the interval E [44-1, is a7. Finally, in (9) can be represented as the integrated form
v(l,
Σ Vill, X1) (p
(115)
Therefore, the optimal control for V in (9) becomes att,x) = Ei= (t, ()-.))))))
B.4 PROOF OF LEMMA 3.4
We first restate the Feynman- Kac formula (Oksendal, 1992; Baldi, 2017), which gives a probabilistic representation of the solution to certain PDEs using expectations of stochastic processes. It relies on the fact that the conditional expectation is martingale process.
Lemma B.5 (The Feynman-Kac formula). Let define f E CRR() and C(RR).. Then, function i(t,x,) = Ep S f(s.X.)ds g( ((T)) x1 is solution of the following linear PDE:
Oyh4 + A.h, fht = 0, 0st<T. h(t,x) = g(Xr), =
(116)
(117)
X,). Since h is conditional expectation with respect to P. implies that Y, is martingale process. By applinyg Itô formula, we have:
Proof. Define the process Y, e- f(+ (()x..) )ds
dY, = - f(t, X,)e e S (o.... h(t, Xt)dt + e Se f(x.X.)de dh(tt X.).
(118)
Next, we apply Itō formula to h(t, X.):
Öh Ache dt + Vxhi(t,X,) ddW.
dh(t, Xx)
(119)
Thus, by substituting equation (119) into equation (118), we get
( Oh
dY, - f(t, Y.) +e - 1?" fls,
Adte) dt + St f(s.8,) h dWr
0t
(120)
For Y, to be martingale process, will have zero drift. Therefore, it implies that
Oh
Tt (1, + A,h((,, f(1, X,),L, = O,
(121)
where ((T,XT) = g (XT) by definition. It concludes the proof.
Now let us provide the proof the Lemma 3.4.
Lemma 3.4 (Hopf-Cole Transformation). The h function satisfying the following linear PDE:
Ochi, + Achi, =0, 4-1
(122)
halli, x) Jely (x)h++ (41, 1= li, Vi [1 kk.
(123)
Moreover, for logarithm transformation
log h, satisfying the HJB equation in (10-11).
Proof. The linaer PDE presented in (122-123) can be directly derived from the function hx(t,x;)
Ep T (X,,, )x, Xt This is done by applying the Feynman-Kac formula in Lemma B.5
with f: 0 and II fy (X,,). Now, let consider the function Vi x) = log (t,
(or he(t, = -V((xx)) for all [1 k] and compute
Ochi = -hic...V. Vxlit -hi
VxViss Vxchia haea-haallVV
Vxx vix) (124)
25

---

Published as a conference paper at ICLR 2025
Then, it is straightforward to compute:
( A.hia
hi Vi = -0phi,
(125)
(Vxht.))
Tb T
Trace [xxhtal
(126)
--hii Vite bb Trace (1..ellVVVi his vxx Vi.e)
Vix)
(127)
2
= ---h, be + Trace [he. VV |P] Trace ( VxxVi.l
1
(128)
2
2
where (i) follows from (122). Now, we can simplify (128) by dividing both sides with h >
Vi, = (-V)) be Trace [IV
1
Trace [VxxV1.)
(129)
= + 2 Veell
(130)
Therefore, combining the above results, we have
0,V4, + A.Vie
V =0, Vi(t,, llo f+(y,,lx) Ve+1 (ti, x).
(131)
Since min
(V>V To.e + [oa_e]]]] = 11Vxlll this concludes the proof.
B.5 PROOF OF COROLLARY 3.5
Corollary 3.5 (Optimal Control). For optimal control a induced by the cost function (7) with dynamics (6), it satisfies a* Vx log In other words, we can simulate the conditional SDEs in (5) by simulating the controlled SDE (6) with optimal control o².
Proof. Lemma 3.4 implies that we can obtain the relation (t, log (t, x). Moreover, by following the definition of the optimal control a* in (12) and the value function V in (1 15), it suggest that a* (1, x) = Vx (1, x) for all E [o, T]. Finally, combining the definition of the function in (4) and the results from Lemma 3.4, we can conclude that a (4, - (1,
Vx log (tt x) for all f [0, T7.
B.6 PROOF OF THEBREM 3.6
The Donsker-Varadhan variational principle provides variational formula for the large deviations of functionals of Brownian motion, often related to free-energy minimization problems (Boué & wide range of
Dupuis, 1998). Moreover, through Girsanov's theorem, this principle extends Markov processes, including Itô diffusion processes (Hartmann et al., 2017; Tzen & Raginsky, 2019).
Varaaaan Variational Principle). For a bounded and measurable functions Lemma B.6 (Donsker- Varadhan
W: C(0, T), Rdd) + I, following relation holds:
logg Ex-P (e
min [Ey~a (W(Yo.)) DRr(Q Q
(132)
Proof. The proof relies on the change of measure and the Jensen's inequality:
-log Ex-.
W(Xor) log Ey~
e -W(Yor)
dPP
Y Yo.T)
(133)
dPp
< Ex
W(YnT)) - log
Y (Y.T))
(134)
= ExaO [W(YoT)] + DKL (Q)
(135)
where the equality holds if and only if
Y Yor) c log Bxmo -
(136)
dIP
It concludes the proof.
26

---

Published as a conference paper at ICLR 2025
Theorem 3.6 (Tight Variational Bound). Let us assume that the path measure P induced by (6) for any a E A satisfies ( (P*) < oC. Then, for a cost function J in (7) and μō in (5), it holds:
DRL(P (Ro Exx-uo (J, Xo, o)] [(o) + log (HHc))))))))))))))))))))))))(g 2 0,
(137)
where the objective functiio ((a) (negative ELBO) given by:
((a:) = Ep J'
2 la(s, (;)] ds Σ log (y,, xX)) 2 log loogg(HH.)) ((H..).
(138)
Moreover. assume that [(a) has global minimum a² arg min C(a). Then the equality holds in (16) L(a") = log g ((a,)) and pb = po almost everywhere with respect to po-
Proof. We begin by deriving the KL--ivergence between pa and PP:
DKL (P DKL (Ho MD) + Ex
[KL (-|x)) (-[6)) (PT (-Xx)) [X = Xo.
(139)
dla
-DKL (4)) + 1xx Bpp
log
[ D:T [X
(140)
dlP
dlo
dl
-Duu(Holute + Ex Ex Ep
log
0.T)
+
loy
(XiT))X3
X
(141)
dPp
dTP
where (t) follows from the disintegration theorem (Léonard, 2013). For the second term of the RHS of (141), applying the Girsanov's 's theorem in B.4, we have:
d.a
[ asdWe + 1/²
p,x0
log
(o:T)
= 100x0
2 llael1 dt = =Rpee
2 llov dt
dP
(142)
Moreover, by the definition of P we get:
log
(6.T)
log [x),,
(143)
dF+
Combining the results from (142) and (143), we obtain:
dlpa
dPP
/-- |[a(t,x;))1 dt
log
(XOT) + log
O:T)
= 0.x0
log fi
(x))
dP
dl
(144)
1-- ||att |aatt dt Σ log fi (xx]
0.xo
(145)
= [0 x). a).
(146)
Moreover, by the definition of (ʃi (fi] we get
T
J(0 Xb, a) = Ez..
la(t, X() dt
log f(y., X
(147)
2
T
(
Epa
||a(t, X 2 dt
log (y,, X X = Xo
+ log Z(Hr.),
2
(148)
where (i) follows from normalizing property in (44). Hence, it result the that:
DKL = DKL (pu (Spuullp (m)) + Exo [J (0. Xg. a:]
(149)
DRL (palpp + [(a) log Z (HH,,).
(150)
27

---

Published as a conference paper at ICLR 2025
For KL-divergence term DXL (Holat)) since (xv
hı (to, (x))d obtain
log dua (Xx) duo
(Ko(H))))) = Exprupo
Ex [- log h1 (0, xo)]
(151)
= Exprrp
[10, xo)]
(152)
Exprmo
[J(0, fo, a²)]
(153)
(0) Exorua
[ xo, xo, a²)
+ log ((Hes)
(154)
L(aa)) + log (HHr)):
(155)
where (i) follows from:
k Σ
[(0, Xo, a) = Ex..
la |a(ss 2 ds -
log (x = Xo
(156)
2
T 1
= Epa
la(s, (.) )2 2 ds - Σ log (y X Xg |X Xo
+ +log Z(H,.))
2
5 (0.xx..))
(157)
Thus, we obtain a* = arg mingminae... J(0, xo, a) arg min [0, xo, a)) as log ((H.,) is con- stant. It result that J(, a*) =0,,x00 a²) log Z(H,,) ). Additionally, (ii) follows from
Exp min J(0, Xo. ax) oεA
min Exo
[5 Xp, a)}
min L(a) L(a")
(158)
aEE
since the minimization is independent of the initial condition, as implied by the disintegration theorem (Léonard, 2013). Hence, since the Donsker-Varadhar variational principle in Lemma B.6 with W(Xor)
0:r = E log (y (x;) yields
Egi [
[ 2 [l((t,
k [y,,
x() dt Σ log Iy x
+ log Ep
2 0,
(159)
it implies that DxL (polaç) = C(a*) + log Z(H,,) = 0 for the optimal control a*. In other words, Mo = no almost everywhere with respect to μo and from the variational bound
DKI PP*) L(a) log Z(He,, 2
(160)
the equality holds if and only if a - a* It concludes the proof.
B.7 DERIVATION OF AMORTIZED ELBO IN (28).
Let Og.r is given time-series data. Then, for an auxiliary variable Yo.7 ~ (o(YorloaT)) the ELBO given
log I Py (o,, (r,))(N
log (Do:T) 2 Eg (yorrloor))
(161)
li=1 (v, lor.)
log lo .v.)) log logo(yoor)) log
II [o((e,lon)) ==1
(162)
2 E Ry )
log log (y,)) - L(ax) log II
go go (oe,
(163)
f=1
K
= =Ega(yarroor (our)
log Pg (v.) L(a)
Σ
log go (y 100s)
(164)
1=1
log log.p((o..yu)) L(a)
(165)
28

---

Published as a conference paper at ICLR 2025
where (i) follows from Eg Σi= log (yy,,o)) = 2 0 since do is Gaussian
distribution with constant covariance matrix.
B.8 PROOF OF THEOREM 3.8
Theorem 3.8 (Simulation- free estimation). Let us consider sequences of SPD matrices {A, [Aiheta
that admit the digen-decompositionn =D=ED ED,EE with E E Rdxd and D, diag (RR 2 for all i [1 k), control vectors {as) (es)se( and following control-affine SDEs for all [1 kk:
dX.- = A, + oz) dt odW., E (t(ii1,))))
(166)
Then, with Xo ~ (mo: Eg), the solution of (166) a Gaussian process N(m Evi) with:
m,, = E e 2 m)Dmm
DD elle ds
e
k=1
E (- -ty- D, E10
1
(t)-
D. (1-2200-t))))))) E
2
k=1
where m. = E m Σ.
E ET
E, and à; E E for all E A.).
Proof. Note that since (Aihena are SPD matrices, we can apply the transformation outlined in Remark 3.7 and express the original dynamics (166) in the projected form using the cigenbasis E. Then, for any t E [ti, 11+1), the solution to (166) at time t is given as
[t t, for t t; [ A.lDD ands [ A.D)) A...D..... A.(1) (0, for Sta,
X, - X,
(167)
where m, E Mc Et, = E, à; E for all E [1: k] and W. E W. Given that we have defined Xo [mo, Eg), X = E X,, is Gaussian process for any E [1: k). The first two moments of Gaussian process can be computed from (167). First, since D is diagonal, the integral can be computed as Jr eA.())D:ds ds DD, [I - and M,( J eA.())D. ddW, martingale process with respect pa i.e., Ep- [M] 0. Hence, since O; is time-invariant vector, the mean Ep [X,] = m, for E [L,, (4+1) can computed as
m (A) mi - A. (I (I- cA(DD:) )ā,
(168)
Secondly, for a covariance Ep[(X; m)) m = Ec, we can compute
Σ Eps
2..())) (x.- m + Mi()) (x., - m.
-mi, - -Mi(0)
(169)
Exa (xx, - - ((, m. ) +|MM)
(170)
i
(1). Er, 2
22.((..) D. (I (I- (5.0))
(171)
where (i) follows from the fact that (t(() is martingale and we Itô isometry in (if):
Ex [10 Ep-
[/IAAA 1 ds = 2... D. (1-2((01)))) (172)
-
Hence, we get the Gaussian law of X, at time t E (..). A N(m, Σ.) Furthermore, given recurrence forms of mean (168) and covariance (171), the first moments of Gaussian distribution
29

---

Published as a conference paper at ICLR 2025
for each time steps can be computed sequentially. For mean my have,
m =e- (e-a(ti)Dimi
D (e
(173)
m =e Ee- Az-1 m
(174)
- -i= )--1(1) ) Di
D (I Aa(li)Di)án - D,, '(I (I- (
(175)
(176)
m = 2و ( mo -e (I- in
(177)
k=1
Moreover, for a covariance En,, similar calculation yields
Σn =E - E1. 2Do(t))DD (I- -e2Ao((1))D))
(178)
Σn = c -22 D 115 ((: Σ1
(179)
1 _ -223 A b)-1 ( DJ (50(t))))) 1 (t.).. Dr' (
(180)
2
(181)
Σ. = 21-1
E.
Σ k=1
(1- (1_22e)(())))))
(182)
Now, since D; EA,E" and the orthonormality of E, we can express the mean and covariance in the original canonical basis. For the mean, we get
Em,, E c Ej As
Σ A0
(I- (I-een-((u)D))) âz
(183)
(i)
E- E m
Σ
Ee A (I c
E az
k=1
(184)
(1) e (g-1tyy-(t))A maa
Ei Ay-(t))A ( (I- ack
(185)
ma;
(186)
where (i) follows from mn.
mm. and à; = E a4 for all [L k] and (ii) follows from E e E. Similarly, for the covariance, get
D = E A, - and e- D.
EΣ, E T = E e -2E A, E
-t (I (. E
k=
(187)
2 - 1,10t) A E
STE 2
E
1
Σ
-3 A
A A 'E I 20 2
E
k=
(188)
= A)-(((((
Eio - A A (I- (()))))
(189)
-E.,>
(190)
where (i) follows from Σ, = E En E. By applying the procedure from (168-171) to the original SDE (166), recover the mean and covariance expressions in (185) and (189). This shows that the transformed mean and covariance in (177) and (182) can indeed projected back into the original basis, completing the proof.
30

---

Published as a conference paper at ICLR 2025
Algorithm Parallel Scan for Mean and Covariance
1: Input. Given time stamps T f2..,...... tx). initial mean mi and covariance Eig control policies (e(ss)) matrices (D)e(()))
2: Compute (A,(t,) D,, C. D, Ciheel
3: Set (M)))( (D,, C,as) hee(t:s] and (Sike(tu)) = {(D,, hejLa)
4: Parallel Scan {Mr S (1)))) ParallelScan| (ariicannn,,,,,aa,,,,,,, St, 8)
5: + Algorithm 2 for ParallelScan
6: for 1 to K do in parallel
7: mc = (1)) mc + +
8: Σ, S Ei +
9: end for
Return m.ET, ΣTET
C PARALLEL SCAN
Given an associative operator & and a sequence of elements [st,.".. the parallel sean al- gorithm (Blelloch, 1990) computes the all-prefix-sum which returns the sequence [st, (8t, Ø $t2),... ... (5t, @ Sta @ ... & Stx )] in (log K) time. Since we have verified that moments {m, Er, (e(1iA) of the controlled distributions be estimated by the recurrences in 1768;
m =Dm + C,a,
(191)
Σ. = D, + C,1,
(192)
where, we define
Dj = Ar-1()DD
C - Ai- (I-AAi--( (2d,-1(t)DD)).
(193)
D. e =(
C - -= -2. (I (.-
(194)
and = (1,..., 1) E Rdd For parallel scan, we will define the sequence of tuple [M.)eelu:) such that each element M. (D;, C,a;)) and (Sihee such that each element is S, (Di, C,1) for {m, ]e(l:: and [Σ., ve(s:)) respectively. Now, let us define a binary operator 8:
M, s M, (D. DD, D. o, C,âr)
(195)
S, S, (D, D.. D. C,I+C
(196)
We verify that @ is associative operator since it satisfying:
(M, @ M2) & M, (D, D,, D. o C, + C,a,) @ (D.. Cobu)
(197)
D2 D D.) D. o (D, o o + Cras) + Cuâu)
(198)
(D. D.) D., (D, o D.) C.a. + DD, C, + Crarr)
(199)
M. @ (M, & M.)
(200)
Thus we get (M, @ M.) Ø M, = M, & (M, 8 M,). Moreover, we can get similar results for S,
(S. 2 S.) 8 S. = (D, o D D, C,1+ C,1) & (D., C,1)
(201)
= (D, (D, o D.) D OD. C,1) + C.1) = (DD, D. o D. (D, o D.) C.1 + D. C,1 + C.1)
(202)
(203)
= S, s SS 8 S.))
(204)
Now, both means and covariances along the interval he and { E,, hee(ia can be computed (parallel in time K) using parallel scan algorithm described in Algorithm 1.
31

---

Published as a conference paper at ICLR 2025
Algorithm 2 ParallelSca
1: Input. Sequence of tuples {F1. F2...., Fx). associative operator X
2: Up-Sweep Stage.
3: for d = to [log2 K] do
4: for each sub-tree of height d do in parallel
5:
Let 2d k 2 1 for 0, I,...
6:
if i < K
7:
F; = F.-20 24 F,
8: end if
9: end for
10: end for
Down Down-Sweep Sweep Stage.
12: FÆ = I, where is the identity element for @.
13: for = [log, K - to 0 do
14:
for each sub-tree of height d do in parallel
15:
Let 2d+1 + 2d+1
for 0, 1....
16:
if < K
17:
F Fi-2d = Fi @ F;
F, F,
19: end iF
end for
21: end for
for
1 to do in parallel
23:
F-= F F2@
F.
24:
end for
25: Return Scanned sequence (F F2.....
Fx
D IMPLEMENTATION DETAILS
D.1 DATASETS
Human Activity The Human Activity dataset consists of time series data collected from five individuals performing different activities. Following the preprocessing steps described in (Rubanova al., 2019), we obtained 554 sequences, each with 211 time points and a fixed sequence length of 50 irregularly sampled time stamps. The time range was rescaled to [0, 1]. The classification task involves assigning each time point to one of seven categories: "walking". "falling", "lying", "sitting", "standing up". "on all fours", or "sitting on the ground". The dataset was split into 4, 4,194 sequences for training, 1,049 for validation, and 1,311 for testing.
00-0
000
00000
8:00
8:00
00 600
00
oóo
20
40
60
80
100
Figure 3: Example of the pendulum sequence. (Up) The input image sequences (ohET observed irregular time stamps. (Down) The angular values of (0,) and cos(O.) where 0, represents the angle of the pendulum at time E [0, 100), are used as regression targets.
https:://oo1 http:://///i../r/////// http:///....44//5552.3 under CC BY 4.0
32

---

Published as a conference paper at ICLR 2025
Pendulum The pendulum images were algorithmically generated through numerical simulation as outlined in (Becker et al., 2019). We followed the setup described in (Schirmer et al., 2022), where 4,000 image sequences were generated. Each sequence consists of 50 time stamps, irregularly sampled from T = 100, with each image being a 24 x 24 pixel representation. The sequences was further corrupted by correlated noise process, detailed in (Becker al., 2019). For our experiments, we used 2,000 sequences for training and 1,000 sequences for validation and testing.
USHCN The USHCN dataset (Menne et al., 2015)² includes daily measurements from 2,8 weather stations across the US, covering five variables: precipitation, snowfall, snow depth, and minimum and maximum temperature. We follow the pre-processing steps outlined in (De Brouwer et al., 2019), but select a subset of 1, 168 stations over four-year period starting from 1990, consistent with (Schirmer et al., 2022). Moreover, we make the time series irregular by subsampling 50% of the time points and randomly removing 20% of the measurements. We normalize the features to lie within the range 00, 1] and split into 60% for training, 20% for validation, and 20% for testing.
The Physionet dataset (Silva et al., 2012) contains 8000 multivariate clinical time-series obtained from the intensive care unit (ICU). Each time-series includes various clinical features recorded during the first 48 hours after the patient's admission to the ICU. We preprocess the data as in (Rubanova et al., 2019). Although the dataset contains a total of 41 measurements, we eliminate 4 static features, i.e., age, gender, height, and ICU-type, leaving 37 time-varying features. We round the time-steps to 6-minute intervals, following (Schirmer et al., 2022). We normalize the features to lie within the range [0, 1] and split into 60% for training, 20% for validation, and 20% for testing.
Physionet
Table 4: Training Hyper- parameters
Dataset
Learning Rate
Train Epoch
Time Scale
Batch Size
Rd
of base matrices (L.)
# of parameters
Human Activity
1 10
400
1/221
256
288
256
1.65M
Pendulum USHCN
x 10-3
500
0.1
50
20
15
19.6K
1 10-3
500
0.2
50
20
18.5K
Physionet
1 10-3
500
0.3
100
24
20
28.5K
D.2
MASKING SCHEME
Figure D.2 provides a detailed illustration of the masked attention mechanism described in the assimilation schemes introduced in Section 3.3.
D.3 TRAINING DETAILS
Training For all experiments, except for human activity classification, we followed the same experimental setup as CRU (Schirmer et al., 2022)". For the human activity classification task, used the setup deseribed in mTAND (Shukla & Marlin, 20. For a fair comparison, kept the number of model parameters similar to mTAND for the Human Activity dataset and CRU for the other datasets. The model was trained using the Adam optimizer (Diederik, 2014) in all experiments. For the per-point classification and regression tasks, we applied a weight decay of 1 x 10- and applied gradient clipping for classification task, while no weight decay was used for other tasks. Additionally, to prevent overfitting, we limited the training epochs to 200 for the Physionet extrapolation experiments. The remaining training hyper-parameters are detailed in Table 4.
To estimate the objective function (28). for the decoder (- = (v) Σp), we used Gaus- sian likelihood with a fixed variance of Σ, = 0.01 I. following the approach outlined in (Rubanova et al., 2019) for the Pendulum, USHCN, and Physionet datasets. For the Human Activity dataset, we employed a categorical likelihood. Additionally, for L(a), the initial conditions (mo. Σo) of the latent state were trainable parameters, initialized randomly. The covariance Eg was set using expo- nential transformation, with a small constant (e = 10 10-6 6) added to ensure positivity. The latent state
/ttp /t13/////////////// NDP019 19 under CC
BY 4.0
under ODC-BY 1.0 under AGPL.-3.0
license 10
httpps: htt.ggttthbb...eeeeeeeeeeeeee coob.coomm/mmmmmmm 11bbm TAN under MIT license
33

---

Published as a conference paper at ICLR 2025
Attention
Mask
Masked Attention
Mt.
Z
Weight
Y1
Y:V6
y
Time
History
Attention
Mask
Masked Attention
Zig
Weight
Ail
Vta
Full
yo
Figure Illustraion of Masked Attention. Given the observed time stamps T= (4)) (o) and the unseen time stamps Tu (x), attention scores computed based on the latent observations (ycheT. (Up) The mask of the History assimilation scheme consists of two components: one masks attention scores related future time stamps, and the other masks those related 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 zt, include information up to time 4, while for Full assimilation scheme, &t, incorporates all available information. Finally, latent variables corresponding unseen time stamps are filled with the nearest past latent variable value.
Algorithm Training ACSSM
Input. Time-series o observed over the entire time stamps T TUTa = (h, bx) with the observed time stamps T and unseen time stamps Tu, encoder neural network la. decoder neural network ps, trainable latent parameters (mo. En) and neural networks fo and Tg. for m = -,....
do
Compute 9 (V (VEE7)OIET) by using (25) on observed time stamps T.
Sample latent observations YieT
Parallel computation of objective function [(a) in (16)
if history assimilation then
Estimate latent variables 7IET such that (THH)) illustrated in Fig (D.2)
else if full assimilation then
Estimate latent variables zIET such that TT(HT) illustrated in Fig (D.2)
end if
do in parallel
Compute A; (z(,).) and a; Boz, for all [1 k]] by using (22).
Estimate (m,, E, HET by parallel sean algorithm described in Appendix C.
Sample X BIeT:N(m, Ex) entire time stamps T'.
Compute C(0) E Σt (bi ||a;|| log g, (y,, X by using (16)
Sample latent predictions ÿIeT g(yte7 XigT)) entire time stamps T'. Compute Pa (o (Rte)) by using (26) on entire time stamps T
Optimize ELBO( () by using (28) with gradient descent.
end for
transition matrix A(-) was initialized A E' ED,ET, where was initialized orthonormally following (Lezcano Casado, 2019), and the diagonal matrices (Dihe(uaL) were initialized randomly and passed through negative exponential to keep the values negative. For the potentials {9s {e:)))): we used Gaussian likelihood ge (- = -N(--lr H.), Eg), where rx(x) Mom, + 2, with trainable matrix Mg. transformer output To(H)), and covariance Σ,,
34

---

Published as a conference paper at ICLR 2025
Algorithm 4 Inference with ACSSM
Input. Time-series OtE observed only on the observed time stamps T, unseen time stamps Tu trained model parameters {ó. v,,
Parallel Sampling of latent dynamics
Sample latent observations YxET ~ (0(yie7 by using (25) on observed time stamps if history assimilation then
Estimate latent variables ZieT such that 2, = (TH))) illustrated in Fig (D.2)
else if full assimilation then
Estimate latent variables 2IeT such that Z, = (T(Hrr illustrated in Fig (D.2)
end if
do in parallel
Compute A; = (Z), and ar = Baz, for all i E [1 k] ] by using (22).
Estimate (m,, hET with parallel scan algorithm described in Appendix C.
Sample N trajectorics X i..d 8teT Σx) on entire time stamps T', Sample latent predictions YIeT ~ g(verr [xie)) on entire time stamps T'.
Sample predictions GeT Po 1y (ET))) on entire time stamps T'.
Architecture In all experiments except for the Pendulum dataset, the time series Oo,7 was provided with the observation mask concatenated. We was used dropout rate of 0.2 for Human Activity, while no dropout rate was used for the other experiments. For our method, the networks used for each dataset are listed in below, where d is the dimension of latent space Rd as described in Table 4.
Human Activity (Input size, I=12)
Encoder network qa= Input (2 x2X + Linear (d) + ReLU () >> Linear (d)
Transformer Tg: Input (d) + Masked (d) + FFN (d) + GRU (d)
Weight network fg:Input (d) + Linear (d) + Softmax ()
Decoder network Py: Input (d) + Linear
Pendulum (Input size, I=1x24x24)
Encoder network qs Input (I) + Conv2d 12, kernel_size=5, stride=4, (expool2d(keznel.sis
stride=2 Conv2d (12, 12, kernel.size=3, stride=2, padding=1) + R)LU() + MaxPool22 (kernel.size=2, stride=2) Flatten (108) Linear (d)
padding=2) + ()
+ ReLU() ()
Transformer Ty: Input (d) 4 x Masked Attn (d) FFN (d)] GRU (d)
Weight network Tn:Tnp (d) Linear (d) + Softmax ()
Decoder network Py: Input (d) + Linear (d) + ReLU () + Linear (d)
USCHCN & Physionet (Input size, I=5 for USHCN and I=37 for Physionet)
Encoder network qø: Input (2xI) + Linear (d) () + LayerNorm (d)
+ Linear (d) + ReLU + (ayerNorm(d) + Linear (d) + ReLU +
LayerNorm (d) Linear (d)
Transformer Tę: Input (d) x Masked Attn (d) FFN (d) GRU
Weight network fg:Input (d) + Linear (d) + Softmax
Decoder network pw: Input (d) + Linear (d) + ReLU() + sinear(d) 7
() + Linnear (d) (e Linear
FFN
Input (d) + (ayerNorm(d) + Linear (d) + GeLu() 4 Linear (d) +
Residual (Tnput )
Masked Attn
Input (0, K, V, mask=M) + Normalize (0) Linear (0) Linear (K)
+ Linear (V) + Attention (Q, K) + Masking (M) + Softmax (d)
Dropout + Matmul (V) LayerNorm (d) Linear (d)
Resi dual (Q)
35


Reviews:
[{"review_id": "hGZTpu0xu7", "reviewer": "Reviewer_jj93", "summary": "The authors propose ACSSM approach for modeling irregular time series which uses continuous-discrete state space models (CD-SSMs).  The authors extend doob's-h transform to multi-marginal case and solve the problem using stochastic optimal control. The formulation leads to a  evidence lower bound (ELBO) and the authors propose a VI based loss function to model the irregular time series. \n\nAuthors make linear approximation for the SDEs, which allows them to be able to perform simulation free estimation and exploit transformers for parallel computations. Overall, the proposed approach is faster and performs better than existing simulation based approaches, neural differential equations and other RNN based approaches tailored to handle irregular time series.", "questions": "- How many learnable parameters are use for each method? Was each method trained with a similar paramter budget?\n- Given several approximations, how does the method perform on larger datasets ?", "limitations": "", "rating": 8, "confidence": 2, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- The paper makes significant theoretical contributions which should be relevant other researchers in this field. \n- The proposed approach has impressive empirical results showing better results with faster training while being theoretically grounded.", "weaknesses": "One of the weakness of this work that I see is that it makes several simplifying approximations to make the solution faster/tractable. Authors already acknowledge this but it would have been nice to understand if there are any practical trade-off due to these approximation.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "RjklifnMoH", "reviewer": "Reviewer_6X8u", "summary": "The paper introduces the Amortized Control of Continuous State Space Model (ACSSM) to handle irregularly sampled time-series. ACSSM aims to model the path measure of trajectories in a latent space, conditioned on observations in the data space. By using a latent space, the model captures a flexible and structured representation of the underlying dynamics that generate observed data, which is especially useful for irregularly sampled observations. To construct this conditional path measure, the authors introduce a novel multi-marginal Doob’s h-transform. This extension of the traditional Doob’s transform induces a class of stochastic differential equations (SDEs) that define the desired path measure in the latent space. However, simulating these SDEs directly is computationally infeasible due to the need for expensive normalization constants and conditional expectations. To overcome this challenge, the authors leverage stochastic optimal control to define a variational objective that approximates the optimal control needed to produce the conditioned dynamics. To further enhance computational efficiency, they propose working with affine linear SDEs with known Gaussian perturbation kernels, allowing simulation-free estimation of the latent trajectories and significantly speeding up the inference process.", "questions": "- The introduction lists the contributions, but it could benefit from more intuition to guide the reader through the chain of thought. For instance, there is no explanation of what a Feynman-Kac model is or how it facilitates sequential analysis.\n\n- The multi-marginal Doob’s  h-transform is a central component of the approach, but its presentation lacks intuitive guidance, which is reflective of the overall style in the paper. Adding more accessible explanations would enhance understanding.\n\n- Some sentences have syntax errors or missing words. I recommend proofreading the text to improve readability.", "limitations": "", "rating": 8, "confidence": 2, "soundness": 4, "presentation": 2, "contribution": 4, "strengths": "The diffusion and stochastic optimal control (SOC) perspective for time-series modeling offers a compelling alternative to classical recurrent network methods, as it allows for continuous, flexible representations of latent dynamics that are suited to handle irregular timesteps.\n\nThe paper demonstrates superior performance on real-world datasets. Additionally, the efficiency gains in training time make it feasible for high-dimensional data, where classical SDE-based methods might be expensive.", "weaknesses": "The absence of a Related Work section limits the reader’s ability to understand how ACSSM compares to existing time-series modeling approaches, particularly those used in irregular time-series contexts, and those compared in the experiments section (e.g., recurrent networks, attention mechanisms, and previous SDE-based models).\n\nThe paper leans heavily on measure-theoretic concepts and complex SDE formulations, which could make it difficult for readers not specialized in these areas. More accessible explanations or visual intuitions could enhance understanding.\n\nIt remains to be seen how the model scales in practice with very high-dimensional or complex latent spaces, as affine SDE simplifications may reduce the expressiveness of the dynamics in these settings.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "7HbizlBOy5", "reviewer": "Reviewer_UjeK", "summary": "The authors introduce a way of amortizing the controller of a state space model (ACSSM) to make it compatible with irregular time series. To do this, they generalize the single-marginal Doob’s $h$-transform to the multi-marginal case. To simulate the resulting continuous dynamics, they use VI to get an ELBO which is then optimized. To make this tractable, they assume that the latent dynamics are locally linear, and use a neural network to get expressive latent dynamics this way. The authors provide a theoretical analysis for this work, showing that the ELBO they obtain is tight, and then offer several experiments where this method shows improved performance over comparable baselines, for both per time classification/regression as well as sequence interpolation/extrapolation.", "questions": "-Did you guys try using the affine linear drift function for the latent dynamics? How much better does the learned NN drift function do? \n\n-It is sometimes unclear where you are using the full and when you are using the history attention mechanism. It is claimed that the authors perform better than CRU because they aren’t just using past information, but also future information. However, this seems like an unfair comparison, am I misunderstanding something?\n\n-How much does the \"Parallel Computation\" stuff increase the speed of inference? Is this a major contributor of why the method is so fast compared the other methods.", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "-Novel method of solving Continuous State Space model which learns linear dynamics to accommodate irregular time series. \n\n-Theoretical analysis for the provided algorithm, including the derivation of an ELBO used to solve the dynamics.\n\n-Demonstration of the algorithm on real time series, showing improved performance over other methods, in both Test MSE and also compute time (<5 secs).\n\n-Addresses limitations, such as the errors accumulated from the linear approximation.", "weaknesses": "-Some of the background was a bit hard to follow. As someone who is relatively unfamiliar with the literature in this area, I found that some of the stuff in the methods section were not explained (it’s possible that it was common knowledge). I tried looking at the appendix for a more fleshed out explanation and i still can’t say I’m confident I understand everything going on.\n\n-The paper argues that this method is scalable, and one thing I wish I could understand better is how this would work on environments with more complex dynamics, and possibly even partially observable environments. Did you guys try anything along this route?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "pyQMWkl67M", "reviewer": "Reviewer_vbhs", "summary": "The author models irregular time series through a latent controlled SDE. The author first showed that the problem can be solved by a conditioned SDE, and then established the link between such conditioned SDE to a controlled SDE that optimizes a specific form of cost function. Finally, solving the HJB equation for the cost function can be obtained by simulating the controlled SDE through the Feynman-Kac theorem. The simulation was then simplified by analytically solving for the first and second moments at each timestep. Overall the paper is theoretically well-motivated with good empirical results.", "questions": "1. Can the authors elaborate on the choice of the cost function in Eq(7)? How should one interpret this cost function? Aside from the theoretical benefit, is there any intuition on the choice of this cost function? \n2. Can the author explain the usage of full attention in this scenario? If the full attention is applied to estimate the latent dynamic, does it mean the controls $\\alpha_{t_i}$ for the controlled SDE is informed by the future y-observations $o_{t_j}, j > i$? Intuitively the masked attention makes sense to me, but I am unsure about the application of full attention. \n3. It seems like the main experiments are all obtained with the full attention scheme. Can the author provide the result of those experiments using masked attention as an ablation study? When is the masked attention used?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper is theoretically motivated with novel ideas and methods. \n2. The presentation of the paper is clear. \n3. The paper provides improvements on empirical results. \n4. The paper provides clear instructions on replicating the experiment.", "weaknesses": "1. The motivation for specific implementation details is not clear", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Oral)