Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Conditional Lagrangian Wasserstein Flow for Time Series Imputation

Weizhu Qian, Dalin Zhang, Yan Zhao, Yunyao Cheng

TL;DR

This work introduces Conditional Lagrangian Wasserstein Flow (CLWF) for fast, accurate time series imputation by modeling data generation as a conditional optimal transport problem in Wasserstein space and learning a velocity field via the principle of least action. It combines flow matching to estimate the velocity, a time-dependent denoising autoencoder to approximate the gradient of a task-specific potential, and a resampling strategy to reduce variance, with Rao-Blackwellization further improving sampling efficiency. CLWF provides a simulation-free training regime and an ODE-based sampler, enabling faster convergence than diffusion-based imputation methods while maintaining competitive accuracy on real and synthetic datasets. The approach establishes links between optimal transport, stochastic control, and path measures, and demonstrates robust performance with ablations highlighting the benefits of variance reduction, fewer sampling steps, and the resampling trick for improved imputations.

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 (CLWF). Following the principle of least action in Lagrangian mechanics, we learn the velocity by minimizing the corresponding kinetic energy. Moreover, to enhance the model's performance, we estimate the gradient of a task-specific potential function using a time-dependent denoising autoencoder and integrate it into the base estimator to reduce the sampling variance. Finally, the proposed method demonstrates competitive performance compared to other state-of-the-art imputation approaches.

Conditional Lagrangian Wasserstein Flow for Time Series Imputation

TL;DR

This work introduces Conditional Lagrangian Wasserstein Flow (CLWF) for fast, accurate time series imputation by modeling data generation as a conditional optimal transport problem in Wasserstein space and learning a velocity field via the principle of least action. It combines flow matching to estimate the velocity, a time-dependent denoising autoencoder to approximate the gradient of a task-specific potential, and a resampling strategy to reduce variance, with Rao-Blackwellization further improving sampling efficiency. CLWF provides a simulation-free training regime and an ODE-based sampler, enabling faster convergence than diffusion-based imputation methods while maintaining competitive accuracy on real and synthetic datasets. The approach establishes links between optimal transport, stochastic control, and path measures, and demonstrates robust performance with ablations highlighting the benefits of variance reduction, fewer sampling steps, and the resampling trick for improved imputations.

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 (CLWF). Following the principle of least action in Lagrangian mechanics, we learn the velocity by minimizing the corresponding kinetic energy. Moreover, to enhance the model's performance, we estimate the gradient of a task-specific potential function using a time-dependent denoising autoencoder and integrate it into the base estimator to reduce the sampling variance. Finally, the proposed method demonstrates competitive performance compared to other state-of-the-art imputation approaches.

Paper Structure

This paper contains 41 sections, 4 theorems, 59 equations, 4 figures, 10 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Preliminaries
  3. Stochastic Differential Equations
  4. Optimal Transport
  5. Shrödinger Bridge
  6. Lagrangian Mechanics
  7. Conditional Lagrangian Wasserstein Flow for Time Series Imputation
  8. Time Series Imputation
  9. Interpolation in Wasserstein Space
  10. Velocity Estimation via Flow Matching
  11. Gradient of Potential Function
  12. Resampling Trick
  13. The Algorithms
  14. Discussion
  15. Experiments
  16. ...and 26 more sections

Key Result

Theorem 2.1

Pontryagin's Maximum Principle (PMP) evans2024introduction. If the $u^*$ is the optimal solution to the optimal control problem Eq. eq:soc_cost, then there exists a function $\lambda$ solution of the costate/adjoint equation for which

Figures (4)

  • Figure 1: The overall training process of Conditional Lagrangian Wasserstein Flow.
  • Figure 2: Visualization of the test imputation results on the synthetic data, green dots are the conditions, blue dots are the imputation results, and red dots are the ground truth.
  • Figure 3: The overall theoretical framework.
  • Figure 4: The proposed resampling process for inference.

Theorems & Definitions (5)

  • Theorem 2.1
  • Theorem 2.2
  • Definition 3.1: Sufficient statistic
  • Theorem 3.2
  • Theorem 3.3