Causal K-Means Clustering

TL;DR

This paper introduces Causal k-Means Clustering, a novel method to identify unknown subgroup structures in causal effect studies using clustering algorithms tailored for counterfactual functions.

Contribution

It develops a simple plug-in estimator and a bias-corrected estimator based on nonparametric efficiency theory, enabling effective subgroup detection in causal analysis.

Findings

The plug-in estimator is easy to implement with off-the-shelf algorithms.

The bias-corrected estimator achieves fast root-n convergence and asymptotic normality.

Simulations and a real study demonstrate the method's practical utility.

Abstract

Causal effects are often characterized with population summaries. These might provide an incomplete picture when there are heterogeneous treatment effects across subgroups. Since the subgroup structure is typically unknown, it is more challenging to identify and evaluate subgroup effects than population effects. We propose a new solution to this problem: \emph{Causal k-Means Clustering}, which leverages the k-means clustering algorithm to uncover the unknown subgroup structure. Our problem differs significantly from the conventional clustering setup since the variables to be clustered are unknown counterfactual functions. We present a plug-in estimator which is simple and readily implementable using off-the-shelf algorithms, and study its rate of convergence. We also develop a new bias-corrected estimator based on nonparametric efficiency theory and double machine learning, and show that this estimator achieves fast root-n rates and asymptotic normality in large nonparametric models. Our proposed methods are especially useful for modern outcome-wide studies with multiple treatment levels. Further, our framework is extensible to clustering with generic pseudo-outcomes, such as partially observed outcomes or otherwise unknown functions. Finally, we explore finite sample properties via simulation, and illustrate the proposed methods using a study of mobile-supported self-management for chronic low back pain.