Conditional Lagrangian Wasserstein Flow for Time Series Imputation

Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 13:29:21

Abstract: Time series imputation is important for numerous real-world applications. To overcome the limitations of diffusion model-based imputation methods, e.g., slow convergence in inference, we propose a novel method for time series imputation in this work, called Conditional Lagrangian Wasserstein Flow. The proposed method leverages the (conditional) optimal transport theory to learn the probability flow in a simulation-free manner, in which the initial noise, missing data, and observations are treated as the source distribution, target distribution, and conditional information, respectively. According to the principle of least action in Lagrangian mechanics, we learn the velocity by minimizing the corresponding kinetic energy. Moreover, to incorporate more prior information into the model, we parameterize the derivative of a task-specific potential function via a variational autoencoder, and combine it with the base estimator to formulate a Rao-Blackwellized sampler. The propose model allows us to take less intermediate steps to produce high-quality samples for inference compared to existing diffusion methods.  Finally, the experimental results on the real-word datasets show that the proposed method achieves competitive performance on time series imputation compared to the state-of-the-art methods.

Main paper: 
 introduction: 
 Time series imputation is essential for various practical scenarios in many fields, such as transportation, environment, and medical care, etc. Deep learning-based approaches, such as RNNs, VAEs, and GANs, have been proved to be advantageous compared to traditional machine learning methods on various complex real-words multivariate time series analysis tasks [18] . More recently, diffusion models, such as denoising diffusion probabilistic models (DDPMs) [20] and score-based generative models (SBGMs) [43] , have gained more and more attention in the field of time series analysis due to their powerful modelling capability [26, 32] .Although many diffusion model-based time series imputation approaches have been proposed and show their advantages compared to conventional deep learning models [44, 11, 12] , they are limited to slow convergence or large computational costs. Such limitations may prevent them being applied to real-world applications. To address the aforementioned issues, in this work, we leverage the optimal transport theory [47] and Lagrangian mechanics [3] to propose a novel method, called Conditional Lagrangian Wasserstein Flow (CLWF), for fast and accurate time series imputation.In our method, we treat the multivariate time series imputation task as a conditional optimal transport problem, whereby the random noise is the source distribution, the missing data is the target distribution, and the observed data is the conditional information. To generate new data samples efficiently and accurately, we need to find the shortest path in the probability space according to the optimal transport theory. To this end, we first project the original source and target distributions into the Wasserstein Submitted to 38th Conference on Neural Information Processing Systems (NeurIPS 2024). Do not distribute.space via sampling mini-batch OT maps. Afterwards, we construct the time-dependent intermediate samples through interpolating the source distribution and target distribution. Then according to the principle of least action in Lagrangian mechanics [3] , the optimal velocity function moving the source distribution to the target distribution is learned in a self-supervised manner by minimizing the corresponding kinetic energy. Moreover, to further improve the model's performance, we learn the task-specific potential function by training a Variational Autoencoder (VAE) model [22] on the observed time series data to build a Rao-Blackwellized trajectory sampler.Finally, CLWF is assessed on two real-word multivariate time series datasets. The obtained resultsshow that the proposed method achieves competitive performance and admits fast convergence compared with other state-of-the-art time series imputation methods.The contributions of the paper ares summarized as follows:• We present Conditional Lagrangian Wasserstein Flow, a novel conditional generative framework based on the optimal transport theory and Lagrangian mechanics;• We propose a Rao-Blackwellized trajectory sampler to enhance the data generation performance by incorporating the prior information;• We develop the practical algorithms to solve the time series imputation problem via a conditional generative approach;• We demonstrate that the proposed method has competitive performance on time series imputation tasks compared other state-of-the-art methods. 
