# Causal Clustering for Conditional Average Treatment Effects Estimation and Subgroup Discovery

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

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 19 of its 53 references that are in the map.
- First posted on arXiv 6 Sep 2025.

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

- [Meta-learners for Estimating Heterogeneous Treatment Effects using Machine Learning](https://mlsanity.com/paper/1706.03461/) (2017)
- [Quasi-Oracle Estimation of Heterogeneous Treatment Effects](https://mlsanity.com/paper/1712.04912/) (2017)
- [Generalized Random Forests](https://mlsanity.com/paper/1610.01271/) (2016)
- [Recursive Partitioning for Heterogeneous Causal Effects](https://mlsanity.com/paper/1504.01132/) (2015)
- [Bayesian Inference of Individualized Treatment Effects using Multi-task Gaussian Processes](https://mlsanity.com/paper/1704.02801/) (2017)
- [Orthogonal Statistical Learning](https://mlsanity.com/paper/1901.09036/) (2019)
- [Orthogonal Random Forest for Causal Inference](https://mlsanity.com/paper/1806.03467/) (2018)
- [Splitting Methods for Convex Clustering](https://mlsanity.com/paper/1304.0499/) (2013)
- [Convex Clustering: Model, Theoretical Guarantee and Efficient Algorithm](https://mlsanity.com/paper/1810.02677/) (2018)
- [Validating Causal Inference Methods](https://mlsanity.com/paper/2202.04208/) (2022)
- [Local Linear Forests](https://mlsanity.com/paper/1807.11408/) (2018)
- [Causal Rule Ensemble: Interpretable Discovery and Inference of Heterogeneous Treatment Effects](https://mlsanity.com/paper/2009.09036/) (2020)
- [Robust Recursive Partitioning for Heterogeneous Treatment Effects with Uncertainty Quantification](https://mlsanity.com/paper/2006.07917/) (2020)
- [Convex Clustering through MM: An Efficient Algorithm to Perform Hierarchical Clustering](https://mlsanity.com/paper/2211.01877/) (2022)
- [Identifying Heterogeneous Treatment Effects in Multiple Outcomes using Joint Confidence Intervals](https://mlsanity.com/paper/2212.01437/) (2022)
- [Optimal Clustering with Missing Values](https://mlsanity.com/paper/1902.09694/) (2019)
- [Handling missing data in model-based clustering](https://mlsanity.com/paper/2006.02954/) (2020)
- [Model-based Clustering with Missing Not At Random Data](https://mlsanity.com/paper/2112.10425/) (2021)
- [SubgroupTE: Advancing Treatment Effect Estimation with Subgroup Identification](https://mlsanity.com/paper/2401.12369/) (2024)

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

- [From Unsupervised Subgroups to Hypothetical State-Intervention Policies: An Evaluation of Selected Subgrouping Methods in Observational Health Data](https://mlsanity.com/paper/2607.26521/) (2026)

## Abstract (arXiv)

Estimating heterogeneous treatment effects is critical in domains such as personalized medicine, resource allocation, and policy evaluation. A central challenge lies in identifying subpopulations that respond differently to interventions, thereby enabling more targeted and effective decision-making. While clustering methods are well-studied in unsupervised learning, their integration with causal inference remains limited. We propose a novel framework that clusters individuals based on estimated treatment effects using a learned kernel derived from causal forests, revealing latent subgroup structures. Our approach consists of two main steps. First, we estimate debiased Conditional Average Treatment Effects (CATEs) using orthogonalized learners via the Robinson decomposition, yielding a kernel matrix that encodes sample-level similarities in treatment responsiveness. Second, we apply kernelized clustering to this matrix to uncover distinct, treatment-sensitive subpopulations and compute cluster-level average CATEs. We present this kernelized clustering step as a form of regularization within the residual-on-residual regression framework. Through extensive experiments on semi-synthetic and real-world datasets, supported by ablation studies and exploratory analyses, we demonstrate the effectiveness of our method in capturing meaningful treatment effect heterogeneity.