Conditional Lagrangian Wasserstein Flow for Time Series Imputation

Many thanks for the insighful comments.

Weakness 1:

Our proposed method can also be implemented via an SDE sampler by adding a time-dependent volatility term (please see Eq. (12), Sec 3.2, Page 4), but in order to achieve the optimal empirical performance in terms of RMSE and MAE on the time series imputation tasks, we opt for an ODE sampler.

Weakness 2:

In practice, we found that the Euler method has similar performance to other numerical methods while requiring less memory and computational time. Therefore, we chose to use the Euler method in our implementation.

Weakness 3:

Yes, the proposed method can be used for time series forecasting tasks as well. To this end, we can view the forecasting task as a conditional generation problem, in which one can just use the future states as the data generation target and the current state as the conditional information.

Weakness 4:

Eq. (15) is derived based on the assumption of an overdamped Langevin dynamic system, which still represents the drift term in Eq.(1). For computational convenience, we choose to omit the diffusion term to formulate an ODE sampler. Nevertheless, the theoretical analysis given in the paper still applies to this case.

Questions:

Although the VAE-based potential function can be seen as a sufficient statistic for the imputation target, it is generally not accurate enough to be used alone to represent the data generation dynamics in time series imputation tasks.

Limitations:

Since our method is implemented via an ODE sampler CRPS may not be a suitable evaluation metric.
Furthermore, [1] and [2] point out that CRPS may overlook the model's performance on each dimension, which could lead to an incorrect assessment of overall model performance.

Reference: [1] Koochali, Alireza, Peter Schichtel, Andreas Dengel, and Sheraz Ahmed. "Random noise vs. state-of-the-art probabilistic forecasting methods: A case study on CRPS-Sum discrimination ability." Applied Sciences 12, no. 10 (2022): 5104.

[2] Biloš, Marin, Kashif Rasul, Anderson Schneider, Yuriy Nevmyvaka, and Stephan Günnemann. "Modeling temporal data as continuous functions with stochastic process diffusion." In International Conference on Machine Learning, pp. 2452-2470. PMLR, 2023.

Many thanks for the valuable feedback .

Weakness 1:

To project the intermediate samples from the Euclidean space to the Wasserstein space, We adopt the technique descried in [1] by drawing mini-batches of the initial noise and target samples to compute the corresponding optimal transport maps.
In terms of Wasserstein metrics, we use the earth mover distance (EMD) for compute the deterministic optimal transport maps, which turns out to perform better than other metrics. Alternatively, we can also adopt the Sinkhorn distance to compute the stochastic optimal transport maps with entropy regularization. For the practical implementation, we use the Python Optimal Transport (POT) library [2].

Weakness 2:

This is because we adopt an ODE sampler in our method. Although ODE samplers are deterministic, they have other advantages compared to SDE samplers. For example, they are usually easier to implement than SDE samplers. Moreover, compared to CSDI and other diffusion-based methods, our method requires less simulation steps and samples to achieve good performance on time series imputation tasks.

Weakness 3:

We can also implement the proposed framework via an SDE sampler. The optimization objective becomes the KL divergence between the target SDE and the prior Wiener process, which is aligned with the frame work of the Schrodinger Bridge problem. However, in practice, we found that the ODE sampler has better empirical performance in terms of RMSE and MAE compared to the SDE sampler.

Weakness 4:

We have proofread the manuscript again and corrected the typos.

Questions:

Since the Euclidean distance is not suitable for measuring the distance between two probabilistic distributions, meaning that we cannot find the shortest path between =the initial distribution and the target distribution. As a result, performing the data generation processes in Euclidean space may lead to slow convergence and poor sample quality. Therefore, we need to perform the data generation processes in the Euclidean space instead of the Wasserstein space.

Limitations:

We conduct additional experiment using the synthetic dataset. And the experimental results demonstrates the effectiveness of the Rao-Blackwellization. Please take a look at the uploaded PDF for the details.