preliminaries: 
 In this section, we concisely introduce the fundamentals of stochastic differential equations, optimal transport, Shrödinger Bridge, and Lagrangian mechanics.We treat the data generation task as an initial value problem (IVP), in which X 0 ∈ R d is the initial data (e.g., some random noise) at the initial time t = 0, and X T ∈ R d is target data at the terminal time t = T . To solve the IVP, we consider a stochastic differential equation (SDE) defined by a Borel measurable time-dependent drift function µ t : [0, T ] × R d → R d , and a positive Borel measurable time-dependent diffusion function σ t : [0, T ] → R d >0 . Accordingly, the Itô form of the SDE can be described as follows [36] :EQUATIONwhere W t is a Brownian motion/Wiener process. When the diffusion term is not considered, the SDE degenerates to an ordinary differential equation (ODE). However, we will use the SDE for the theoretical analysis as it is more general.The Fokker-Planck equation (FPE) [40] describing the evolution of the marginal density p t (X t ) reads:EQUATIONwhere ∆p t = ∇ • (∇p t ) is the Laplacian. In fact, both Eq. eq:sde and Eq. eq:fpe reveal the dynamics of the system and serve as the boundary conditions for the optimization problems we will introduce in later sections with different focuses. The differences are when the constraint is Eq. ( 1), the formalism is Lagrangian, which depicts the movement of each individual particle; while when the constraint is Eq.( 2), the formalism is Eulerian, which depicts the evolution of population.The optimal transport (OT) problem aims to seek the optimal transport plans/ maps that moves the source distribution to the target distribution [47, 41, 38] . In the Kantorovich's formulation of the OT problem, the transport costs are minimized with respect to some probabilistic couplings/joint distributions [47, 41, 38] . Let p 0 and p T be two Borel probability measures with finite second moments on the space Ω ∈ R d . Π(p 0 , p T ) denotes a set of transport plans between these two marginals. Then, the Kantorovich's OT problem is defined as follows:EQUATIONwhere Π(p 0 , p T ) = π ∈ P(X × Y) : (π x ) # π = p 0 , (π y ) # π = p T , with π x and π x being two projections of X × Y on Ω. The minimizer of Eq .( 3), π * , always exist and referred to as the optimal transport plan.Note that the R.H.S of Eq. ( 3) can also include an entropy regularization term D KL (π∥p 0 ⊗ p T ), then the original OT problem transforms into the entropy-regularized optimal transport (EROT) problem with Eq. (2) as the constraint, which frames the transport problem better in terms of convexity and stability [13] In particular, from a data generation perspective, p 0 is some random initial noise and p T is the target data distribution, and we can sample the optimal transport maps in a mini-batch manner [46, 45, 39] .The transport problem in Sec. 2.2 can be further viewed from a distribution evolution perspective, which is particularly suitable for developing the flow models that model the data generation process. For this reason, the Shrödinger Bridge (SB) problem is introduced [25] . Assume thatΩ ∈ C 1 ([0, T ], R d ), P(Ω)is a probability path measure on the path space Ω, then the goal of the SB problem aims to find the following optimal path measure:P * = arg min P∈P(Ω)D KL (P∥Q) subject to P 0 = q 0 and PEQUATIONwhere the Kullback-Leibler (KL) divergenceD KL (P∥Q) = log dP dQ dP, if P ≪ Q, +∞,otherwise, and Q is a reference path measure, e.g., Brownian motion or Ornstein-Uhlenbeck process. Moreover, the distribution matching problem in Eq. ( 3) can be cast as a dynamical SB problem as well [19, 24, 28] :EQUATIONsubject to Eq. (1) or Eq. ( 2), where θ is the parameters of the variational drift function µ t .In this section, we formulate the data generation problem under the framework of Lagrangian mechanics [3] . Let p t and ṗt = dpt dt be the density and law of the generalized coordinates X t , respectively. K(p t , ṗt , t) is the kinetic energy, and U(p t , t) is the potential energy, then the corresponding Lagrangian isEQUATIONWe assume that Eq. ( 6) is lower semi-continuous (lsc) and strictly convex in ṗt in the Wasserstein space. The kinetic energy K(x t , µ t , t) and potential energy U(p t , t) are defined as follows, respectively:EQUATIONEQUATIONwhere U t (X t ) is the potential function. Then the action in the context of Lagrangian mechanics is defined as follow:EQUATIONAccording to the principle of least action, the shortest path is the one minimizing the action, which is aligned with Eq. ( 4) in the SB theory as well. Therefore, we can leverage the Lagrangian dynamics to tackle the OT problem for data generation. To solve Eq. ( 6), we need to satisfy the stationary condition, i.e., the Euler-Lagrangian equation:EQUATIONwith the boundary condition dXt dt = µ(X t , t), q 0 = p 0 , q T = p T . 
