Learning Counterfactual Distributions via Kernel Nearest Neighbors

TL;DR

This paper introduces a kernel-based nearest neighbors method for estimating multivariate distributions across units and outcomes, effectively handling missing not at random data and unobserved confounding in distributional matrix completion.

Contribution

It proposes a novel distributional matrix completion framework using kernel mean embeddings and maximum mean discrepancies, ensuring consistent recovery under challenging data conditions.

Findings

Consistent distribution recovery despite missing not at random data.

Robustness to heteroscedastic noise with multiple measurements.

Effective estimation in the presence of unobserved confounding.

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.