Test imputation results on synthetic data ($5$ trials, values are multiplied by $10^2$.

References:

[1] Alexander Tong, Kilian FATRAS, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector-Brooks, Guy Wolf, and Yoshua Bengio. Improving and generalizing flow-based genera431 tive models with minibatch optimal transport. Transactions on Machine Learning Research,432, 2024.

[2] RFlamary, Rémi, Nicolas Courty, Alexandre Gramfort, Mokhtar Z. Alaya, Aurélie Boisbunon, Stanislas Chambon, Laetitia Chapel et al. "Pot: Python optimal transport." Journal of Machine Learning Research 22, no. 78 (2021): 1-8. Website: https://pythonot.github.io/

Many thanks for the insightful comments.

Weakness 1:

The value of the potential function's variance do effect the Rao-Blackwellized sampler. If the variance is too small the effect of the Rao-Blackwellization is negligible, if the variance is too large the performance will decrease. In future research, we plan to explore more advanced methods for formulating the potential function for time series imputation tasks.

Weakness 2:

Since the VAE model is trained on clean target data, we can leverage its denoising/self-correction effect to force the reconstructed intermediate samples closer to the ground truth during the sampling process. An ideal choice of the potential function should be convex.

Weakness 3:

If the VAE model's performance is too poor, it cannot be considered a sufficient statistic and, consequently, will not aid in the data generation process.

Questions:

Unfortunately, in the current implementation we cannot compute the likelihood of the samples obtained because the sampler is deterministic. We will leave this problem for future research.

Limitations:

This is true, although we tried the methods proposed in existing literature, we do not find a satisfactory method to estimate the diffusion term to achieve competitive performance on time series imputation tasks. Therefore, we plan to explore better methods for estimating the diffusion term in the flow matching model as part of our future work.

Limitations:

''' Other ODE models have had no problem with diversity. I would suggest that this is actually due to performing gradient ascent on the density of the VAE.''' Other ODE models also have the problem with diversity to some extent. The usage of the VAE is not relevant on this issue, which, on the contrary, will slightly increase the randomnesses of the sampling procedure as the VAE is a probabilistic model. To clarify the statement, let's consider an ODE sampler described by

$

X_T = X_0 + \int_0^T \mu (X_t,t)dt,

$

and an SDE sampler described by

$

X_T = X_0 + \int_0^T \mu (X_t,t)dt + \int_0^T \sigma (X_t,t)dW_t.

$

From the above equations you can see that the randomness of the ODE sampler only arises from $X_0$ (the initial random noise) as the sampling process is deterministic. Thus the diversity of the samples obtained is limited as the velocity model $\mu (X_t,t)$ is biased towards the training data, which may suffer the overfitting issue. Please note that the training targets are only some finite samples of the target distribution but not the target distribution itself. On the contrary, the SDE-based sampling procedure can result in different samples even the initial value $X_0$ is the same datapoint due to the existence of the diffusion term $\sigma_t$.

Test imputation results on synthetic data ($5$ trials, values are multiplied by $10^2$.

Test imputation results on PM 2.5 with different simulation steps ($5$ trials)

Test imputation results on PhysioNet 0.1 with different simulation steps ($5$ trials).

Test imputation results on PhysioNet 0.5 with different simulation steps ($5$ trials).

Test imputation results with error bars ($5$ trials).

  How is this implemented practically?

Yes, conditioned on an initial point $x_0$, an ODE has limited diversity, however conditioned on $p_0$, it does not.

How does this prove that the VAE is a sufficient statistic for $X_{t+1}$? Or does this not need to be shown.

There are no additional experiments on the "Rao-Blackwellization" that I can see.

$\nabla_x U(X_t) = \nabla \log N(t, 1) = t - x_t$.

any non-zero quantity will bias our estimate $X_{t+1} | X_t = \mu_t^\star(x) = 1.$ 

Thank you very much for the constructive feedback and insightful comments

Weaknesses (Quality) 1:

We have conducted the ablation study on the impacts of sampling steps of CSDI and CLWF. The results show that CLWF has better performance when the simulation steps are small. Please also take a look at the uploaded PDF for the detailed experimental results.

Weaknesses (Quality) 2:

We have conducted additional experiment using the synthetic dataset. And the experimental results demonstrates the effectiveness of the Rao-Blackwellization, please see the uploaded PDF.

Weaknesses (Quality) 3:

We have added the error bars to the experimental results as suggested, please see the uploaded PDF.

Weaknesses (Quality) 4:

Compared to the SOTA methods, whose performances are already nearly optimal, the performance of our method has been significantly improved. In this situation, the Rao-Blackwellization can further enhance performance. Even without substantial improvements in values, it still demonstrates the effectiveness of our method.

Weaknesses (Clarity):

As far as I can tell this “Rao-Blackwellized” sampler is fundamentally biased (while the original sampler is not), and it therefore cannot possibly be a Rao-Blackwellized sampler. 

The Rao-Blackwell theorem only requires the base/original sampler is unbiased and the statistic is sufficient (not necessarily unbiased). According to the law of Total Expectation, the new sampler is also unbiased. Hence, one cannot say the new sampler is not Rao-Blackwellized.

Justification for Rao-Blackwellization:

Let's first consider a dynamic system which can be described by the following SDE:

$

dX_t &= \mu_t(X_t,t)dt + \sigma_t dW_t,

$

where $\mu_t(X_t,t)$ is the velocity and can be solved via flow matching.

Alternatively, if we assume the above system is a over-damped Langevin dynamic system, then system can be described by the following SDE:

$

dX_t &= \nabla V_t(X_t,t)dt + \sigma_t dW_t

$

where $\nabla V_t(X_t,t)$ is the drift.

Essentially, we have two different mathematical interpretations for the same system.
Note that $\mu_t(X_t,t)$ (using flow matching) and $\nabla \mu_t(X_t,t)$ (using the potential function) are computed via different numerical methods, but ideally, they both suffice to represent the same system independently. We also assume the Gaussian potential function constructed via the VAE model is accurate enough.

Since $X_{t+1}$ is the parameter we want to estimate, according to the Rao-Blackwell theorem, we construct the base sampler $\mathcal{S}$ based on the flow matching model $dX_t = \mu_t(X_t,t)dt + \sigma_t dW_t$ which is unbiased for $X_{t+1}$. And $\mathcal{T}(X_t)$ based on $dX_t = \nabla V_t(X_t,t)dt + \sigma_t dW_t$ is a sufficient statistic for $X_{t+1}$, meaning that $\mathcal{T}(X_t)$ contains all the necessary information for estimating $X_{t+1}$ but does not dependent on $X_{t+1}$. Afterwards, we define the new Rao-Blackwellized sampler $\mathcal{S}^* = \mathbb{E}[\mathcal{S}|\mathcal{T}]$ by letting the base sampler condition on the sufficient statistic.

Now we prove that the new sampler is unbiased, i.e., $\mathbb{E}[\mathcal{S}^*] = \mathbb{E}[\mathbb{E}[\mathcal{S}|\mathcal{T}]] = \mathbb{E}[\mathcal{S}]$, based on the Law of Total Expectation (or the Tower Rule).

$

\mathbb{E}[\mathcal{S}^*] = \mathbb{E}[\mathbb{E}[\mathcal{S}|\mathcal{T}]] & = \iint s p(s|t)ds p(t)dt \nonumber \ & = \iint s p(s|t) p(t)dsdt = \iint s p(s,t) dsdt = \int s p(s) dt = \mathbb{E}[\mathcal{S}].

$

Since $\mathbb{E}[\mathcal{S}]$ is unbiased, $\mathbb{E}[\mathbb{E}[\mathcal{S}|\mathcal{T}]]$ is unbiased as well.

By doing so, we can leverage the denoising/self-correction effect of the VAE model (considering that it is trained on clean data), which enables the reconstructed intermediate samples closer to the ground truth during the sampling process.
If the above explanation is still not intuitive enough, one can think about a Kalman filter, whereby the first model is the system's dynamics and the second model is the noisy observation. The noise observation can be used to improve the trajectory estimation during tacking as long as it is sufficient. Please see new experimental results reported in the uploaded PDF. From the experiments results, we can see that the variances are reduced, and the improvement can be verified by calculating the RMSE. For the future work, we will explore new techniques to further improve the stability of the Rao-Blackwellized sampler.

Weaknesses (Significance):

In this paper, we propose a novel framework based on Lagrangian dynamics and optimal transport for time series imputation and a novel sampler to enhance its performance. Our contributions are both theoretical and applied. We have also added new experimental results as the reviewer suggest, please refer to the uploaded PDF for more details.

Questions 1:

```
    How is this a Rao-Blackwellization, and what does the paper 
```

Since the comment does not fully display, we will only reply the first part of the question. Please see the above response to the issue regarding the Rao-Blackwellization.

Questions 2:

We have corrected the typos as suggested.

Questions 3:

paragraph line 146 is confusing. this equation is strange given that $p(x^tar)$ and $p(x^input)$ are dependent if taking batches. 

The authors are confused by the comment as paragraph line 146 is not an equation. Doe the reviewer refers to Eq. (14)? This is the probabilistic coupling in context of optimal transport. Please refer to [1].

Reference:

[1] Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. In The Eleventh International Conference on Learning Representations, 2022.

Thank you for your valuable feedback.

Weakness 1:

Indeed, CLWF is trained through the simulation-free training of the Schrödinger Bridge. However, in this work, based on the Lagrangian mechanics framework, we specifically proposed the Rao-Blackwellized sampler to further enhance the model's performance for time series imputation tasks.

Weakness 2:

The ODE sampler is adopted to optimize the performance of our method for the time series imputation tasks. However, an SDE sampler is also applicable within our proposed framework.

Since our method is implemented via an ODE sampler CRPS may not be a very suitable evaluation metric.
[1] and [2] suggest that CRPS could lead to an incorrect assessment of overall model performance due to the overlook of the model's performance on each dimension of the data.

For the future work, we plan to explore more advanced methods to improve the diversity of the samples to have better empirical performance regarding CRPS .

Weakness 3:

The datasets used in our experiments, PM 2.5 and Physio, are both highly representative in the field of time series imputation. We have also included additional experimental results using synthetic data to further demonstrate the advantages of our methods. Please refer to the uploaded PDF for more detailed experimental information.

References:

[1] Koochali, Alireza, Peter Schichtel, Andreas Dengel, and Sheraz Ahmed. "Random noise vs. state-of-the-art probabilistic forecasting methods: A case study on CRPS-Sum discrimination ability." Applied Sciences 12, no. 10 (2022): 5104.

[2] Biloš, Marin, Kashif Rasul, Anderson Schneider, Yuriy Nevmyvaka, and Stephan Günnemann. "Modeling temporal data as continuous functions with stochastic process diffusion." In International Conference on Machine Learning, pp. 2452-2470. PMLR, 2023.