conditional lagrangian wasserstein flow for time series imputation: 
 In the section, building upon the optimal transport theory, the Shrödinger Bridge problem, and Lagrangian mechanics introduced in Sec. 2, we propose Conditional Lagrangian Wasserstein Flow, which is a novel conditional generative method for time series imputation.Our goal is to impute the missing time series data points based on the observations. For training, we adopt adopt a conditionally generative approach for time series imputation in the sample space R K×L , where K represents the dimension of the multivariate time series and L represents sequence length. In our self-supervised learning approach, the total observed data x obs ∈ R K×L are partitioned into the imputation target x tar ∈ R K×L and the conditional datax cond ∈ R K×L .As a result, the missing data points x tar can be generated based on the conditions x cond joint with some uninformative initial distribution x 0 ∈ R K×L (e.g., Gaussian noise) at time t = 0, then the imputation task can be described as:x tar ∼ p(x tar |x cond 0), where the total input of the model isx input 0 := (x cond , x 0 ) ∈ R K×L×2 .To solve Eq. ( 7), we need to sample the intermediate variable X t in the Wasserstein space first. To do so, the interpolation method is adopted to construct the intermediate samples. According to the OT and SB problems introduced in Sec. 2, we define the following time-differentiable interpolant:EQUATIONwhere Ω ∈ R d is the support of the marginals p 0 (X 0 ) and p T (X T ), as well as the conditional p(X t |X 0 , X T , t).For implement I t , first, we independently sample some random noise X 0 ∼ N (0, σ 2 0 ) at the initial time t = 0 and the data samples X T ∼ p(x tar ) at the terminal time t = T , respectively. Afterwards, the interpolation method is used to construct the intermediate samples X t ∼ p(X t |X 0 , X T , t), where t ∼ uniform(0, T ) [30, 2, 45] . More specifically, we design the following sampling approach:EQUATIONwhere γ t ∼ N (0, σ 2 γ ) is some random noise with variance σ γ injected to the target data samples for improving the coupling's generalization property, and α(t) ≥ 0 is a time-dependent scalar.Note that Eq. ( 12) can only allow us to generate time-dependent intermediate samples in the Euclidean space but not the Wasserstein space, which can lead to slow convergence as the sampling paths are not straightened. Hence, to address this issue, we need to project the interpolations in the Wasserstein space before interpolating to strengthen the probability flow. To this end, we leverage the method adopted in [46, 45, 39] to sample the optimal mini-batch OT maps between X 0 and X T first, and perform the interpolations according to Eq. ( 12) afterwards. Finally, we have the joint variablex input t := (x cond , x t )as the input for computing the velocity of the Wasserstein flow.To estimate the velocity of the Wasserstein flow µ t (X t , t) in Eq. ( 1), the previous methods that require trajectory simulation for training can result in long convergence time and large computational costs [9, 37] . To circumvent the above issues, in this work we adopt a simulation-free training strategy based on the OT theory introduce in Sec. 2.2 [30, 46, 2] , which turns out to be faster and more scalable to large time series datasets.Since we can now draw mini-batch interpolated samples of the source distribution and target distribution in the Wasserstein space using Eq. ( 12), we can model the variational velocity function using a neural network with parameters θ. Then, according to Eq. ( 1), the target velocity can be computed by the difference between the source distribution and target distribution. Therefore, the variational velocity function µ θ (x input t , t) can be learned byEQUATION≈ arg minEQUATION)Eq. ( 14) can be solved by drawing mini-batch samples in the Wasserstein space and performing stochastic gradient descent accordingly. In this fashion, the learning process is simulation-free as the trajectory simulation is not needed.Moreover, note that that Eq. ( 13) also obeys the principle of least action introduced in Sec. 2.4 as it minimizes the kinetic energy described in Eq. ( 7). Therefore, it indicates that the geodesic that drives the particles from the source distribution to the target distribution in the OT problem described in Sec. 2 is found as well, which enables us to generate new samples with less simulation steps compared to standard diffusion models.So far, we have demonstrated how to leverage the kinetic energy to estimate the velocity in the Lagrangian described by Eq. 6. Apart from this, we can also incorporate the prior knowledge within the task-specific potential energy into the dynamics, which enables us to further improve the data generation performance. To this end, let U (X t ) : R d × [0, T ] → R be the task-specific potential function depending on the generalized coordinates X t [48, 37, 34] . Therefore, we can compute the dynamics of the system byEQUATIONSince the data generation problem in our case can also be interpreted as a stochastic optimal control (SOC) problem [4, 17, 35, 50, 21, 5] , then the existence of such U t (X t ) is assured by Pontryagin's Maximum Principle (PMP) [16] .To estimate v t (X t , t), according to the Lagrangian Eq. ( 6), we assume that the potential functiontakes the form U t (X t ) ≈ -log N (X t | Xt , σ 2 p ), where Xt the learned mean and σ 2 p is the pre-defined variance. As a result, the derivative is∇ x U (X t ) = Xt-Xt σ 2 p. In terms of practical implementation, we parameterize ∇ x U (X t ) via a Variational Autoencoder (VAE) [22] . More specifically, we pre-train a VAE on the total observed time series data X obs . Afterwards, the reconstruction discrepancies of the VAE are used to approximate the task-specific v ϕ (X t , t) depending on X t :EQUATIONwhere VAE(X t ) represents the reconstruction output of the pre-trained VAE model with input X t , and σ 2 p is treated as a positive constant for simplicity. In this manner, we can incorporate the prior knowledge learned from the accessible training data into the sampling process formulated by Eq. ( 14)to enhance the data generation performance.To generate the missing time series datapoints, we first formulate an unbiased ODE sampler S(X t , µ θ t (X t , t), t) for X t+1 with the Euler method and µ θ t (X t , t) learned by Eq. ( 14) (which means the diffusion term in Eq. 1 is omitted). Alternatively, one can also adopt the SDE sampler by using the Euler-Maruyama method. Nevertheless, to ensure achieve the best imputation performance, we choose the ODE sampler for implementation. Note that the ODE sampler alone is good enough to generate high-quality samples for time series imputation.Now we can construct a Rao-Blackwellized trajectory sampler [8] for time series data imputation using Eq. 14 and Eq. 16. To this end, we first treat S(X t+1 |X t , µ θ t (X t , t), t) be the base estimator for X t+1 with E[S 2 ] < ∞ for all X t+1 . And we assume T (X t , v ϕ t (x t , t), t) is a sufficient statistic for X t+1 based on Eq. 16, even it is not a very accurate estimator for X t+1 . As a result, we can formulate a new trajectory sampler S * = E[S|T ] to generate the missing time series data. Then according to the Rao-Blackwell theorem [8] , we haveEQUATION)where the inequality is strict unless S is a function of T . Eq. 17 suggests we can construct a more powerful sampler with smaller errors than the base ODE sampler S using Rao-Blackwellization.The overall training process of CLWF is illustrated in Fig. 1 , which consists of the following stages. First, the total observed data x obs are partitioned into the target data and conditional data for training. Next, the data pairs of x tar and x 0 are sampled from the target dataset and random Gaussian noise, respectively. Then, the data pairs are projected into the Wasserstein space by sampling the corresponding OT maps. After that, the intermediate variable x t is sampled through interpolation using Eq. ( 12). We can approximate the target velocity dXt dt by computing x tar -x0. Subsequently, we use the joint distribution of the conditional information x cond and the intermediate variable x t ,x input as the total input to feed the variational flow model µ θ t to compute the velocity. And the flow matching loss defined by Eq. ( 14) is minimized by stochastic gradient descent.Furthermore, to incorporate the prior information of into the model, we can choose to train a VAE model on the total observed data x obs . This is used to estimate the derivative of the task-specific potential function according to Eq. ( 16), which can be further utilized to construct a more powerful Rao-Blackwellized sampler for inference.For inference, at time t = 0, we sample the initial random noise x 0 and conditional informationx cond to formulate the joint variable x input . Note that during the trajectory sampling x 0 will evolve over time, while x cond remain invariant. We use x input as the input of the flow model µ θ t to compute the velocity. Afterwards, we sample the new x t using the Euler method. If we perform Rao-Blackwellization, then x t is fed to the VAE model for computing the derivative of the potential function, and x t is updated again using the Euler method. The above process will be repeated until reach its convergence. Moreover, we can sample multiple trajectories using different initial random noise, and the averages as the final imputation results. Finally, the proposed training and sampling procedures are presented in Algorithm 1 and Algorithm 2, respectively.Require: Terminal time: T , max epochs, observed data X obs , parameters: θ and ϕ. while epoch < max epochs do sample t, (x 0 , x T ), and OT maps; sample x t according to Eq. ( 12); minimize Eq. ( 14); end while if Rao-Blackwellization then train a VAE model on X obs . end if Algorithm 2 Sampling procedure Require: initial time t = 0, terminal time:T = 1, Euler method step number: N . sample initial noise X 0 ∼ N (0, σ 2 0 ) while t < T do X t = X t + µ θ t (x input , t) T N if Rao-Blackwellization then X t = X t + v ϕ t (x t , t) T N end if t = t + T N end while4 Experiments 
