Title: Conditional Lagrangian Wasserstein Flow for Time Series Imputation
Venue: NeurIPS 2024
Accepted: no
Keywords: dep learning; diffusion model; optimal transport; time series imputation
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.

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.

1 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

space via sampling mini-batch OT maps. Afterwards, we construct the time-dependent intermediate 37 samples through interpolating the source distribution and target distribution. Then according to 38 the principle of least action in Lagrangian mechanics [3], the optimal velocity function moving the 39 source distribution to the target distribution is learned in a self-supervised manner by minimizing 40 the corresponding kinetic energy. Moreover, to further improve the model's performance, we learn 41 the task-specific potential function by training a Variational Autoencoder (V AE) model [22] on the 42 observed time series data to build a Rao-Blackwellized trajectory sampler. 43

44

45

46

Finally, CLWF is assessed on two real-word multivariate time series datasets. The obtained results show that the proposed method achieves competitive performance and admits fast convergence compared with other state-of-the-art time series imputation methods.

47

48

49

50

51

52

53

54

55

56

57

58

59

60

61

62

63

64

65

66

67

68

69

70

71

72

73

74

75

76

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.

2 Preliminaries

In this section, we concisely introduce the fundamentals of stochastic differential equations, optimal transport, Shrödinger Bridge, and Lagrangian mechanics.

2.1 Stochastic Differential Equations

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]:

where 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:

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

2.2 Optimal Transport

The optimal transport (OT) problem aims to seek the optimal transport plans/ maps that moves the 77 source distribution to the target distribution [47, 41, 38]. In the Kantorovich's formulation of the 78 OT problem, the transport costs are minimized with respect to some probabilistic couplings/joint 79 distributions [47, 41, 38]. Let p 0 and p T be two Borel probability measures with finite second 80

moments on the space Ω ∈ R d . Π( p 0 , p T ) denotes a set of transport plans between these two 81 marginals. Then, the Kantorovich's OT problem is defined as follows: 82

where Π( 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].

2.3 Shrödinger Bridge

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:

where the Kullback-Leibler (KL) divergence D KL ( P ∥ Q ) = { log d P d Q d P , 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]:

where θ is the parameters of the variational drift function µ t .

2.4 Lagrangian Mechanics

In this section, we formulate the data generation problem under the framework of Lagrangian mechanics [3]. Let p t and ˙ p t = d p t d t be the density and law of the generalized coordinates X t , respectively. K ( p t , ˙ p t , t ) is the kinetic energy, and U ( p t , t ) is the potential energy, then the corresponding Lagrangian is

We assume that Eq. (6) is lower semi-continuous (lsc) and strictly convex in ˙ p 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:

where U t ( X t ) is the potential function. Then the action in the context of Lagrangian mechanics is 110 defined as follow: 111

83

84

85

86

87

88

89

90

91

92

93

94

95

96

97

98

99

100

101

102

103

104

105

106

107

108

109

According to the principle of least action , the shortest path is the one minimizing the action, which is 112 aligned with Eq. (4) in the SB theory as well. Therefore, we can leverage the Lagrangian dynamics 113 to tackle the OT problem for data generation. To solve Eq. (6), we need to satisfy the stationary 114 condition, i.e., the Euler-Lagrangian equation: 115

with the boundary condition d X t d t = µ ( X t , t ) , q 0 = p 0 , q T = p T .

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

3.1 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 data x 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 is x input 0 := ( x cond , x 0 ) ∈ R K × L × 2 .

3.2 Interpolation in Wasserstein Space

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:

where Ω ∈ 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:

where γ 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 variable x input t := ( x cond , x t ) as the input for computing the velocity of the Wasserstein flow.

3.3 Flow Matching

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

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

[9, 37]. To circumvent the above issues, in this work we adopt a simulation-free training strategy 153 based on the OT theory introduce in Sec. 2.2 [30, 46, 2], which turns out to be faster and more 154 scalable to large time series datasets. 155

