# Learning Counterfactual Distributions via Kernel Nearest Neighbors

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

## Key facts

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

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

- [Kernel Mean Embedding of Distributions: A Review and Beyond](https://mlsanity.com/paper/1605.09522/) (2016)
- [Matrix Completion and Low-Rank SVD via Fast Alternating Least Squares](https://mlsanity.com/paper/1410.2596/) (2014)
- [Cross-validation: what does it estimate and how well does it do it?](https://mlsanity.com/paper/2104.00673/) (2021)
- [Generalized Low Rank Models](https://mlsanity.com/paper/1410.0342/) (2014)
- [Learning Theory for Distribution Regression](https://mlsanity.com/paper/1411.2066/) (2014)
- [Kernel Thinning](https://mlsanity.com/paper/2105.05842/) (2021)
- [Causal Matrix Completion](https://mlsanity.com/paper/2109.15154/) (2021)
- [Semiparametric counterfactual density estimation](https://mlsanity.com/paper/2102.12034/) (2021)
- [Kernel Methods for Causal Functions: Dose, Heterogeneous, and Incremental Response Curves](https://mlsanity.com/paper/2010.04855/) (2020)
- [Counterfactual Mean Embeddings](https://mlsanity.com/paper/1805.08845/) (2018)
- [Missing Not at Random in Matrix Completion: The Effectiveness of Estimating Missingness Probabilities Under a Low Nuclear Norm Assumption](https://mlsanity.com/paper/1910.12774/) (2019)
- [Generalized Kernel Thinning](https://mlsanity.com/paper/2110.01593/) (2021)
- [Synthetic Interventions](https://mlsanity.com/paper/2006.07691/) (2020)
- [Distribution Compression in Near-linear Time](https://mlsanity.com/paper/2111.07941/) (2021)
- [Estimating Barycenters of Measures in High Dimensions](https://mlsanity.com/paper/2007.07105/) (2020)
- [Local nearest neighbour classification with applications to semi-supervised learning](https://mlsanity.com/paper/1704.00642/) (2017)
- [Causal effects based on distributional distances](https://mlsanity.com/paper/1806.02935/) (2018)
- [Design-based Analysis in Difference-In-Differences Settings with Staggered Adoption](https://mlsanity.com/paper/1808.05293/) (2018)
- [Matrix completion with data-dependent missingness probabilities](https://mlsanity.com/paper/2106.02290/) (2021)
- [Counterfactual inference in sequential experiments](https://mlsanity.com/paper/2202.06891/) (2022)
- [On counterfactual inference with unobserved confounding](https://mlsanity.com/paper/2211.08209/) (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)
- [Entrywise Inference for Missing Panel Data: A Simple and Instance-Optimal Approach](https://mlsanity.com/paper/2401.13665/) (2024)
- [Kernel-based independence tests for causal structure learning on functional data](https://mlsanity.com/paper/2311.08743/) (2023)
- [Distributional Bellman Operators over Mean Embeddings](https://mlsanity.com/paper/2312.07358/) (2023)
- [Doubly Robust Inference in Causal Latent Factor Models](https://mlsanity.com/paper/2402.11652/) (2024)
- [Jackknife multiplier bootstrap: finite sample approximations to the U-process supremum with applications](https://mlsanity.com/paper/1708.02705/) (2017)

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

- [N2: A Unified Python Package and Test Bench for Nearest Neighbor-Based Matrix Completion](https://mlsanity.com/paper/2506.04166/) (2025)
- [Causal Matrix Completion under Multiple Treatments via Mixed Synthetic Nearest Neighbors](https://mlsanity.com/paper/2603.11942/) (2026)

## Abstract (arXiv)

Consider a setting with multiple units (e.g., individuals, cohorts, geographic locations) and outcomes (e.g., treatments, times, items), where the goal is to learn a multivariate distribution for each unit-outcome entry, such as the distribution of a user's weekly spend and engagement under a specific mobile app version. A common challenge is the prevalence of missing not at random data, where observations are available only for certain unit-outcome combinations and the observation availability can be correlated with the properties of distributions themselves, i.e., there is unobserved confounding. An additional challenge is that for any observed unit-outcome entry, we only have a finite number of samples from the underlying distribution. We tackle these two challenges by casting the problem into a novel distributional matrix completion framework and introduce a kernel based distributional generalization of nearest neighbors to estimate the underlying distributions. By leveraging maximum mean discrepancies and a suitable factor model on the kernel mean embeddings of the underlying distributions, we establish consistent recovery of the underlying distributions even when data is missing not at random and positivity constraints are violated. Furthermore, we demonstrate that our nearest neighbors approach is robust to heteroscedastic noise, provided we have access to two or more measurements for the observed unit-outcome entries, a robustness not present in prior works on nearest neighbors with single measurements.