datasets: 
 We use two public multivariate time series datasets for validation. The first dataset is the PM 2.5dataset [51] from the air quality monitoring sites for 12 months. The missing rate of the raw data is 13%. The feature number K is 36 and the sequence length L is 36. In our experiments, only the observed datapoints are masked randomly as the imputation targets.The other dataset we use is the PhysioNet dataset [42] collected from the intensive care unit for 48hours. The feature number K is 35 and the sequence length L is 48. The missing rate of the raw data is 80%. In our experiments, 10% and 50% of the datapoints are masked randomly as the imputation targets, which are denoted as PhysioNet 0.1 and PhysioNet 0.5, respectively.For comparison, we select the following state-of-the-art timer series imputation methods as the baselines: 1) GP-VAE [18] , which is combines a VAE model and a Gaussian Process prior; 2) CSDI [44] , which is based on the conditional diffusion model; 3) CSBI [12] , which is based on the Schrödinger bridge diffusion model; 4) DSPD-GP [7] , which combines the diffusion model with the Gaussian Process prior.In terms of the choices of architectures, tboth the flow model and the VAE model are built upon Transformers [44] . We use the ODE sampler for inference and sample the exact optimal transport maps for interpolations to achieve the optimal performance. The optimizer is Adam and the learning rate: 0.001 with linear scheduler. The maximum training epochs is 200. The mini batch size for training is 64. The total step number of the Euler method used in CLWF is 15, while the total step numbers for other diffusion models. i.e., is CSDI, CSBI, and DSPD-GP are 15 (as suggested in their papers). The number of the Monte Carlo samples for inference is 50. The standard deviation σ 0 for the initial noise X 0 is 0.1, and the standard deviation σ γ for the injected noise γ t 0.001. The coefficient σ 2 p in the derivative of the potential function is 0.01.We assess the proposed method on PM 2.5, PhysioNet 0.1 and PhysioNet 0.5, respectively. The root means squared error (RMSE) and mean absolute error (MAE) are used as the evaluation metrics.From the test results shown in Table 1 and Fig. 2 , we can see that our method CLWF outperforms the existing deep learning-based method (GP-VAE) and the recent state-of-the-art diffusion methods (CSDI, CSBI, and DSPD-GP). Moreover, CLWF uses only 15 sampling steps for inference, while the 5 Related Work 
diffusion models: 
 Diffusion models, such as DDPMs [20] and SBGM [43] , are considered as the new contenders to GANs on data generation tasks. But they generally take relatively long time to produce high quality samples. To mitigate this problem, the flowing matching methods have been proposed from an optimal transport. For example, ENOT uses the saddle point reformulation of the OT problem to develop a new diffusion model [19] The flowing matching methods have also been proposed based on the OT theory [27, 29, 31, 2, 1] . In particular, mini-batch couplings are proposed to straighten the probability flows for fast inference [39, 45, 46] .The Schrödinger Bridge have also been applied to diffusion models for improving the data generation performance of diffusion models. Diffusion Schrödinger Bridge utilizes the Iterative Proportional Fitting (IPF) method to solve the SB problem [14] . SB-FBSDE proposes to use forward-backward (FB) SDE theory to solve the SB problem through likelihood training [10] . GSBM formulates a generalized Schrödinger Bridge matching framework by including the task-specific state costs for various data generation tasks [28] NLSB chooses to model the potential function rather than the velocity function to solve the Lagrangian SB problem [24] . Action Matching [33, 34] leverages the principle of least action in Lagrangian mechanics to implicitly model the velocity function for trajectory inference. Another classes of diffusion models have also been proposed from an stochastic optimal control perspective by solving the HJB-PDEs [35, 50, 5, 28] .Many diffusion-based models have been recently proposed for time series imputation [26, 32] . For instance, CSDI [44] combines a conditional DDPM with a Transformer model to impute time series data. CSBI [12] adopts the FB-SDE theory to train the conditional Schrödinger bridge model to for probabilistic time series imputation. To model the dynamics of time series from irregular sampled data, DSPD-GP [7] uses a Gaussian process as the noise generator. TDdiff [23] utilizes self guidance and learned implicit probability density to improve the time series imputation performance of the diffusion models. However, the time series imputation methods mentioned above exhibit common issues, such as slow convergence, similar to many diffusion models. Therefore, in this work, we proposed CLWF to tackle thess challenges. 
conclusion, limitation, and broader impact: 
 In this work, we proposed CLWF, a novel time series imputation method based on the optimal transport theory and Lagrangian mechanics. To generate the missing time series data, following the principle of least action, CLWF learns a velocity field by minimizing the kinetic energy to move the initial random noise to the target distribution. Moreover, we can also estimate the derivative of a potential function via a VAE model trained on the observed training data to further improve the performance of the base sampler by Rao-Blackwellization. In contrast with previous diffusion-based models, the proposed requires less simulation steps and Monet Carlo samples to produce high-quality data, which leads to fast inference. For validation, CWLF is assessed on two public datasets and achieves competitive results compared with existing methods.One limitation of CLWF is that the samples obtained are not diverse enough as we use ODE for inference, which results in slightly higher test (continuous ranked probability score) CRPS compared to previous works, e.g., CSDI. Therefore, for future work, we will seek suitable approaches to accurately model the diffusion term in the SDE. Moreover, we will also try to design better taskspecific potential functions for sparse multivariate time series data. We plan to explore the potential of the Lagrangian Wasserstein Flow model for other time series analysis tasks, such as anomaly detection and uncertainty quantification.In terms of broader impact, our study on time series imputation has the potential to address important real-world challenges and consequently make a positive impact on daily lives.The data generation task can also be interpreted as a stochastic optimal control (SOC) problem [4, 17, 35, 50, 21, 5] whose cost function J is defined as:EQUATIONwhere 1 2 ∥∇ x Ψ(X t , t)∥ 2 denotes the running cost, and Ψ(X T ) denotes the terminal cost. The above SOC problem can be solved by dynamic programming [4, 6] .Further, let V (X t , t) = inf J (X t , t) be the value function/optimal-cost-to-go of the SOC problem, then the corresponding Hamilton-Jacobi-Bellman (HJB) partial differential equation (PDE) [15, 49] is given byEQUATIONwith the terminal condition: V (X t , T ) = Ψ(X t ).B Rao-Blackwell Theorem Theorem 1 (Rao-Blackwell) Let S be an unbiased estimator of some parameter θ ∈ Θ, and T (X)the sufficient statistic for θ, then: 1) S * = E[S|T (X)],is an unbiased estimator for θ, and 2)V θ [S * ] ≤ V θ [S]for all θ. The inequality is strict unless S is a function of T . 
 
 Appendix: For the hardware environment of the experiments, we use a single NVIDIA A100-PCIE-40GB GPU and an Intel(R) Xeon(R) Gold-6248R-3.00GHz CPU. For the software environment, the Python version is 3.9.7, the CUDA version 11.7, and the Pytorch version is 2.0.1. 


