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

Authors: Byoungwoo Park, Hyungi Lee, Juho Lee
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
Keywords: stochastic optimal control, variational inference, state space model, irregular time series

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.

Decision: accept (oral)
Reviewer scores: 8, 8, 8, 8

Reviews and discussion:
<thread>
creation date:
1729180614522

thread_type:
official_review

<thread_object>
creator: reviewer_vbhs

<content>
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.

strengths:
1. The paper is theoretically motivated with novel ideas and methods. 
2. The presentation of the paper is clear. 
3. The paper provides improvements on empirical results. 
4. The paper provides clear instructions on replicating the experiment.

weaknesses:
1. The motivation for specific implementation details is not clear

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? 
2. 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. 
3. 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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730698273023

thread_type:
official_review

<thread_object>
creator: reviewer_ujek

<content>
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.

strengths:
-Novel method of solving Continuous State Space model which learns linear dynamics to accommodate irregular time series. 

-Theoretical analysis for the provided algorithm, including the derivation of an ELBO used to solve the dynamics.

-Demonstration of the algorithm on real time series, showing improved performance over other methods, in both Test MSE and also compute time (<5 secs).

-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.

-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?

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? 

-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?

-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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730824153451

thread_type:
official_review

<thread_object>
creator: reviewer_6x8u

<content>
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.

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.

The 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).

The 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.

It 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.

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.

- 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.

- Some sentences have syntax errors or missing words. I recommend proofreading the text to improve readability.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731312942632

thread_type:
official_review

<thread_object>
creator: reviewer_jj93

<content>
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. 

Authors 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.

strengths:
- The paper makes significant theoretical contributions which should be relevant other researchers in this field. 
- 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.

questions:
- How many learnable parameters are use for each method? Was each method trained with a similar paramter budget?
- Given several approximations, how does the method perform on larger datasets ?
</content>
</thread_object>

</thread>
<thread>
creation date:
1731766423949

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Author response to all reviewers

comment:
We sincerely thank the reviewers for thoughtful feedback. We are pleased that all reviewers recognized the importance of our theoretical contributions, and acknowledged the novelty and impact of our proposed method. We are also encouraged by the multiple reviewers who found the paper to be clearly presented and well-structured.


We believe our work makes a significant contribution by integrating stochastic processes and stochastic optimal control frameworks to advance continuous state space modeling. While our primary focus is on addressing the challenges associated with irregular time-series data, we believe our approach has the potential to benefit a broad range of sequential data applications.

----
Although all reviewers acknowledged the technical novelty of our work, they raised concerns regarding (1) the scalability to larger datasets and (2) the insufficient clarification of key theoretical concepts. In response to their valuable feedback, we have provided some general responses below.


**1. The scalability to larger datasets.**

* While approximations (such as locally linear approximation) may seem to limit the expressiveness of the model, our empirical results indicate that, for the datasets we tested (e.g., Physionet, USHCN, and Human Activity), which are (we believe) already large-scale (e.g., Physionet has 8000 observations with 37 dim) real-world datasets benchmarked by several prior works, it does not lead to any sacrifice in model performance. Moreover, the Pendulum dataset contains 4,000 observations with a 576 dimension (24x24 pixels), and exhibits partial observability due to being corrupted by a correlated noise process, as illustrated in Figure 3.

* This is because, by leveraging amortized inference, our approach efficiently scales to high-dimensional and complex latent spaces by decoupling representation learning from the latent dynamics. This helps mitigate some of the trade-offs associated with these approximations, thereby preserving both efficiency and accuracy.

* Additionally, while linear approximations might appear to reduce expressiveness, state-space models utilizing linear dynamics have already been successfully applied in [1] to large-scale datasets. Moreover, [2] observed that the performance of generative models may depend not so much on the linearity of the forward process, but rather on the complexity of the backward generation process, this aligns with our observation that linearizing the prior distribution is less critical compared to capturing the complexity of the posterior distribution. To address this, we leverage the expressiveness of transformer architectures to achieve sufficient flexibility in modeling the posterior. 

* For more complex modeling needs, we believe that utilizing more suitable architectures will be the key to further improvements. We believe that the balance we have struck between computational efficiency and model flexibility can extend to even more complex time-series data, similar to these successful cases. In fact, we plan to adapt our method to large-scale medical datasets in future work.

**2. Insufficient clarification of key theoretical concepts.**

* In line with the reviewer’s suggestion, we have included a comprehensive Related Work section and added further brief explanations in **Appendix A** in the revision. We sincerely hope this section will improve the grasp of both the reviewers and potential readers, thereby increasing their confidence in understanding our paper.

----
Notable changes are highlighted in magenta in the revised manuscript. We have made every effort to address all concerns raised, as detailed in the individual responses below.

----

    [1] Smith et al., “Simplified state space layers for sequence modeling”
    [2] Deng et al., “Variational Schrödinger Diffusion Models”
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper proposes the Amortized Control of continuous State Space Model for continuous dynamical modeling of time series for irregular and discrete observations. It extends Doob's h-transform to the multi-marginal setting, and defines a variational inference algorithm with a tight ELBO. To speed up training and inference, ACSSM assumes locally linear latent dynamics and employs transformers for parallel computation.

All reviewers praise its novel and significant theoretical contribution to modeling continuous state space with irregular observations. The empirical studies on real-world datasets show superior performance in both computational efficiency and accuracy.

I recommend accepting this paper based on the consensus from reviewers.