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

We thank the reviewer for appreciating our research and acknowledging its significance. The questions raised have been addressed in the responses provided below.

1. Motivation for specific implementation.

• To address the reviewer’s concern, we have included a brief note on the key concepts and related works in Appendix A. We hope that this section helps improve the reviewer’s understanding of our approach. If there are still specific aspects or details the reviewer’s find unclear, please let us know, and we would be happy to address them further in the revised version.

2. Cost functional for SOC problem.

• The cost function in Eq (7) is central to our approach as it serves to balance two objectives: minimizing control effort and aligning the generated dynamics with observed data. The intuition behind this choice comes from SOC theory, where the goal is to control a system's trajectory in such a way that it closely follows a desired path (in our case, the latent dynamics inferred from observations) while minimizing control energy.
• The first term, $\frac{1}{2}||\alpha_t||^2$, represents the control energy, which penalizes large deviations in control inputs for the principle of least action to regularize the control effort, thereby encouraging smoother trajectories. In practice, we can interpret it as a regularization term to help generalization by discouraging overly complex control signals and ensuring that the model relies on inherent patterns in the data rather than excessively adjusting controls to fit noise.
• The second term, $-\log f_i(\mathbf{y}_{t_i} | \cdot)$, ensures that the controlled dynamics are consistent with the observed data points over the time interval [0, T] by adjusting the dynamics to maximize the likelihood of observations.
• This interpretation is closely related to variational inference, where the objective is to approximate a target posterior distribution by optimizing a variational distribution while including a regularization term (often a Kullback-Leibler divergence). In our formulation, we extend this approach to the path space, where the prior distribution is given by the dynamics in Eq (1), and the target posterior distribution is defined in Eq (2). Here, the control function \alpha plays the role of adjusting the prior dynamics to approximate the posterior path measure, in SOC problem Eq (7) (technically Eq (16)) often referred to as KL-control problem, where the first term acts as KL-regularization term while the second term acts as target log-likelihood. The cost function in Eq (16) acts as a variational bound to achieve this alignment while maintaining regularization over the entire trajectory.

3. Assimilation scheme.

• The use of full assimilation means that the controls for the controlled SDE indeed incorporate information over a full observation $\mathbb{o}_{t \in [0, T]}$. The type of assimilation is inspired by filtering/smoothing in the standard SSM algorithm, where the smoothing algorithm typically incorporates full observation. This design choice was made to allow for non-causal data assimilation, which improves the expressiveness and flexibility of our model in capturing long-range dependencies. We provided a conceptual illustration of the proposed information assimilation (full/history) scheme in Figure 4.

4. Use of history attention scheme.

• We would like to clarify that, in our implementation, we utilize a history assimilation scheme (which we believe the reviewer refers to as "masked attention") for the sequence extrapolation task, because the task involves forecasting based solely on past data. It prevents the use of the full assimilation scheme, which leverages future observations. With history assimilation, our approach effectively behaves like a traditional filtering method which typically relies on past observation to infer the latent states. As a result, in sequential extrapolation, the performance gap between our model and CRU narrows. Because CRU is genuinely a filtering method that typically relies on past observation and with history assimilation, our approach effectively behaves like a filtering method. In this setting, it is worth noting that our method achieves comparable performance with CRU, while avoiding numerical trajectory inference by leveraging parallel state estimation based on Theorem 3.8.

We sincerely appreciate the reviewer’s interest in our research and acknowledgment of its significant contributions. We have addressed the concerns raised by the reviewer below.

1. Background

• We have added further explanations in Appendix A to provide a deeper understanding of the theoretical foundations. We hope these revisions help clarify the concepts and address the reviewer’s concerns.

2. Scalability on more complex dynamics

• Thank you for your insightful suggestions to improve our work. We have addressed this concern in the general response. For a more detailed explanation, kindly refer to that section.

3. Neural network drift function