Decision: This paper deals with the problem of time series imputation. It introduces a method that uses tools from transport theory, Lagrangian mechanics, and optimal transport to generate samples.

This paper received mixed reviews (1x accept, 2x weak accept, 1x borderline accept, 1x strong reject). After the authors' rebuttal, some reviewers stated that their concerns had been addressed and maintained their score. Reviewer 3hN3, however, points to some remaining flaws of the paper. After the rebuttal period, only Reviewer 3hN3 participated in the follow-up discussion between reviewers and AC, and thus a consensus among the reviewers was not reached.

Because of the mixed opinions, I had a careful look at the paper myself. Unfortunately, I found it does not meet the acceptance bar and therefore I encourage the authors to improve their manuscript using the received feedback and submit to another conference.

Besides the issues found by Reviewer 3hN3, I noted that the paper contains basic mistakes in the mathematical equations. As an example, consider Eq. (7) (in the "Background" section), which is:
$$\mathcal{K}(x_t, \mu_t, t) = \mathbb{E}_{p(X_t)} \Bigg[ \int_0^T \int  \frac{1}{2} ||\mu_t(x_t, t)||^2 dx dt$$
(where the inner integral is over $\mathbb{R}^d$).
+ First, note that the square bracket isn't closed - I assume that's a simple typo and that the closing bracket should be after $dt$.
+ Second, note that the inner integral has $dx$ but there is no $x$ in the integrand - I assume that it should be $dx_t$ instead.
+ Third, if we integrate out $x_t$, then the expectation $\mathbb{E}_{p(X_t)}[\ldots]$ would not depend on $x_t$ at all, and therefore it can be safely ignored as everything inside is a constant - which should obviously not be the case.
+ Fourth, after we integrate out both $x_t$ and $t$, the result does not depend on either $x_t$ or $t$, whereas the LHS of the equation clearly depends on both - it is written as $\mathcal{K}(x_t, \mu_t, t)$.

As another example of wrong mathematics, Eq. (8) suffers from similar issues.

Besides these mistakes, I would like to note that the paper contains some unsupported claims. Two examples are:
+ In Section 3.5 ("Rao-Blackwellized sampler"), the text says: *"Note that the ODE sampler alone is good enough to generate high-quality samples for time series imputation"*. This leaves it unclear why switching from SDE to ODE is okay, as it's never discussed further.
+ In Section 3.4 ("Potential function"), the text says: *"In terms of practical implementation, we parameterize $\nabla_x U(X_t)$ via a Variational Autoencoder (VAE)"*. Once again, it remains unclear why that is an okay thing to do, since the (time-independent) VAE is fitted to the observed data only, and it might suffer from some biases that would propagate further. Reviewer 3hN3 also raised this point: *"More fundamentally I'm unclear why effectively taking steps towards data density of a VAE trained on a related (but fundamentally different) data distribution is even a good idea [...]"*.

Another point, raised by Reviewer 3hN3 (and which I agree with), is the lack of justification of the Rao-Blackwellized sampler. Section 3.5 does not clarify why this is a valid sampler, since the expectation $\mathbb{E}[\mathcal{S} | \mathcal{T}]$ is not correct. Unfortunately, despite this being one of the main contributions of the paper, this concern remained unaddressed after the rebuttal phase.

Even if the Rao-Blackwellized sampler were not theoretically valid, it could still be empirically useful; however, the experimental results do not support that conclusion (there is only a single result on a single dataset with marginally better results). It is also suspicious that, despite considering two datasets, the Rao-Blackwellized sampler is consctructed for one of them only, ignoring the second one.

To conclude, I think the paper would need a substantial changes before it is ready for acceptance at a conference like NeurIPS.

Reviews: review: paper_summary: paper_summary: The paper presents Conditional Lagrangian Wasserstein Flow (CLWF), a new method for time series imputation. Using (entropic) optimal transport theory and Lagrangian mechanics, CLWF generates high-quality samples. Enhanced with a Rao-Blackwellized sampler, CLWF incorporates prior information through a variational autoencoder. Experiments on real-world datasets show that CLWF performs competitively.
strength_weakness: strength_weakness: 1. The conditional generation process is trained through the simulation-free training of the Schrödinger Bridge, which limits the novelty of the CLWF.
2. The paper mentions that the samples obtained using the ODE-based inference method may lack diversity, potentially leading to higher continuous ranked probability scores (CRPS) compared to some previous works.
3. While the method is tested on two real-world datasets, broader evaluation across more diverse and challenging datasets could strengthen the validation of the approach.
questions: questions: I have no other questions.
limitations: limitations: Yes
paper_summary: paper_summary: This work trains a conditional flow model from noise to time series data. A VAE is used to estimate the data density then perform interleaved flow and density gradient ascent steps to generate new time series. This is referred to as a Rao-Blackwellization procedure. It is shown empirically that the model performs well on time series imputation tasks across two datasets, and that the VAE-based guidance helps in one case.
strength_weakness: strength_weakness: - Quality: Experiments are limited in scope and the it is not clear whether the “Rao-Blackwellization” (the main methodological novelty) step reliably improves performance.
    - In order to claim “Less sampling steps” it would be great to understand how performance changes with the number of steps for both diffusion and flow-based models.
    - More than 2 time series datasets would allow an understanding of when this “Rao-Blackwellized” sampler is valid.
    - No error bars (even though this is claimed in the checklist)
    - The empirical support is lacking for the effectiveness of the “Rao-Blackwellization” step. Currently it is shown to slightly benefit (although without error bars this is difficult to tell how much) on a single dataset. In this case, an improvement in RMSE from 18.2 to 18.1 on a single dataset seems insignificant without additional detail. While it is claimed that “the PhysioNet dataset does not have enough non-zero datapoints to train a valid VAE model”, a Rao-Blackwellized sampler should always work regardless of the performance of the estimate? It would be good to see if it helps at all on this dataset and others.
- Clarity: It is not clear to me that this even is a “Rao-Blackwellized” sampler. 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. It would be great if the authors could prove that Algorithm 2 is really a Rao-Blackwellization (at least under some conditions).
- Significance: Without further experimental benchmarking especially with regards to the novel sampler, the applicability of this method is very limited.
questions: questions: How is this a Rao-Blackwellization, and what does the paper 

Minor comments:

Shrodinger —> Schrodinger throughout

paragraph line 146 is confusing. this equation is strange given that p(x^tar) and p(x^input) are dependent if taking batches.
limitations: limitations: I'm a bit confused about the limitations. It is claimed that 

> One limitation of CLWF is that the samples obtained are not diverse enough as we use ODE for
> inference, which results in slightly higher test (continuous ranked probability score) CRPS compared
> to previous works, e.g., CSDI. Therefore, for future work, we will seek suitable approaches to
> accurately model the diffusion term in the SDE.

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.
paper_summary: paper_summary: Inspired by optimal transport and Lagrangian dynamics, this work proposes to use the conditional Lagrangian Wasserstein flow to impute time series data. The method requires less model evaluation steps to generate high quality samples compared to existing diffusion models. Moreover, a task-specific energy function is used to further improve the model’s performance.
strength_weakness: strength_weakness: - How should the values of the hyperparameters (e.g., the variance of the potential function) be chosen? How will these choice affect the model’s performance?
- The reason for choosing the VAE model to construct the potential function is unclear. Are there other functions can also be used as the task-specific potential function?
- The authors proposed to use a VAE model to enhance the data sampling procedure. However, if the VAE model performs poorly, can it still help with the data generation process?
questions: questions: Can the proposed method be used to compute the likelihood of the generated data samples?
limitations: limitations: The deterministic sampling process may introduce bias, potentially failing to accurately reflect the target marginal distribution.
paper_summary: paper_summary: In this paper, the authors introduce a new time series imputation model based on the conditional Lagrangian Wasserstein flow. Different from previous diffusion-based models, the proposed model leverages the optimal transport theory and Lagrangian dynamics to improve the data generation performance.
strength_weakness: strength_weakness: 1. In Sec 3.2, the technique used for projecting the interpolants into the Wasserstein space is unclear. The authors should elaborate more on this.
2. Compared to CSDI and other diffusion-based imputation methods, the proposed method’s data generation process seems deterministic. Hence, it cannot be used to quantify the uncertainty of the prediction.
3. The proposed sampler is implemented through an ODE, which is inconsistent with the Schrodinger Bridge problem in sec 2.3.
4. There are some inconsistencies in the notations, please double-check.
questions: questions: What is the difference between the data generation processes in Euclidean space and Wasserstein space?
limitations: limitations: The ablation study on the effectiveness of Rao-Blackwellization should be conducted on more than one dataset (PM 2.5).
paper_summary: paper_summary: The authors introduce the conditional Lagrangian Wasserstein flow method for time series imputation.
The time series imputation task is treated as a conditional data generation problem. The authors use flow matching to learn an ODE sampler for generating the missing time series data. They further propose to enhance the imputation performance via Rao-Blackwellization. The method is tested on the real-world datasets.
strength_weakness: strength_weakness: W1. The imputation process is implemented using an ODE sampler. Why not use an SDE sampler?

W2. Why use the Euler method instead of other higher-order solvers, such as the Runge-Kutta methods, for solving the ODE?

W3. Can the proposed method also be applied to time series forecasting tasks?

W4. The dynamic described in Eq (15) is deterministic, which is inconsistent with Eq (1). Why the diffusion term is missing in Eq (15)?
questions: questions: Please see my weaknesses.

And there is another question,
Q. Can the VAE-based potential function be used alone for time series imputation?
limitations: limitations: CRPS is not used as the evaluation metric in the experiments.

Comments: title: Thanks for rebuttal
reply
reply to rebuttal
continuation of rebuttal
continuation of rebuttal
comment: Thanks for your rebuttal. The authors have addressed all my concerns and I will maintain my score at this stage.
Thanks for the quick reply.
    
      How is this implemented practically?

At each step we first update the sample via $x_{t+1} = x_{t} +  v_t^\phi(x_t, t) \frac{T}{N}$, then obtain the new sample using the previous sample with $x_{t+1} = x_{t} +   \mu_t^\theta(x_t, t) \frac{T}{N}$ (you can take this as a correction step).


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

It still does, because in this case the diversity of the samples obtained only relies on the initial noise (which also appears in the training samples), which certainly causes that 1) it has lost some level of diversity during the sampling process compared SDE, 2) and the overfitting issue is more serious (especially when the variance of  $p_0$ is small).


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

