Learning Counterfactual Distributions via Kernel Nearest Neighbors
Kyuseong Choi, Jacob Feitelberg, Caleb Chin, Anish Agarwal, Raaz Dwivedi
TL;DR
The paper addresses learning counterfactual distributions for unit-outcome entries under MNAR missingness with limited samples by formulating a distributional matrix completion problem and introducing kernel-NN, a nearest-neighbors method operating on kernel mean embeddings. It combines a latent-factor model with MMD-based distances to estimate full distributions, providing instance-dependent guarantees that hold under MNAR and non-positivity, including staggered adoption and propensity-based missingness. Theoretical results show consistent distributional recovery and a distributional treatment effect ($\mathrm{iDTE}$) estimator, while experiments on simulated data and the HeartSteps mobile health study demonstrate accurate distributional imputation and favorable compared to scalar baselines. This approach offers a principled, scalable route for learning and comparing multivariate counterfactual distributions in causal panel data with complex missingness patterns.}$
Abstract
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.