Since we can now draw mini-batch interpolated samples of the source distribution and target distribu156 tion in the Wasserstein space using Eq. (12), we can model the variational velocity function using a 157 neural network with parameters θ . Then, according to Eq. (1), the target velocity can be computed 158 by the difference between the source distribution and target distribution. Therefore, the variational 159 velocity function µ θ ( x input t , t ) can be learned by 160

Eq. (14) can be solved by drawing mini-batch samples in the Wasserstein space and performing 161 stochastic gradient descent accordingly. In this fashion, the learning process is simulation-free as the 162 trajectory simulation is not needed. 163

164

165

166

167

168

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.

169

170

171

172

173

174

175

3.4 Potential Function

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 by

Since the data generation problem in our case can also be interpreted as a stochastic optimal control 176 (SOC) problem [4, 17, 35, 50, 21, 5], then the existence of such U t ( X t ) is assured by Pontryagin's 177 Maximum Principle (PMP) [16]. 178

To estimate v t ( X t , t ) , according to the Lagrangian Eq. (6), we assume that the potential function 179 takes the form U t ( X t ) ≈ -log N ( X t | ˆ X t , σ 2 p ) , where ˆ X t the learned mean and σ 2 p is the pre-defined 180 variance. As a result, the derivative is ∇ x U ( X t ) = X t -ˆ X t σ 2 p . In terms of practical implementation, we 181 parameterize ∇ x U ( X t ) via a Variational Autoencoder (VAE) [22]. More specifically, we pre-train a 182 VAE on the total observed time series data X obs . Afterwards, the reconstruction discrepancies of the 183 VAE are used to approximate the task-specific v ϕ ( X t , t ) depending on X t : 184

where VAE( X t ) represents the reconstruction output of the pre-trained V AE 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.

3.5 Rao-Blackwellized Sampler

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

185

186

187

188

189

190

191

Observed data

x

obs

Target data

x

obs

Initial noise

Intermediate samples

x

t

Joint input

x

tar

Conditional data

Target velocity

d

Xt/

d

t

Flow

matching

loss

VAE model

Flow model

x

cond

x

0

x

cond

x

t

Potential

function

means the diffusion term in Eq. 1 is omitted). Alternatively, one can also adopt the SDE sampler by 192 using the Euler-Maruyama method. Nevertheless, to ensure achieve the best imputation performance, 193 we choose the ODE sampler for implementation. Note that the ODE sampler alone is good enough to 194 generate high-quality samples for time series imputation. 195

196

197

198

199

200

201

202

203

204

205

206

207

208

209

210

211

212

213

214

215

216

217

218

219

220

221

222

223

224

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 have

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.

3.6 The Algorithms

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 d X t d t by computing x tar -x 0 T . 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 V AE 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 information x 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 RaoBlackwellization, then x t is fed to the V AE 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

225

226

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.

Algorithm 1 Training procedure

228

229

230

231

232

233

234

235

236

237

238

239

240

241

242

243

244

245

246

247

248

249

250

251

252

253

254

255

256

257

258

259

260

4 Experiments

4.1 Datasets