It suggests that the potential function does not need to be time-dependent function (as in our case). And that the derivative of the potential function can represent the sampling dynamics, thus it is a sufficient statistic for $X_{t+1}$.


    There are no additional experiments on the "Rao-Blackwellization" that I can see.
Please see the new results from the first Table (in the previous comment and the uploaded PDF) with CLWF(no RB) and CLWF(with RB).

## Previous comments
    $\nabla_x U(X_t) = \nabla \log N(t, 1) = t - x_t$.
It should be $\nabla_x U(X_t) = \nabla \log N(x_t|1) = 1 - x_t$.

    any non-zero quantity will bias our estimate $X_{t+1} | X_t = \mu_t^\star(x) = 1.$ 
in this case the velocity is optimal, then the derivative of the  potential function will be zero, which will not affect the sampling result.
Thank you for the quick response.

> The proposed sampling step is not $\mathbb{E}[\mathcal{S} | \mathcal{T}] = \mu_t^\theta(x^{input}, t) \frac{T}{N} + v_t^\phi(x_t, t) \frac{T}{N}$ 
(which the reviewer framed), but, $\mathbb{E}[\mathcal{S} | \mathcal{T}] = \mu_t^\theta(x_t + v_t^\phi(x_t, t) \frac{T}{N} , t) \frac{T}{N}$, which is fundamentally different

