Flow Matching with Missing Data

The gist The Missing-Data Flow Matching method trains flow matching from partially observed data by treating missing coordinates as latent variables, proving that this correction is exact rather than approximate.

How it works

Flow matching assumes fully observed training data, but real-world applications often have missing coordinates which standard flow matching cannot handle because the loss requires a complete data point as an endpoint <ref:2607.28698#pg2>. Missing-Data Flow Matching (MDFM) treats these missing coordinates as latent variables and averages the flow matching loss over the values they could take <ref:2607.28698#pg3>. When data are missing completely at random (MCAR) with true completions, training incomplete data is exactly equivalent to training from complete data, meaning missingness changes nothing about what flow matching learns and the entire difficulty relocates to the completion model <ref:2607.28698#pg3>.

Key Theoretical Results

The paper proves that the oracle objective equals the complete-data objective when using true conditional completions under MCAR <ref:2607.28698#pg2>. Furthermore, it derives a finite-sample variance decomposition showing that one oracle completion matches complete-data variance and is optimal under a fixed evaluation budget <ref:2607.28698#pg2>. For learned completions, the population-objective bias is bounded by the conditional Wasserstein distance to the truth <ref:2607.28698#pg3>.

Variance Decomposition and Finite Samples

The estimator variance for the oracle method can be decomposed into three components: the between observation variance, the base Monte Carlo variance, and the conditional completion variance <ref:2607.28698#pg2>. The analysis shows that missingness transfers estimator variance rather than adding it, and that one completion per example already matches complete-data variance exactly <ref:2607.28698#pg2>. Under a fixed evaluation budget, extra completions help only when the number of examples is fixed, suggesting that for scarce data, spending resources on more incomplete observations is always better <ref:2607.28698#pg2>.

Learned Completion Model Analysis

When using a learned completion model qφ, the error splits into estimation and bias components <ref:2607.28698#pg4>. The bias term is bounded by the conditional Wasserstein distance to the truth via Lemma 5.7 <ref:2607.28698#pg4>. The total bound on the error combines an estimation term that shrinks with evaluation data and completions, and a bias term that shrinks only with auxiliary data and modeling effort spent on qφ <ref:2607.28698#pg4>.

Experimental Validation

Experiments validate the theoretical claims on synthetic data, showing that stochastic completion preserves generated conditional variability <ref:2607.28698#pg2>. On real tabular data, MDFM is competitive with strong imputation pipelines and performs best on covariance preservation <ref:2607.28698#pg2>. The results show that every stochastic strategy, resampled or frozen and oracle or fitted, recovers 97% or more and matches complete-data flow matching <ref:2607.28698#pg2>.

Missingness Mechanisms

The paper details how missingness mechanisms affect the objective. Under MCAR and MAR, the weight is identically one and the gap should vanish, while under MNAR, the naive objective is a reweighted complete data objective with a weight that averages to one but need not be constant <ref:2607.28698#pg2>. The practical content of this finding is that the method does not silently fail under MNAR, it optimizes a tilted objective that overweights regions of the data space where the observed coordinates make missingness patterns more likely <ref:2607.28698#pg4>.

Conclusion

The work introduces MDFM, which trains flow matching directly from partially observed data by treating missing coordinates as resampled latent variables <ref:2607.28698#pg2>. The correction is exact rather than approximate, since under MCAR with true completions the incomplete-data objective equals the complete-data objective <ref:2607.28698#pg2>. The open case remains missing not at random (MNAR), where the required posterior conditions on the mask and is generally not identifiable without extra structure <ref:2607.28698#pg2>. The research shows that a learned completion model contributes a single irreducible bias, which its expected conditional Wasserstein distance to the truth bounds <ref:2607.28698#pg2>. The method is competitive on real data and best on covariance preservation <ref:2607.28698#pg2>.

--- Page 30 ---

The gist The Missing-Data Flow Matching method trains flow matching from partially observed data by treating missing coordinates as latent variables, proving that this correction is exact rather than approximate.

Key Theoretical Results

The paper proves that the oracle objective equals the complete-data objective when using true conditional completions under MCAR <ref:2607.28698#pg2>. Furthermore, it derives a finite-sample variance decomposition showing that one oracle completion matches complete-data variance and is optimal under a fixed evaluation budget <ref:2607.28698#pg2>. For learned completions, the population-objective bias is bounded by the conditional Wasserstein distance to the truth <ref:2607.28698#pg3>.

Conclusion

The work introduces MDFM, which trains flow matching directly from partially observed data by treating missing coordinates as latent variables <ref:2607.28698#pg2>. The correction is exact rather than approximate, since under MCAR with true completions the incomplete-data objective equals the complete-data objective <ref:2607.28698#pg2>. The open case remains missing not at random (MNAR), where the required posterior conditions on the mask and is generally not identifiable without extra structure <ref:2607.28698#pg2>. The research shows that a learned completion model contributes a single irreducible bias, which its expected conditional Wasserstein distance to the truth bounds <ref:2607.28698#pg2>. The method is competitive on real data and best on covariance preservation <ref:2607.28698#pg2>.

Key Theoretical Results

The paper proves that the oracle objective equals the complete-data objective when using true conditional completions under MCAR <ref:2607.28698#pg2>.

1. Missing-Data Flow Matching (MDFM) will train flow matching directly from partially observed data by treating missing coordinates as latent variables and averaging over completions to ensure the incomplete-data objective equals the complete-data objective under MCAR. This allows the model to learn the full generative flow, rather than collapsing a range of plausible values into a single point, thus preserving the true conditional spread.

2. The system will utilize a learned completion model, qφ(x−Λ xΛ,Λ), which is trained via self-supervised masking to sample plausible missing parts from the observed part. This completion model acts as an engine, and its quality bounds the final bias by its conditional Wasserstein distance to the truth, ensuring that improvement depends on both more data and a better completion model.

3. The system will employ a variance decomposition strategy where it tracks three components: between-observation variance, base Monte Carlo variance, and conditional completion variance, allowing for an optimal allocation of resources. Specifically, the analysis shows One oracle completion per example already matches complete-data variance, meaning collecting more incomplete examples is always better than drawing more completions under a fixed budget.

4. When data are Missing Not At Random (MNAR), the system will optimize a tilted objective where the weight w(x) is identified as Ep⋆ [w(X)] = 1, allowing it to learn from observed patterns without silently failing, as the naive objective is shown to be a reweighted complete data objective.

5. The system will employ a dynamic completion strategy based on budget constraints, where it determines that a single completion is the right choice whenever more incomplete observations are available by showing that the estimator variance is minimized at K=1 for a fixed evaluation budget.

Flow matching assumes fully observed training data, which many real-world applications rarely provide. We study 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 that the oracle correction is exact: completing from the true pattern-conditioned posterior recovers the complete-data flow-matching objective under any missingness mechanism. Under missing completely at random (MCAR) and missing at random (MAR), this posterior equals the ordinary complete-data conditional, while under missing not at random (MNAR) it generally depends on the missingness mechanism and may require additional structure to identify. Our finite-sample analysis then answers design questions that the algorithm leaves open. Missingness transfers estimator variance rather than adding it, and for the fixed-parameter loss estimator one completion per example matches complete-data variance and is optimal under a fixed evaluation budget. A learned completion model contributes population completion 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 reduces generated conditional variability, and place our method alongside strong classical and deep imputation baselines on real tabular data.