Conditional Lagrangian Wasserstein Flow for Time Series Imputation

We treat the data generation task as an initial value problem (IVP), in which $X_{0}\in\mathbb{R}^{d}$ is the initial data (e.g., some random noise) at the initial time $t=0$, and $X_{T}\in\mathbb{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 $\mu_{t}:[0,T]\times\mathbb{R}^{d}\rightarrow\mathbb{R}^{d}$, and a positive Borel measurable time-dependent diffusion function $\sigma_{t}:[0,T]\rightarrow\mathbb{R}^{d}_{>0}$. Accordingly, the Itô form of the SDE can be described as follows Oksendal (2013):

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) Risken and Frank (2012) describing the evolution of the marginal density $p_{t}(X_{t})$ reads:

where $\Delta p_{t}=\divergence(\nabla 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 Villani et al. (2009); Santambrogio (2015); Peyré et al. (2019). In the Kantorovich’s formulation of the OT problem, the transport costs are minimized with respect to some probabilistic couplings/joint distributions Villani et al. (2009); Santambrogio (2015); Peyré et al. (2019). Let $p_{0}$ and $p_{T}$ be two Borel probability measures with finite second moments on the space $\Omega\in\mathbb{R}^{d}$. $\Pi(p_{0},p_{T})$ denotes a set of transport plans between these two marginals. Then, the Kantorovich’s OT problem is defined as follows:

where $\Pi(p_{0},p_{T})=\big{\{}\pi\in\mathcal{P}(\mathcal{X}\times\mathcal{Y}):(\pi^{x})_{\#}\pi=p_{0},(\pi^{y})_{\#}\pi=p_{T}\big{\}}$, with $\pi^{x}$ and $\pi^{x}$ being two projections of $\mathcal{X}\times\mathcal{Y}$ on $\Omega$. The minimizer of Eq .3, $\pi*$, 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_{\mathrm{KL}}(\pi\|p_{0}\otimes 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 Cuturi (2013) 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 Tong et al. (2023, 2024); Pooladian et al. (2023).

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 Léonard (2012). Assume that $\Omega\in C^{1}([0,T],\mathbb{R}^{d})$, $\mathcal{P}(\Omega)$ is a probability path measure on the path space $\Omega$, then the goal of the SB problem aims to find the following optimal path measure:

where the Kullback–Leibler (KL) divergence $D_{\mathrm{KL}}(\mathbb{P}\|\mathbb{Q})=\begin{cases}\log\frac{\mathrm{d}\mathbb{P}}{\mathrm{d}\mathbb{Q}}\mathrm{d}\mathbb{P},&\mbox{if}~{}\mathbb{P}\ll\mathbb{Q},\\ +\infty,&\mbox{otherwise},\end{cases}$ and $\mathbb{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 Gushchin et al. (2024); Koshizuka and Sato (2022); Liu et al. (2024):

where $\theta$ is the parameters of the variational drift function $\mu_{t}$.

In this section, we formulate the data generation problem under the framework of Lagrangian mechanics Arnol’d (2013). Let $p_{t}$ and $\dot{p_{t}}=\frac{\mbox{d}p_{t}}{\mbox{d}t}$ be the density and law of the generalized coordinates $X_{t}$, respectively. $\mathcal{K}(p_{t},\dot{p_{t}},t)$ is the kinetic energy, and $\mathcal{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 $\dot{p_{t}}$ in the Wasserstein space. The kinetic energy $\mathcal{K}(x_{t},\mu_{t},t)$ and potential energy $\mathcal{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 defined as follow:

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

with the boundary condition $\frac{\mathrm{d}X_{t}}{\mathrm{d}t}=\mu(X_{t},t),~{}q_{0}=p_{0},~{}q_{T}=p_{T}$.

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 $\mathbb{R}^{K\times 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^{\mathrm{obs}}\in\mathbb{R}^{K\times L}$ are partitioned into the imputation target $x^{\mathrm{tar}}\in\mathbb{R}^{K\times L}$ and the conditional data $x^{\mathrm{cond}}\in\mathbb{R}^{K\times L}$.

As a result, the missing data points $x^{\mathrm{tar}}$ can be generated based on the conditions $x^{\mathrm{cond}}$ joint with some uninformative initial distribution $x_{0}\in\mathbb{R}^{K\times L}$ (e.g., Gaussian noise) at time $t=0$, then the imputation task can be described as: $x^{\mathrm{tar}}\sim p(x^{\mathrm{tar}}|x^{\mathrm{cond}}_{0})$, where the total input of the model is $x^{\mathrm{input}}_{0}:=(x^{\mathrm{cond}},x_{0})\in\mathbb{R}^{K\times L\times 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:

where $\Omega\in\mathbb{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}\sim\mathcal{N}(0,\sigma^{2}_{0})$ at the initial time $t=0$ and the data samples $X_{T}\sim p(x^{\mathrm{tar}})$ at the terminal time $t=T$, respectively. Afterwards, the interpolation method is used to construct the intermediate samples $X_{t}\sim p(X_{t}|X_{0},X_{T},t)$, where $t\sim\mathrm{uniform}(0,T)$ Liu et al. (2022); Albergo and Vanden-Eijnden (2023); Tong et al. (2024). More specifically, we design the following sampling approach:

where $\gamma_{t}\sim\mathcal{N}(0,\sigma^{2}_{\gamma})$ is some random noise with variance $\sigma_{\gamma}$ injected to the target data samples for improving the coupling’s generalization property, and $\alpha(t)\geq 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 Tong et al. (2023, 2024); Pooladian et al. (2023) 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^{\mathrm{input}}_{t}:=(x^{\mathrm{cond}},x_{t})$ as the input for computing the velocity of the Wasserstein flow.

To estimate the velocity of the Wasserstein flow $\mu_{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 Chen et al. (2018); Onken et al. (2021). 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 Liu et al. (2022); Tong et al. (2023); Albergo and Vanden-Eijnden (2023), 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 $\theta$. 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 $\mu_{\theta}(x^{\mathrm{input}}_{t},t)$ can be learned by

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}):\mathbb{R}^{d}\times[0,T]\rightarrow\mathbb{R}$ be the task-specific potential function depending on the generalized coordinates $X_{t}$ Yang and Karniadakis (2020); Onken et al. (2021); Neklyudov et al. (2023b). 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 (SOC) problem Bellman (1966); Fleming and Rishel (2012); Nüsken and Richter (2021); Zhang and Chen (2022); Holdijk et al. (2023); Berner et al. (2024), then the existence of such $U_{t}(X_{t})$ is assured by Pontryagin’s Maximum Principle (PMP) Evans (2024).

To estimate $v_{t}(X_{t},t)$, according to the Lagrangian Eq. 6, we assume that the potential function takes the form $U_{t}(X_{t})\approx-\log\mathcal{N}(X_{t}|\hat{X_{t}},\sigma^{2}_{p})$, where $\hat{X_{t}}$ the learned mean and $\sigma^{2}_{p}$ is the pre-defined variance. As a result, the derivative is $\nabla_{x}U(X_{t})=\frac{X_{t}-\hat{X_{t}}}{\sigma^{2}_{p}}$. In terms of practical implementation, we parameterize $\nabla_{x}U(X_{t})$ via a Variational Autoencoder (VAE) Kingma and Welling (2014). More specifically, we pre-train a VAE on the total observed time series data $X^{\mathrm{obs}}$. Afterwards, the reconstruction discrepancies of the VAE are used to approximate the task-specific $v^{\phi}(X_{t},t)$ depending on $X_{t}$:

where $\mathrm{VAE}(X_{t})$ represents the reconstruction output of the pre-trained VAE model with input $X_{t}$, and $\sigma^{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},\mu_{t}^{\theta}(X_{t},t),t)$ for $X_{t+1}$ with the Euler method and $\mu_{t}^{\theta}(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 Casella and Robert (1996) for time series data imputation using Eq. 14 and Eq. 16. To this end, we first treat $\mathcal{S}(X_{t+1}|X_{t},\mu_{t}^{\theta}(X_{t},t),t)$ be the base estimator for $X_{t+1}$ with $\mathbb{E}[\mathcal{S}^{2}]<\infty$ for all $X_{t+1}$. And we assume $\mathcal{T}(X_{t},v_{t}^{\phi}(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 $\mathcal{S}^{*}=\mathbb{E}[\mathcal{S}|\mathcal{T}]$ to generate the missing time series data. Then according to the Rao-Blackwell theorem Casella and Robert (1996), we have

where the inequality is strict unless $\mathcal{S}$ is a function of $\mathcal{T}$. Eq. 17 suggests we can construct a more powerful sampler with smaller errors than the base ODE sampler $\mathcal{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^{\mathrm{obs}}$ are partitioned into the target data and conditional data for training. Next, the data pairs of $x^{\mathrm{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 $\frac{\mbox{d}X_{t}}{\mbox{d}t}$ by computing $\frac{x^{\mathrm{tar}}-x_{0}}{T}$. Subsequently, we use the joint distribution of the conditional information $x^{\mathrm{cond}}$ and the intermediate variable $x_{t}$, $x^{\mathrm{input}}$ as the total input to feed the variational flow model $\mu_{t}^{\theta}$ 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^{\mathrm{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^{\mathrm{cond}}$ to formulate the joint variable $x^{\mathrm{input}}$. Note that during the trajectory sampling $x_{0}$ will evolve over time, while $x^{\mathrm{cond}}$ remain invariant. We use $x^{\mathrm{input}}$ as the input of the flow model $\mu_{t}^{\theta}$ 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.

We use two public multivariate time series datasets for validation. The first dataset is the PM 2.5 dataset Zheng et al. (2013) 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 Silva et al. (2012) 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.

For comparison, we select the following state-of-the-art timer series imputation methods as the baselines: 1) GP-VAE Fortuin et al. (2020), which is combines a VAE model and a Gaussian Process prior; 2) CSDI Tashiro et al. (2021), which is based on the conditional diffusion model; 3) CSBI Chen et al. (2023), which is based on the Schrödinger bridge diffusion model; 4) DSPD-GP Biloš et al. (2023), 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 Tashiro et al. (2021). 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 $\sigma_{0}$ for the initial noise $X_{0}$ is $0.1$, and the standard deviation $\sigma_{\gamma}$ for the injected noise $\gamma_{t}$ $0.001$. The coefficient $\sigma_{p}^{2}$ in the derivative of the potential function is $0.01$.

Diffusion models, such as DDPMs Ho et al. (2020) and SBGM  Song et al. (2020), 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 Gushchin et al. (2024) The flowing matching methods have also been proposed based on the OT theory Lipman et al. (2022); Liu (2022); Liu et al. (2023); Albergo and Vanden-Eijnden (2023); Albergo et al. (2023). In particular, mini-batch couplings are proposed to straighten the probability flows for fast inference Pooladian et al. (2023); Tong et al. (2024, 2023).

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 De Bortoli et al. (2021). SB-FBSDE proposes to use forward-backward (FB) SDE theory to solve the SB problem through likelihood training Chen et al. (2022). GSBM formulates a generalized Schrödinger Bridge matching framework by including the task-specific state costs for various data generation tasks Liu et al. (2024) NLSB chooses to model the potential function rather than the velocity function to solve the Lagrangian SB problem Koshizuka and Sato (2022). Action Matching Neklyudov et al. (2023a, b) 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 Nüsken and Richter (2021); Zhang and Chen (2022); Berner et al. (2024); Liu et al. (2024).

Many diffusion-based models have been recently proposed for time series imputation Lin et al. (2023); Meijer and Chen (2024). For instance, CSDI Tashiro et al. (2021) combines a conditional DDPM with a Transformer model to impute time series data. CSBI Chen et al. (2023) 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 Biloš et al. (2023) uses a Gaussian process as the noise generator. TDdiff Kollovieh et al. (2024) 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.