Source-linked AI summary
Missing Data Imputation using Optimal Transport
Boris Muzellec, Julie Josse, Claire Boyer, Marco Cuturi
TL;DR
Missing values complicate machine-learning analyses, especially when incomplete rows cannot simply be discarded. The paper turns distributional agreement between random batches into an optimal-transport loss and develops non-parametric and parametric imputers. Across varied missingness settings, the proposed OT-based methods match or outperform established imputation approaches, including in challenging structured-data and benchmark comparisons.
Problem
Missing values are pervasive, and deleting incomplete rows becomes impractical in high-dimensional datasets, motivating methods that produce plausible imputations.
Method
The paper minimizes Sinkhorn divergences between randomly sampled batches using direct value optimization or differentiable parametric imputers trained round-robin.
Results
The OT-based methods match or outperform baseline and deep-learning imputers across datasets and missingness settings, while faithfully recovering complex non-linear structures in toy data.
Takeaways & Limitations
Non-parametric OT imputation can preserve complex global data structure without parametric assumptions, while trained parametric models additionally support out-of-sample imputation.
Takeaways & Limitations
Direct imputation cannot generate imputations for out-of-sample points, and its initialization criterion is a priori legitimate only under MCAR.
Abstract
from arXiv · showhide
Missing data is a crucial issue when applying machine learning algorithms to real-world datasets. Starting from the simple assumption that two batches extracted randomly from the same dataset should share the same distribution, we leverage optimal transport distances to quantify that criterion and turn it into a loss function to impute missing data values. We propose practical methods to minimize these losses using end-to-end learning, that can exploit or not parametric assumptions on the underlying distributions of values. We evaluate our methods on datasets from the UCI repository, in MCAR, MAR and MNAR settings. These experiments show that OT-based methods match or out-perform state-of-the-art imputation methods, even for high percentages of missing values.
1. Introduction
Missing values are pervasive, motivating imputation methods that preserve the data’s joint and marginal distributions. This paper proposes optimal transport as a scalable, differentiable basis for defining imputation losses and algorithms.
- Motivation: Missing values are widespread in collected datasets, and more than 150 implementations address missing-data handling.Methods differ by analytical objective, variable type, and assumptions about the data.
- Motivation: Imputation fills missing entries with plausible values while supporting downstream machine-learning analyses on completed data.Desirable imputations preserve both joint and marginal data distributions.
- Optimal transport rationale: Optimal transport offers geometrically meaningful distributional distances that can be computed scalably and differentiably through regularization.These properties have enabled OT-based losses in several machine-learning applications.
- Contributions: The proposed loss translates the intuition that two random batches from the same dataset should share a distribution.The paper develops both non-parametric and parametric algorithms to minimize this loss.
- Contributions: The non-parametric method combines global data shape with local features, while the parametric method supports out-of-sample imputation after training.The parametric framework can be combined with imputers such as Multi-Layer Perceptrons.
2. Background
Missingness mechanisms affect what can be learned from observed data, while deleting incomplete rows becomes impractical in high dimensions. The paper reviews optimal transport and motivates differentiable Sinkhorn-based losses for imputation.
- Missing-data mechanisms: MCAR missingness is independent of the data, whereas MAR missingness depends only on observed values and MNAR missingness depends on unobserved values.MNAR values can introduce important bias because their missingness probability depends on the missing values themselves.
- Why imputation: Deleting incomplete observations becomes impractical as dimensionality increases: with 1% independent entrywise missingness, only around 5% of rows remain complete when d = 300.Around 95% of rows remain complete when d = 5 under the same assumption.
- Existing approaches: Existing imputation methods include explicit joint models, low-rank models, conditional sequential imputation, nearest neighbors, and random forests.The proposed methods are positioned at the intersection of global and local approaches.
- Optimal transport: Optimal transport compares weighted empirical distributions by minimizing transport cost over plans that match their mass vectors.Quadratic ground costs yield Wasserstein distances, but exact computation is costly and non-differentiable.
- Sinkhorn divergences: Entropic regularization makes OT differentiable and solvable with Sinkhorn iterations, while debiasing produces Sinkhorn divergences that are positive and tractable.Sinkhorn divergences provide a differentiable proxy for Wasserstein distances with little additional computational cost.
- Gradient-based imputation: Gradients of OT between batches move imputed points toward target points specified by the optimal transport plan.This barycentric transport-map interpretation motivates gradient-based imputation updates.
3. Imputing Missing Values using OT
The paper formulates missing-value imputation as minimizing Sinkhorn divergences between batches, then optimizes either individual imputations or differentiable parametric imputers. It also extends the approach to round-robin variable-wise updates and practical training refinements, while noting limits for out-of-sample use and non-MCAR settings.
- Direct imputation: Direct imputation initializes missing entries near observed means, then minimizes batch Sinkhorn divergences with gradient updates.RMSprop is used in the described implementation.
- OT-based imputation: Sinkhorn divergences between batches provide the loss for imputing missing values, with one parameter corresponding to each imputed value.The loss is optimized over sampled batches to encourage distributional agreement.
- Parametric imputation: Parametric imputers must be differentiable so the batch Sinkhorn loss can back-propagate through imputations and train reusable out-of-sample models.The trained parameter and completed data are both returned.
- Round-robin imputation: Round-robin imputation fits one parameter set per variable, cyclically predicts each missing coordinate from the remaining variables, and repeats the process.Algorithm 3 performs gradient updates for each variable before moving through the cycle and repeating until convergence or a fixed iteration limit.
- Practical refinements: Several refinements average losses over multiple batch pairs, adapt variable order, and use validation data to monitor training.A proposed MCAR-specific sampling strategy uses one batch with the target variable observed and another with it missing.
- Limitations: Direct imputation cannot directly generate imputations for out-of-sample incomplete points, while the MCAR-specific sampling refinement is not generally legitimate under MAR or MNAR.Under MAR or MNAR, non-missing values generally do not share the unknown distribution of missing values.
4. Experimental Results
Across toy and UCI experiments, OT-based imputers handle nonlinear structure and difficult missingness mechanisms without requiring a parametric model. They match or outperform established and deep-learning baselines, while parametric round-robin models also retain strong out-of-sample performance.
- Toy experiments: Sinkhorn faithfully imputes complex nonlinear toy distributions where conditional-expectation ice imputations fail.The comparison uses 20% MCAR missingness on S-shapes, half-moons, and concentric circles.
- Parametric models: MLP RR has MAE and RMSE comparable to Linear RR but generally better W2 scores, supporting more expressive imputers for optimizing the OT objective.The comparison concerns the round-robin parametric models evaluated with the batched Sinkhorn loss.
- Method comparison: The direct non-parametric imputer is generally the best-performing OT-based method across datasets and metrics, with its advantage increasing as missingness grows.Its imputation model is not restricted by a parametric assumption.
- MAR and MNAR: Under MAR and MNAR mechanisms, the proposed methods remain robust despite not explicitly modeling the missingness mechanisms.Exceptions remain concentrated on the same datasets across mechanisms and missingness rates.
Conclusion
The paper defines missing-data imputation through an OT-based loss and proposes both unrestricted and parametric algorithms. Across varied MCAR, MAR, and MNAR settings, the methods are highly competitive, while the two algorithm classes trade distributional flexibility against out-of-sample capability.
- The loss is the expected Sinkhorn divergence between randomly sampled batches.
- The unrestricted algorithm estimates one parameter per imputed value without parametric assumptions on the data distribution.
- The parametric algorithm fits a model that enables out-of-sample imputation after training.
- Experiments across varied datasets and MCAR, MAR, and MNAR settings found the proposed models very competitive, including against recent deep-learning methods.
- The results support the loss as an optimizable proxy for imputation metrics, while further theoretical study remains future work.
A. Appendix
The appendix reports experiments across multiple missingness regimes and finds the proposed algorithms competitive with state-of-the-art methods. Runtime comparisons identify Sinkhorn divergences as the main computational bottleneck.
- Experimental regimes: The appendix covers 30% MAR, two 30% MNAR mechanisms, and MCAR experiments at 10%, 30%, and 50% missingness.
- Evaluation protocol: Error bars represent ±1 standard deviation across 30 runs on each dataset.
- Performance: The proposed methods, including Algorithm 1 and Algorithm 3 with linear or shallow MLP imputers, are competitive with state-of-the-art methods across missing-data regimes.
- Runtime: Algorithm 1 runs about as fast as VAEAC and faster than GAIN and MIWAE, whereas round-robin methods are slowest overall.
- Runtime: The number of computed Sinkhorn batch divergences is the computational bottleneck, suggesting that more complex round-robin imputers may add little runtime cost.