We use two public multivariate time series datasets for validation. The first dataset is the PM 2.5 dataset [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 48 hours. 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.

4.2 Baselines

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.

4.3 Experimental Settings

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 .

4.4 Experimental Results

4.4.1 Imputation Results

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

Algorithm 2 Sampling procedure

0

5

10

15

20

25

30

35

Time

10

20

30

40

50

60

70

Value

0

5

10

15

20

25

30

35

Time

50

100

150

200

250

300

350

400

Value

0

5

10

15

20

25

30

35

Time

5

10

15

20

25

Value

baseline diffusion method uses only 50 sampling steps. This suggests that CLWF is faster and more 261 accurate than the existing methods on time series imputation tasks. 262

263

264

265

266

267

268

4.4.2 Ablation Study

Single-sample Imputation Result. We compare the time series imputation performance of CLWF with CSDI using only one Monte Carlo sample. The test results shown in Table 2 shows that CWFL outperforms CSDI, which suggests that CWFL exhibits lower imputation variances compared to diffusion-based models. This indicates that CWFL is more efficient and computationally economical for inference.

Effect of Rao-Blackwellzation. We compare the test imputation CLWF wth and without using RaoBlackwellzation. Note that the PhysioNet dataset does not have enough non-zero data points to train a valid V AE model, therefore we only construct the Rao-Blackwellized sampler for the PM 2.5 dataset. The results showed in Table 3 indicates

that the Rao-Blackwellized sampler can further improve the time series imputation performance of the base sampler.

269

270

271

272

273

274

275

276

277

278

279

280

281

282

283

284

285

286

287

288

289

290

291

292

293

294

295

296

297

298

299

300

301

302

303

304

305

306

307

308

309

310

311

312

313

314

315

316

5 Related Work

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

5.2 Time Series Imputation

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.

6 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 V AE 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.

317

318

319

320

321

322

323

324

References

Michael S Albergo, Nicholas M Boffi, and Eric Vanden-Eijnden. Stochastic interpolants: A unifying framework for flows and diffusions. arXiv preprint arXiv:2303.08797 , 2023.

Michael Samuel Albergo and Eric Vanden-Eijnden. Building normalizing flows with stochastic interpolants. In The Eleventh International Conference on Learning Representations , 2023.

Vladimir Igorevich Arnol'd. Mathematical methods of classical mechanics , volume 60. Springer Science & Business Media, 2013.

Richard Bellman. Dynamic programming. science , 153(3731):34-37, 1966.

Julius Berner, Lorenz Richter, and Karen Ullrich. An optimal control perspective on diffusion325 based generative modeling. Transactions on Machine Learning Research , 2024. 326

327

328

Dimitri Bertsekas. Dynamic programming and optimal control: Volume I , volume 4. Athena scientific, 2012.

329

330

331

332

333

334

335

336

337

338

339

340

341

342

343

344

345

346

347

348

349

350

351

352

353

354

355

356

357

358

359

360

361

Marin Biloš, 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 , pages 2452-2470. PMLR, 2023.

George Casella and Christian P Robert. Rao-blackwellisation of sampling schemes. Biometrika , 83(1):81-94, 1996.

Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems , 31, 2018.

Tianrong Chen, Guan-Horng Liu, and Evangelos Theodorou. Likelihood training of schrödinger bridge using forward-backward sdes theory. In International Conference on Learning Representations, 2022 , 2022.

Yongxin Chen, Tryphon T Georgiou, and Michele Pavon. Stochastic control liaisons: Richard sinkhorn meets gaspard monge on a schrodinger bridge. Siam Review , 63(2):249-313, 2021.

Yu Chen, Wei Deng, Shikai Fang, Fengpei Li, Nicole Tianjiao Yang, Yikai Zhang, Kashif Rasul, Shandian Zhe, Anderson Schneider, and Yuriy Nevmyvaka. Provably convergent schrödinger bridge with applications to probabilistic time series imputation. In International Conference on Machine Learning , pages 4485-4513. PMLR, 2023.

Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in neural information processing systems , 26, 2013.

Valentin De Bortoli, James Thornton, Jeremy Heng, and Arnaud Doucet. Diffusion schrödinger bridge with applications to score-based generative modeling. Advances in Neural Information Processing Systems , 34:17695-17709, 2021.

Lawrence C Evans. Partial differential equations , volume 19. American Mathematical Society, 2022.

Lawrence C Evans. An introduction to mathematical optimal control theory spring, 2024 version. Lecture notes available at https://math.berkeley.edu/ evans/control.course.pdf , 2024.

Wendell H Fleming and Raymond W Rishel. Deterministic and stochastic optimal control , volume 1. Springer Science & Business Media, 2012.

Vincent Fortuin, Dmitry Baranchuk, Gunnar Rätsch, and Stephan Mandt. Gp-vae: Deep probabilistic time series imputation. In International conference on artificial intelligence and statistics , pages 1651-1661. PMLR, 2020.

Nikita Gushchin, Alexander Kolesov, Alexander Korotin, Dmitry P Vetrov, and Evgeny Burnaev. Entropic neural optimal transport via diffusion processes. Advances in Neural Information Processing Systems , 36, 2024.

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances 362 in neural information processing systems , 33:6840-6851, 2020. 363

364

365

366

Lars Holdijk, Yuanqi Du, Ferry Hooft, Priyank Jaini, Berend Ensing, and Max Welling. Stochastic optimal control for collective variable free sampling of molecular transition paths. Advances in Neural Information Processing Systems , 36, 2023.

367

368

369

370

371

372

373

374

375

376

377

378

379

380

381

382

383

384

385

386

387

388

389

390

391

392

393

394

395

396

397

398

399

400

401

402

403

404

405

406

Diederik P. Kingma and Max Welling. Auto-encoding variational bayes. In The Second International Conference on Learning Representations , 2014.

Marcel Kollovieh, Abdul Fatir Ansari, Michael Bohlke-Schneider, Jasper Zschiegner, Hao Wang, and Yuyang Bernie Wang. Predict, refine, synthesize: Self-guiding diffusion models for probabilistic time series forecasting. Advances in Neural Information Processing Systems , 36, 2024.

Takeshi Koshizuka and Issei Sato. Neural lagrangian schrödinger bridge: Diffusion modeling for population dynamics. In The Eleventh International Conference on Learning Representations , 2022.

Christian Léonard. From the schrödinger problem to the monge-kantorovich problem. Journal of Functional Analysis , 262(4):1879-1920, 2012.

Lequan Lin, Zhengkun Li, Ruikun Li, Xuliang Li, and Junbin Gao. Diffusion models for timeseries applications: a survey. Frontiers of Information Technology & Electronic Engineering , pages 1-23, 2023.

Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling. In The Eleventh International Conference on Learning Representations , 2022.

Guan-Horng Liu, Yaron Lipman, Maximilian Nickel, Brian Karrer, Evangelos Theodorou, and Ricky TQ Chen. Generalized schrödinger bridge matching. In The Twelfth International Conference on Learning Representations , 2024.

Qiang Liu. Rectified flow: A marginal preserving approach to optimal transport. arXiv preprint arXiv:2209.14577 , 2022.

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.

Xingchao Liu, Lemeng Wu, Mao Ye, and Qiang Liu. Learning diffusion bridges on constrained domains. In international conference on learning representations (ICLR) , 2023.

Caspar Meijer and Lydia Y Chen. The rise of diffusion models in time-series forecasting. arXiv preprint arXiv:2401.03006 , 2024.

Kirill Neklyudov, Rob Brekelmans, Daniel Severo, and Alireza Makhzani. Action matching: Learning stochastic dynamics from samples. In International Conference on Machine Learning , pages 25858-25889. PMLR, 2023.

Kirill Neklyudov, Rob Brekelmans, Alexander Tong, Lazar Atanackovic, Qiang Liu, and Alireza Makhzani. A computational framework for solving wasserstein lagrangian flows. arXiv preprint arXiv:2310.10649 , 2023.

Nikolas Nüsken and Lorenz Richter. Solving high-dimensional hamilton-jacobi-bellman pdes using neural networks: perspectives from the theory of controlled diffusions and measures on path space. Partial differential equations and applications , 2(4):48, 2021.

Bernt Oksendal. Stochastic differential equations: an introduction with applications . Springer Science & Business Media, 2013.

Derek Onken, Samy Wu Fung, Xingjian Li, and Lars Ruthotto. Ot-flow: Fast and accurate 407 continuous normalizing flows via optimal transport. Proceedings of the AAAI Conference on 408 Artificial Intelligence , 35(10):9223-9232, 2021. 409

410

411

412

413

414

415

416

417

418

419

420

421

422

423

424

425

426

427

428

429

430

431

432

433

434

435

436

437

438

439

440

441

442

443

444

445

446

447

Gabriel Peyré, Marco Cuturi, et al. Computational optimal transport: With applications to data science. Foundations and Trends® in Machine Learning , 11(5-6):355-607, 2019.

Aram-Alexandre Pooladian, Heli Ben-Hamu, Carles Domingo-Enrich, Brandon Amos, Yaron Lipman, and Ricky TQ Chen. Multisample flow matching: Straightening flows with minibatch couplings. In International Conference on Machine Learning , pages 28100-28127. PMLR, 2023.

Hannes Risken and Till Frank. The Fokker-Planck Equation: Methods of Solution and Applications , volume 18. Springer Science & Business Media, 2012.

Filippo Santambrogio. Optimal transport for applied mathematicians. Birkäuser, NY , 55(5863):94, 2015.

Ikaro Silva, George Moody, Daniel J Scott, Leo A Celi, and Roger G Mark. Predicting inhospital mortality of icu patients: The physionet/computing in cardiology challenge 2012. In 2012 Computing in Cardiology , pages 245-248. IEEE, 2012.

Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations , 2020.

Yusuke Tashiro, Jiaming Song, Yang Song, and Stefano Ermon. Csdi: Conditional score-based diffusion models for probabilistic time series imputation. Advances in Neural Information Processing Systems , 34:24804-24816, 2021.

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

Alexander Tong, Nikolay Malkin, Kilian Fatras, Lazar Atanackovic, Yanlei Zhang, Guillaume Huguet, Guy Wolf, and Yoshua Bengio. Simulation-free schrödinger bridges via score and flow matching. arXiv preprint arXiv:2307.03672 , 2023.

Cédric Villani et al. Optimal transport: old and new , volume 338. Springer, 2009.

Liu Yang and George Em Karniadakis. Potential flow generator with l 2 optimal transport regularity for generative models. IEEE Transactions on Neural Networks and Learning Systems , 33(2):528-538, 2020.

Jiongmin Yong and Xun Yu Zhou. Stochastic controls: Hamiltonian systems and HJB equations , volume 43. Springer Science & Business Media, 2012.

Qinsheng Zhang and Yongxin Chen. Path integral sampler: A stochastic control approach for sampling. In The Tenth International Conference on Learning Representations . OpenReview.net, 2022.

Yu Zheng, Furui Liu, and Hsun-Ping Hsieh. U-air: When urban air quality inference meets big data. In Proceedings of the 19th ACM SIGKDD international conference on Knowledge discovery and data mining , pages 1436-1444, 2013.

A Stochastic Optimal Control 448

The data generation task can also be interpreted as a stochastic optimal control (SOC) problem 449 [4, 17, 35, 50, 21, 5] whose cost function J is defined as: 450

where 1 2 ∥∇ x Ψ( X t , t ) ∥ 2 denotes the running cost, and Ψ( X T ) denotes the terminal cost. The above 451 SOC problem can be solved by dynamic programming [4, 6]. 452

Further, let V ( X t , t ) = inf J ( X t , t ) be the value function/optimal-cost-to-go of the SOC problem, 453 then the corresponding Hamilton-Jacobi-Bellman (HJB) partial differential equation (PDE) [15, 49] 454 is given by 455

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 .

C Experimental Environment

456

457

458

459

460

For the hardware environment of the experiments, we use a single NVIDIA A100-PCIE-40GB GPU 461 and an Intel(R) Xeon(R) Gold-6248R-3.00GHz CPU. For the software environment, the Python 462 version is 3.9.7, the CUDA version 11.7, and the Pytorch version is 2.0.1. 463

464

465

466

467

468

469

470

471

472

473

474

475

476

477

478

479

480

481

482

483

484

485

486

487

488

489

490

491

492

493

494

495

496

497

498

499

500

501

502

503

504

505

506

507

508

509

510

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: See Sec. 6

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [NA]

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer:[Yes]

Justification: See supplemental material

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: See supplemental material

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: See supplementary.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: See Sec. 4

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes] .

