Flow Matching with Missing Data

2026-07-30

AI Summary

This work addresses the challenge of applying flow matching to real-world data with missing values, a setting where traditional methods—reliant on fully observed samples—struggle. The authors treat missing coordinates as latent variables and propose averaging the flow matching loss over their possible values to enable accurate modeling of incomplete observations. Under the missing completely at random (MCAR) assumption, they theoretically establish equivalence between their objective and that of full-data flow matching, showing that missingness affects only the imputation process, not the learning target. Moreover, they prove that a single learnable imputation suffices to attain the variance lower bound achievable with complete data. Experiments corroborate these theoretical insights, demonstrating that deterministic learnable imputation outperforms frozen imputation strategies and matches or exceeds the performance of both classical and deep-learning-based imputation baselines on real-world tabular datasets.

Abstract

Flow matching assumes fully observed training data, which many real-world applications rarely provide. We propose Missing-Data Flow Matching, which treats the missing coordinates of training samples as latent variables and averages the flow matching loss over the values they could take. We first prove the correction is exact rather than approximate. Under missing completely at random with true completions, the incomplete-data objective equals the complete-data objective, so missingness changes nothing about what flow matching learns and the entire difficulty relocates to the completion model. Our finite-sample analysis then answers design questions that the algorithm leaves open, and the answers are not the ones intuition suggests. Missingness transfers estimator variance rather than adding it, one completion per example already matches complete-data variance exactly, and under a fixed evaluation budget one completion is optimal. A learned completion model contributes a single irreducible bias, which we bound by its expected conditional Wasserstein distance to the true completion law. Experiments numerically validate the theoretical predictions, show that deterministic rather than frozen imputation is what collapses the generated distribution, and place our method alongside strong classical and deep imputation baselines on real tabular data.

Problem

• Flow Matching
• Missing Data
• Imputation
• Latent Variables
• Wasserstein Distance

Innovation

• Flow Matching
• Missing Data
• Imputation
• Wasserstein Distance
• Latent Variables