• In general, using an NN drift can be effective for modeling complex latent dynamics. However, employing an NN drift requires a numerical solver, which can lead to instability as the time-series length increases [1]. In such cases, we expect simplifying the dynamics to derive a closed-form solution can actually be more beneficial, as it leads to more stable learning. We believe that our experiments demonstrate this by showing superior performance compared to other baselines that utilize NN drift functions, as highlighted in the Human Activity Classification task in Section 4.1.

4. Assimilation schemes

• With the exception of specific settings such as sequence extrapolation, it is generally expected that utilizing all available data for predictions will result in better performance. In our experiments on classification and regression in Section 4.1, all methods, except for CRU, leverage this advantage. In cases where future data is accessible, not utilizing it for predictions, as in the case of CRU, could lead to inefficiencies.

• In this regard, we do not see this as an unfair comparison, but rather as an indication of the inherent limitations in CRU's modeling approach, which presents a challenge that needs to be addressed. In contrast, our method introduces a control formulation that offers an effective solution to overcome this limitation

5. Parallel computation

The key factors that make our algorithm faster compared to other methods are: (1) the use of locally linear dynamics with a diagonalizable matrix $A$, and (2) the incorporation of a parallel scan algorithm.

• (1) While locally linear dynamics alone do not directly guarantee faster inference, they do provide computational advantages compared to neural differential equation models. By using locally linear dynamics, we can simplify the heavy numerical simulations often required, reducing the problem to two ODEs as described in equations (18-19). This allows us to leverage efficient matrix operation tricks, making computations more efficient. However, it is important to note that this approach still involves ODE integration and matrix operations, which can be computationally intensive.
• (2) The significant speedup in our method comes from applying the parallel scan algorithm, particularly in cases where the matrix $A$ is diagonalizable, as shown in Theorem 3.8. Typically, inferring the Bayesian posterior distribution requires $\mathcal{O}(k)$ computation for $k$ observations due to the sequential update. By leveraging the parallel scan algorithm, we can reduce this time complexity from $\mathcal{O}(k)$ to $\mathcal{O}(\log k)$, allowing the processing of the data simultaneously. This results in a substantial reduction in computation time, especially for large-scale datasets, making our method significantly faster than others that rely on sequential processing.

[1] Iakovlev et al., “Latent Neural ODEs with Sparse Bayesian Multiple Shooting”

We sincerely appreciate the reviewer's interest in our research and acknowledgment of its significant contributions. We have provided detailed responses in the subsequent.

1. Related Works and intuition of background.

• In response to the reviewers' suggestions, we have clarified key concepts (such as probabilistic SSMs, the Feynman-Kac model, and Doob’s h-transform) in the revised manuscript. Additionally, we have added the related work to include relevant literature. For more details, please refer to Appendix A of the revised version. We hope this addresses the reviewer’s concerns and provides the clarity the reviewer was seeking.

2. How the model scales in practice with very high-dimensional or complex latent spaces.

• Thanks for the valuable suggestion to enhance our work. We have addressed the reviewer’s point in the general responses. Kindly refer to that part for our detailed answer.

3. Sentence.

• We appreciate the feedback regarding readability. We will conduct a thorough proofreading pass to address any syntactic errors, enhance sentence flow, and improve overall clarity.

Dear Reviewer 6X8u,

As the discussion period is coming to a close, we kindly invite you to share your feedback on our response and consider revising your score if you find it appropriate.

Best regards,

Authors

We sincerely appreciate the reviewer's interest in our work and recognition of its contributions. Detailed responses are provided in below.

1. Performance on larger datasets.

• We thank the reviewer for the valuable suggestions to enhance our work. This point has been addressed in the general responses section. Please refer to that section for a detailed explanation.

2. Parameter budget.

• To ensure a fair comparison, all methods, including ACSSM and the baseline models, were trained using a comparable parameter budget. As outlined in Table 4, we maintained similar parameter counts to those of the baseline methods. For instance, we matched the parameter count to [1] for the Human Activity dataset and to [2] for the other datasets.

[1] Shukla & Marlin, “ Multi-time attention networks for irregularly sampled time series.”
[2] Schirmer et al., “Modeling irregular time series with continuous recurrent units”