Justification: See Appendix. C

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

10. Broader Impacts

511

512

513

514

515

516

517

518

519

520

521

522

523

524

525

526

527

528

529

530

531

532

533

534

535

536

537

538

539

540

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: See Sec. 6

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA] .

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: See Sec. 4

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [Yes]

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA] .

541

Peer reviews:
comments: 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.

STRENGTHS:
S1. The approach of using Lagrangian Wasserstein flow for time series imputation is interesting and novel.

S2. Compared to other diffusion model-based time series imputation methods, this method seems more efficient.

S3. The experimental results given in the paper are satisfactory.

WEAKNESSES:
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:
Please see my weaknesses.

And there is another question,
Q. Can the VAE-based potential function be used alone for time series imputation?

LIMITATIONS:
CRPS is not used as the evaluation metric in the experiments.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-09

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

STRENGTHS:
1. The overall presentation of paper is clear. 
2. The idea of using the conditional Lagrangian Wasserstein flow for time series imputation is novel.
3. The proposed method is assessed on the real-world datasets and shows competitive performance compared to the state-of-the-art methods.

WEAKNESSES:
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:
What is the difference between the data generation processes in Euclidean space and Wasserstein space?

LIMITATIONS:
The ablation study on the effectiveness of Rao-Blackwellization should be conducted on more than one dataset (PM 2.5).
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-10

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