Okay great. I agree this is fundamentally different. This is closer, but still needs to be proven in my opinion. I also don't understand how this lines up with Algorithm 2 which clearly states $\mathbb{E}[\mathcal{S} | \mathcal{T}] = \mu_t^\theta(x^{input}, t) \frac{T}{N} + v_t^\phi(x_t, t) \frac{T}{N}$. How is this implemented practically?

> Why the optimal velocity is $\mu_t^\star(x) = 1$...

This is the optimal velocity for a flow model which is integrated from time zero to time one.

Because $x_1 = x_0 + \int_{t=0}^1 \mu_t^\star(x) dt =  x_0 + \int_{t=0}^1 1 dt = x_0 + 1$

> An ODE is always deterministic, which certainly limits its diversity to some extent.

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

> Please refer to the concepts of gradient flow and Langevin dynamics.

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

> We added additional experiments, please check the results carefully.

I have. There are no additional experiments on the "Rao-Blackwellization" that I can see.
## Q: Why is $\mathbb{E}[\mathcal{S} | \mathcal{T}] = \mu_t^\theta(x^{input}, t) \frac{T}{N} + v_t^\phi(x_t, t) \frac{T}{N}$? 
### A:    The proposed sampling step is not  $\mathbb{E}[\mathcal{S} | \mathcal{T}] = \mu_t^\theta(x^{input}, t) \frac{T}{N} + v_t^\phi(x_t, t) \frac{T}{N}$ (which the reviewer framed), but $\mathbb{E}[\mathcal{S} | \mathcal{T}] = \mu_t^\theta(x_t + v_t^\phi(x_t, t) \frac{T}{N} , t) \frac{T}{N} $, which is fundamentally different. 

## Q: the optimal $\mu_t^\theta(x, t)$ which we will call $\mu_t^\star(x) = 1$
### A: Why the optimal velocity is $\mu_t^\star(x) = 1$ in your case ($p_0 = N(0,1)$ and $p_1 = N(1, 1)$)? You certainly cannot generate the target distribution via $x_{t+1} = x_t + h$ (not converging, more iterations always result in larger values). 



## Q:  every SDE there is an equivalent ODE
### A: An ODE is always deterministic, which certainly limits its diversity to some extent. 


## Q: It will be over concentrated towards the mean at every step.''' 
### A: That is why this is an ODE sampler.

## Q: It seems impossible as the VAE does not depend on for it to be a sufficient statistic for which is time dependent.
### A: Please refer to the concepts of gradient flow and Langevin dynamics.


## Q: The experimental results are still extremely limited. The fact that this “Rao-Blackwellization” can only be tuned to work on one of the two datasets is not convincing. Also why are the provided results different than those in the paper?"

### A: We added additional experiments, please check the results carefully.

## Q: Also why are the provided results different than those in the paper?
### A: Because the learning process and sampling process are random like any other stochastic learning methods.
I thank the authors for their detailed response, and additional experiments. Particularly when a small number of steps are used. 

> Justification of Rao-Blackwellization

After the authors' response, I am fairly convinced this **not** a valid Rao-Blackwellization. The provided is not a valid proof that this is a valid Rao-Blackwellization, only that Rao-Blackwell samplers in general are valid. 

More specifically I believe at least two things are necessary to show that this is a valid Rao-Blackwellization:

Why is $\mathbb{E}[\mathcal{S} | \mathcal{T}] = \mu_t^\theta(x^{input}, t) \frac{T}{N} + v_t^\phi(x_t, t) \frac{T}{N}$? This does not seem possible to me. Lets take the simple example of an OT-FM model with $p_0 = N(0,1)$ and $p_1 = N(1, 1)$ then the optimal $\mu_t^\theta(x, t)$ which we will call $\mu_t^\star(x) = 1$. $\nabla_x U(X_t) = \nabla \log N(t, 1) = t - x_t$. In fact any non-zero quantity will bias our estimate of $X_{t+1} | X_t = \mu_t^\star(x) = 1.$  Furthermore, following the suggested “Rao-Blackwellized” estimator, will not lead to $p_1$. It will be over concentrated towards the mean at every step. 

Why is $v_t^\psi(x_t, t) \frac{T}{N}$ a sufficient statistic? It seems impossible as the VAE does not depend on $t$ for it to be a sufficient statistic for $X_{t+1} | X_t$ which is time dependent.

