# Distributional Matrix Completion via Nearest Neighbors in the Wasserstein Space

arXiv: https://arxiv.org/abs/2410.13112 · HTML: https://mlsanity.com/paper/2410.13112/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 11 of its 53 references that are in the map.
- First posted on arXiv 17 Oct 2024.

## Builds on (most cited first, 11 of 11)

- [Fast Computation of Wasserstein Barycenters](https://mlsanity.com/paper/1310.4375/) (2013)
- [Causal Matrix Completion](https://mlsanity.com/paper/2109.15154/) (2021)
- [Optimal Transport for Diffeomorphic Registration](https://mlsanity.com/paper/1706.05218/) (2017)
- [Wasserstein barycenters are NP-hard to compute](https://mlsanity.com/paper/2101.01100/) (2021)
- [Matrix Completion via Max-Norm Constrained Optimization](https://mlsanity.com/paper/1303.0341/) (2013)
- [Counterfactual inference in sequential experiments](https://mlsanity.com/paper/2202.06891/) (2022)
- [Distributional synthetic controls](https://mlsanity.com/paper/2001.06118/) (2020)
- [Doubly robust nearest neighbors in factor models](https://mlsanity.com/paper/2211.14297/) (2022)
- [Top-N Recommender System via Matrix Completion](https://mlsanity.com/paper/1601.04800/) (2016)
- [Variance Minimization in the Wasserstein Space for Invariant Causal Prediction](https://mlsanity.com/paper/2110.07064/) (2021)
- [Centered plug-in estimation of Wasserstein distances](https://mlsanity.com/paper/2203.11627/) (2022)

## Built on by (most cited first, 3 of 3)

- [N2: A Unified Python Package and Test Bench for Nearest Neighbor-Based Matrix Completion](https://mlsanity.com/paper/2506.04166/) (2025)
- [Synthetic Survival Control: Extending Synthetic Controls for "When-If" Decision](https://mlsanity.com/paper/2511.14133/) (2025)
- [Low-rank Distributional Matrix Completion](https://mlsanity.com/paper/2606.04176/) (2026)

## Abstract (arXiv)

We study the problem of distributional matrix completion: Given a sparsely observed matrix of empirical distributions, we seek to impute the true distributions associated with both observed and unobserved matrix entries. This is a generalization of traditional matrix completion, where the observations per matrix entry are scalar-valued. To do so, we utilize tools from optimal transport to generalize the nearest neighbors method to the distributional setting. Under a suitable latent factor model on probability distributions, we establish that our method recovers the distributions in the Wasserstein metric. We demonstrate through simulations that our method (i) provides better distributional estimates for an entry compared to using observed samples for that entry alone, (ii) yields accurate estimates of distributional quantities such as standard deviation and value-at-risk, and (iii) inherently supports heteroscedastic distributions. In addition, we demonstrate our method on a real-world dataset of quarterly earnings prediction distributions. We also prove novel asymptotic results for Wasserstein barycenters over one-dimensional distributions.