Technical Summary: Missing-Data Flow Matching (MDFM)

Problem Statement

Flow matching is a powerful generative modeling technique that learns a time-dependent vector field to transport a base distribution to a data distribution via a stable regression objective. However, standard flow matching assumes fully observed training endpoints. In real-world applications (e.g., medical records, sensor data, surveys), data is frequently incomplete.

Existing solutions typically employ point imputation (filling missing values with a single estimate like the conditional mean). While effective for predictive tasks relying on means, point imputation is detrimental to generative learning. Because flow matching learns to reproduce the distribution of its training endpoints, replacing a range of plausible missing values with a single point collapses the conditional distribution. This results in a generated distribution with zero conditional variance, failing to capture the true spread and dependencies of the data.

Methodology: Missing-Data Flow Matching (MDFM)

The authors propose Missing-Data Flow Matching (MDFM), a framework that treats missing coordinates as latent variables rather than attempting to fix them to a single value.

Core Mechanism

Instead of committing to one imputed value, MDFM treats the missing part of a sample, $X_{-\Lambda}$, as a random variable drawn from its conditional distribution given the observed part, $X_\Lambda$.

1. Latent Variable Treatment: For an incomplete observation $O = (\Lambda, X_\Lambda)$, the missing coordinates $X_{-\Lambda}$ are sampled from a completion distribution.
2. Averaged Loss: The flow matching loss is averaged over these sampled completions. The objective function becomes an expectation over the observed patterns, the observed data, and the completions of the missing parts.
3. Oracle vs. Learned:
   • Oracle Objective: Assumes access to the true conditional distribution $p^*(x_{-\Lambda} | x_\Lambda)$.
   • Learned Objective: Uses a learned completion model $q_\phi(x_{-\Lambda} | x_\Lambda, \Lambda)$, trained via self-supervised masking on an independent auxiliary dataset.

Theoretical Foundation

The paper establishes that under Missing Completely At Random (MCAR) conditions, if completions are drawn from the true conditional distribution, the MDFM objective is exactly equivalent to the complete-data flow matching objective.

• Exactness: Unlike score matching with missing data, which requires approximating intractable integrals, MDFM's correction is an exact equality. The regression target ($X_1 - X_0$) becomes explicit once an endpoint is completed.
• Relocation of Difficulty: The difficulty of missing data does not vanish but relocates entirely to the estimation of the conditional completion distribution. If the completion model is perfect, the flow matching model learns the true data distribution exactly.

Key Contributions

1. Exact Correction and Equivalence

The authors prove that under MCAR, training on incomplete data with true completions yields the same population objective and gradients as training on complete data. This holds for any differentiable point in the parameter space. The result extends to Missing At Random (MAR) but not Missing Not At Random (MNAR) without additional structural assumptions.

2. Finite-Sample Variance Analysis

The paper provides a rigorous finite-sample variance decomposition of the estimator, yielding counter-intuitive design insights:

• Variance Transfer, Not Addition: Missingness does not inherently increase the total estimator variance; it transfers variance from the observation level to the completion level. Averaging over completions removes this transferred variance.
• Optimality of One Completion: Under a fixed evaluation budget (total number of loss evaluations), using one completion per example is optimal. Drawing more completions per example is less efficient than collecting more incomplete examples.
• Bias Bound: For learned completion models, the bias in the population objective is bounded by the expected conditional Wasserstein-1 distance between the learned model and the true completion law. This bias is irreducible by increasing the sample size of the flow matching training; it requires a better completion model.

3. Empirical Validation

Experiments on synthetic and real tabular data (UCI datasets) validate the theoretical predictions:

• Stochastic vs. Deterministic: Deterministic imputation (e.g., conditional mean) collapses the generated distribution, tripling distributional error. Stochastic completion (resampling) preserves conditional variability and matches complete-data performance.
• Distributional Metrics: On real data, MDFM is competitive with strong imputation pipelines. Specifically, MICE achieves the best average ranks on sliced Wasserstein distance, energy distance, and MMD, while MDFM achieves the best average rank on covariance preservation among methods trained on incomplete data.
• Imputation Quality: While MDFM is competitive on imputation error (RMSE), its primary strength lies in capturing the joint distribution's second-moment structure.

Significance and Claims

The paper positions MDFM as a theoretically grounded solution that bridges the gap between missing data handling and generative modeling.

• Exactness: The primary significance is the proof that the correction is exact, not approximate, shifting the burden of missing data solely to the quality of the completion model.
• Design Guidance: The finite-sample analysis provides concrete, non-intuitive guidelines: specifically, that one should prioritize collecting more data examples over generating multiple imputations per example under fixed computational budgets.
• Generative Integrity: The work demonstrates that to maintain the integrity of a generative model (specifically its ability to reproduce conditional variance), one must avoid deterministic point imputation and instead utilize stochastic completion.

The authors acknowledge limitations, noting that the method relies on the assumption that the completion model can be trained effectively (requiring identifiable patterns) and that the theoretical guarantees currently focus on loss estimation rather than the direct quality of the generated vector field. The extension to MNAR scenarios remains an open challenge requiring additional structural assumptions.