STRENGTHS:
- The contributions of this paper are significant.
- This work is built on a solid theoretical foundation.
- The experimental results and relevant ablation study findings are provided.

WEAKNESSES:
- 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:
Can the proposed method be used to compute the likelihood of the generated data samples?

LIMITATIONS:
The deterministic sampling process may introduce bias, potentially failing to accurately reflect the target marginal distribution.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-11

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n* Originality: Time series imputation is an important problem and I have not seen any flow matching works in towards this application. Presents a new improvement to the standard flow matching framework with a “Rao-Blackwellization”.\n\nWEAKNESSES:\n- Quality: Experiments are limited in scope and the it is not clear whether the “Rao-Blackwellization” (the main methodological novelty) step reliably improves performance.\n    - 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.\n    - More than 2 time series datasets would allow an understanding of when this “Rao-Blackwellized” sampler is valid.\n    - No error bars (even though this is claimed in the checklist)\n    - 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.\n- 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).\n- Significance: Without further experimental benchmarking especially with regards to the novel sampler, the applicability of this method is very limited.\n\nQUESTIONS:\nHow is this a Rao-Blackwellization, and what does the paper \n\nMinor comments:\n\nShrodinger —> Schrodinger throughout\n\nparagraph line 146 is confusing. this equation is strange given that p(x^tar) and p(x^input) are dependent if taking batches.\n\nLIMITATIONS:\nI'm a bit confused about the limitations. It is claimed that \n\n> One limitation of CLWF is that the samples obtained are not diverse enough as we use ODE for\n> inference, which results in slightly higher test (continuous ranked probability score) CRPS compared\n> to previous works, e.g., CSDI. Therefore, for future work, we will seek suitable approaches to\n> accurately model the diffusion term in the SDE.\n\nOther 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.", "TITLE": "Official Review", "RECOMMENDATION": 2, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 1, "DATE": "2024-07-12"}

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

STRENGTHS:
1. The paper is well-written, clear, and easy-to-follow.
2. The use of a Rao-Blackwellized sampler and a variational autoencoder to integrate prior information enhances the model's performance, providing a more robust imputation process.

WEAKNESSES:
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:
I have no other questions.

LIMITATIONS:
Yes
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-14