> Other ODE models also have the problem with diversity to some extent.
> 

This is simply not true. For every SDE there is an equivalent ODE with equivalent diversity through the probability flow ODE for any diversity metric measured on $p_T$.

The experimental results are still extremely limited. The fact that this “Rao-Blackwellization” can only be tuned to work on one of the two datasets is not convincing. Also why are the provided results different than those in the paper? 

> Note that the PhysioNet dataset does not have enough non-zero data points to train a valid VAE model, therefore we only construct the Rao-Blackwellized sampler for the PM 2.5 dataset.

If this is a true Rao-Blackwellization, then the estimate should be **no worse** than the simple estimator. As the authors state, it may not help but it should not hurt. This is further evidence to me that this theoretical result is invalid. 

My feeling is that that CLWF without (”Rao-Blackwellization”), which is equivalent to OT-CFM (Tong et al. 2023) / Minibatch-OT FM (Pooladian et al. 2023), should work quite well in this setting. However, the “Rao-Blackwellization” is not proven anywhere theoretically and the response does not justify it sufficiently. Furthermore the experimental results are only applied on a single dataset for a method that should be no worse than the baseline. I keep my current score of 2. However, am open to substantially raising it if the authors can show that this is a valid Rao-Blackwellization.
Since all of my questions have been adequately answered by the authors, I recommend acceptance.
The authors' reply addresses most of my issues. I appreciate the clarification made by the authors, and I have no other concerns.
## Test imputation results on synthetic data ($5$ trials, values are multiplied by $10^2$.


| Method  ||   Synthetic 0.4   ||  Synthetic 0.6  ||  Synthetic 0.8    |
| ----------- |----------- | ----------- | ----------- | -----------   | -----------  | -----------  | 
|   | RMSE |  MAE  |RMSE |  MAE  |RMSE |  MAE  |
| CLWF(no RB) |$22.91\pm0.49$ |$15.28\pm0.21$|$25.65\pm0.31$ |$15.54\pm0.22$|$27.41\pm0.27$  |$15.91\pm0.23$ |
|CLWF(with RB)  | $22.72\pm0.48$| $13.23\pm0.42$|$25.44\pm0.30$|$15.28\pm0.17$| $27.32\pm0.27$| $15.79\pm0.23$ |




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


| Method  ||   5      steps   ||   10 steps   ||  15 steps    ||  20 steps     |
| ----------- |----------- | ----------- | ----------- | -----------   | -----------  | -----------  | -----------  | -----------   | 
|   | RMSE |  MAE  |RMSE |  MAE  |RMSE |  MAE  |RMSE |  MAE  |
| CSDI      |$34.21\pm0.16$  |$14.85\pm0.01$  |$29.43\pm0.46$   |$12.48\pm0.08$  |$22.40\pm0.16$ | $10.78\pm0.04$  |$19.22\pm0.13$  | $9.91\pm0.02$ |
| CLWF      |$18.29\pm0.002$ |$9.78\pm0.004$  |$18.28\pm0.003$  |$9.77\pm0.005$  |$18.26\pm0.006$|$9.76\pm0.004$ |$18.21\pm0.002$ |$9.72\pm0.004$ |





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


| Method  ||   5 steps   ||   10 steps   ||  15 steps    ||  20 steps     |
| ----------- |----------- | ----------- | ----------- | -----------   | -----------  | -----------  | -----------  | -----------   | 
|   | RMSE |  MAE  |RMSE |  MAE  |RMSE |  MAE  |RMSE |  MAE  |
| CSDI      |$0.60\pm0.00$|$0.22\pm 0.000$|$0.58\pm0.005$|$0.22\pm0.000$|$0.57\pm0.002$ |$0.22\pm0.000$ |$0.56\pm 0.002$|$0.22\pm0.001$|
| CLWF      |$0.48\pm0.000$|$0.22\pm0.00$|$0.47\pm0.000$|$0.22\pm0.00$|$0.47\pm0.000$|$0.22\pm0.00$  |$0.48\pm0.000$|$0.22\pm0.000$|





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


| Method  ||   5 steps   ||   10 steps   ||  15 steps    ||  20 steps     |
| ----------- |----------- | ----------- | ----------- | -----------   | -----------  | -----------  | -----------  | -----------   | 
|   | RMSE |  MAE  |RMSE |  MAE  |RMSE |  MAE  |RMSE |  MAE  |
| CSDI      |$0.71\pm0.000$  |$0.31\pm0.000$ |$0.69\pm0.001$  |$0.31\pm0.000$ |$0.68\pm0.001$ |$0.31\pm0.000$ |$0.68\pm0.000$  |$0.30\pm0.000$ |
| CLWF      |$0.64\pm0.000$|$0.29\pm0.000$|$0.64\pm0.000$|$0.29\pm0.000$|$0.64\pm0.000$|$0.29\pm0.000$|$0.64\pm0.000$|$0.29\pm0.000$ |




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


| Method  ||   PM 2.5   || PhysioNet 0.1 ||  PhysioNet 0.5    |
| ----------- |----------- | ----------- | ----------- | -----------   | -----------  | -----------  | 
|   | RMSE |  MAE  |RMSE |  MAE  |RMSE |  MAE  |
| CLWF(no RB) |$18.27\pm 0.01$ |$9.76\pm0.01$  |$0.47\pm 0.000$  |$0.22\pm 0.000$  |$ 0.64\pm0.002$  |$0.29\pm0.009$|
|CLWF(with RB) | $18.08\pm 0.02$ |$9.71\pm 0.00$|
## 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
\begin{align}
X_T = X_0 + \int_0^T \mu (X_t,t)dt,    
\end{align}
and an SDE sampler described by
\begin{align}
X_T = X_0 + \int_0^T \mu (X_t,t)dt  +  \int_0^T \sigma (X_t,t)dW_t. 
\end{align}

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