Bayesian Supervised Causal Clustering

Luwei Wang; Nazir Lone; Sohan Seth

stat.MLMar 5, 2026

Bayesian Supervised Causal Clustering

Authors: Luwei Wang, Nazir Lone, Sohan Seth

Organizations: School of Informatics, University of Edinburgh, Edinburgh, UK · Usher Institute, University of Edinburgh, Edinburgh, UK

Abstract

Finding patient subgroups with similar characteristics is crucial for personalized decision-making in various disciplines such as healthcare and policy evaluation. While most existing approaches rely on unsupervised clustering methods, there is a growing trend toward using supervised clustering methods that identify operationalizable subgroups in the context of a specific outcome of interest. We propose Bayesian Supervised Causal Clustering (BSCC), with treatment effect as outcome to guide the clustering process. BSCC identifies homogenous subgroups of individuals who are similar in their covariate profiles as well as their treatment effects. We evaluate BSCC on simulated datasets as well as real-world dataset from the third International Stroke Trial to assess the practical usefulness of the framework.

Figures & tables

Figure 1: (a) Illustrating the conceptual difference between (left) unsupervised clustering with Gaussian Mixture Model ( gmm ) (right) supervised clustering with bscc . gmm fails to recover the subgroups relevant to treatment while bscc can. (b) Difference between bscc clustering samples based on both covariates and treatment effect in a supervised clustering set-up, as opposed to causal clustering or supervised learning based approaches clustering based on potential outcomes or treatment effect only.

Category	Supervised	Probabilistic	Covariate-based	TE	Example
UC	$\times$	$\checkmark$	$\checkmark$	$\times$	gmm
SC	$\checkmark$	$\checkmark$	$\checkmark$	$\times$	sgmm
SA	$\checkmark$	$\times$	$\checkmark$	$\checkmark$	it , mob
SA	$\checkmark$	$\checkmark$	$\checkmark$	$\checkmark$	bscc (ours)
EM	$\checkmark$	$\times$	$\times$	$\checkmark$	cf , bart , DR-l , R-l , FindIt , owe , cc
BM	$\checkmark$	$\checkmark$	$\times$	$\checkmark$	Xu et al. [2016] , Shahn and Madigan [2017]

Table 1: Summary of approaches. TE : treatment-effect-based.

Panel of Figure 2 — Figure 2: Plate Diagram of bscc with feature selection.

Model (Treat Prop.)	ARI	SATE	PEHE
gmm (0.5)	0.768	[ $-0.40$ , $1.32$ ]	3.13
sgmm (0.5)	0.625	[ $-5.54$ , $0.67$ ]	2.65
it (0.5)	0.574	[ $-0.03$ , $-0.03$ ]	3.09
mob (0.5)	0.437	[ $-4.73$ , $2.38$ ]	2.58
bscc (0.2)	0.715	[ $-4.99$ , $3.98$ ]	1.49
bscc (0.5)	0.721	[ $-5.13$ , $3.99$ ]	1.45

Table 2: Summary of performance for simulation hte .

Figure 3: Visualization of UMAP Projection of the Simulated Dataset in simulation hte : Cluster assignments are given by (a) ground truth of constant treatment effects $(0.5,5,-5,0,0)$ , (b) gmm , (c) sgmm , (d) mob , (e) bscc , (f) cf , (g) bart , (h) DR-learner , (i) R-learner and (j) cc .

Figure 4: bscc on IST-3 Dataset: (a) The estimated cluster means and (b) corresponding ORs with $95\%$ CIs for training (in log scale), and (c) the empirical cluster means and (d) corresponding ORs with $95\%$ CIs for test set. Each axis representing a continuous covariate in the radar chart is transformed back to its original scale to enhance interpretability. (e) Cluster-specific feature importance.

Group based on	Task	Example	Covariate Split	Supervision
Clustering covariates	Unsupervised	gmm , lca	Multivariate	None
	Supervised	BayesProfile, sgmm	Multivariate	Outcome
	Supervised	it , mob	One-at-a-time	Treatment effect
	Supervised	bscc	Multivariate	Treatment effect
Clustering treatment effects	Supervised	Meta-learners, cf , bart , FindIt , owe	None	Treatment effect
Clustering potential outcomes	Supervised	Causal Clustering	None	Potential outcomes

Table A.1: Summary of Related Works

ID	Gender	Age	Treated	Control	Treated
				Outcome	Outcome
1	F	61	Yes	NA	2
2	M	35	No	1	NA
3	F	46	Yes	NA	1

Table A.2 : Example Tabular Data for HTE

Size	Data Case	Avg. Rel.	ARI
	(Treated Prop.)	$L_{2}$ Error
$N=720$	L-C (0.5)	0.04	1
$N=740$	L-C (0.2)	0.033	0.996
$N=760$	L-L (0.5)	0.072	1
$N=780$	L-L (0.2)	0.054	0.996

Table B.3 : A Summary of Synthetic Datasets and Model Performance for Sanity Check of bscc : The data case of L-C (i.e. linear-constant) means the data is generated from a linear control outcome function and constant treatment effects for each cluster.

Panel of Figure 10 — Figure B.1 : simulation fs : (a) Simulated dataset and (b) learned cluster-specific feature importance by bscc .

Figure B.2 : True potential outcomes $(y^{0},y^{1})$ for simulation hte , colored by ground-truth cluster membership. The dashed line indicates $y^{1}=y^{0}$ (zero treatment effect).

Panel of Figure 12 — Figure B.3 : UMAP projections of covariate space for simulation hte , colored by cluster assignments from (a) it , (b) FindIt , and (c) owe .

$\beta_{k}$ SD	ARI	Range of $\tau_{k}$	PEHE
0.5	0.750	[-4.810 4.984]	0.847
1	0.733	[-4.766, 5.020]	0.943
5	0.733	[-4.750, 4.927]	0.944
10	0.732	[-4.841, 5.059]	0.983
100	0.739	[-4.812, 5.006]	0.987

Table B.4 : Sensitivity Test of Prior SD of constant

\beta_{k}

on simulation hte .

$\beta_{k}$ Mean	ARI	Range of $\tau_{k}$	PEHE
(0,0,0,0)	0.733	[-4.766, 5.020]	0.943
(0,0,-5,5)	0.751	[-4.790, 5.029]	0.885

Table B.5 : Sensitivity Test of Prior Mean of constant

\beta_{k}

on simulation hte .

K	Train RMSE	Test RMSE	Train PEHE	Test PEHE	Train Range	Test Range
	of Control	of Control			of $\tau_{k}$	of SATE
1	0.217	0.460	0.048	0.048	[0.048, 0.048]	[-0.057,-0.057]
2	0.210	0.460	0.386	0.382	[-0.347, 0.429]	[-0.172, -0.025]
3	0.201	0.529	0.386	0.268	[-0.550, 0.733]	[-6.29, -0.048]
4	0.203	0.524	0.368	0.223	[-0.694, 1.232]	[-0.138,0.029]
5	0.216	0.472	0.361	0.189	[-0.748, 1.240]	[-0.0224,-0.0215]

Table B.6 : Test of Homogeneous and Null Treatment Effects on simulation null . We have only one cluster of individuals with zero treatment effects.

Figure C.4 : IST-3 Data: Oxford Handicap Score (OHS) at 6 months outcome distributions for control and treated groups.

Description of Covariate	Feature Type
Age	Continuous
Delay Time from onset of stroke symptoms to randomisation ( $\mathrm{h}$ )	Continuous
Systolic blood pressure ( $\mathrm{mm}$ )	Continuous
Diastolic blood pressure ( $\mathrm{mm}$ )	Continuous
Glucose ( $\mathrm{mmol}$ )	Continuous
NIHSS (Stroke Score at randomisation)	Continuous

Table C.7: List of Covariates Used for the IST-3 Dataset.

No. Clusters	Test Acc. of	Test Policy Risk
	Control Outcome
$K=2$	0.754	0.648
$K=3$	0.770	0.644
$K=4$	0.762	0.644

Table C.8 : The Evaluation Metrics for Selecting the Number of Clusters of the IST-3 Dataset.

Model Case	Train Range of SATE	Test Range of SATE	Test Policy Risk
bscc	[-0.273, 0.658]	[-0.243, 0.501]	0.644
cf	[-1.942, 2.079]	[-0.249, 0.172]	0.616
bart	[-0.058, 2.521]	[-0.204, 0.636]	0.632
DR-learner	[-3.563, 4.348]	[-0.165, 0.071]	0.621
R-learner	[-1.941, 2.603]	[-0.270, 0.096]	0.618
mob	[-0.621, 0.416]	[-0.724, 0.355]	0.638

Table C.9 : A Summary of Model Performance for Different Supervised Learning Models on IST-3 Data.

	Max	Min	Popu	Cluster1	Cluster2	Cluster3
Age	89.84	66.08	78.07	73.42	81.68	80.54
Delay2Rand	4.37	3.14	3.78	3.98	3.49	3.68
SBloodPres	171.79	139.70	155.27	155.22	155.38	156.17
DBloodPres	92.70	71.42	81.61	84.28	79.36	82.03
Glucose	8.06	6.30	7.29	7.32	7.32	7.00
NIHSS	19.50	7.31	12.25	8.12	17.73	12.22

Table C.10 : IST-3: Summary Statistics (Min/Max/Mean) of Covariates by Clusters for Training Set ( bscc ).

	Max	Min	Popu	Cluster1	Cluster2	Cluster3
Age	88.48	67.16	77.32	74.62	80.44	78.47
Delay2Rand	4.52	3.34	3.93	4.10	3.71	3.88
SBloodPres	171.26	138.54	155.09	155.69	153.93	155.61
DBloodPres	91.99	72.21	82.26	83.63	80.23	82.42
Glucose	8.07	6.30	7.23	7.25	7.34	7.00
NIHSS	18.99	7.26	12.12	8.06	17.26	13.08

Table C.11 : IST-3: Summary Statistics (Min/Max/Mean) of Covariates by Clusters for Test Set ( bscc ).

Cluster	Group	Mortality	Alive	Mortality Prop.
Cluster1	Treated	53	265	16.7%
Cluster1	Control	43	280	13.3%
Cluster2	Treated	92	133	40.9%
Cluster2	Control	100	110	47.6%
Cluster3	Treated	49	86	36.3%
Cluster3	Control	36	122	22.8%

Table C.12 : IST-3: Contingency Table of Clusters for Training Set ( bscc ).

Cluster	Group	Mortality	Alive	Mortality Prop.
Cluster1	Treated	45	287	13.6%
Cluster1	Control	51	255	16.7%
Cluster2	Treated	74	131	36.1%
Cluster2	Control	102	143	41.6%
Cluster3	Treated	49	100	32.9%
Cluster3	Control	30	101	22.9%

Table C.13 : IST-3: Contingency Table of Clusters for Test Set ( bscc ).

Cluster	Group	Mortality	Alive	Mortality Prop.
Cluster1	Treated	244	276	46.9%
Cluster1	Control	248	277	47.2%
Cluster2	Treated	81	488	14.2%
Cluster2	Control	81	489	14.2%
Cluster3	Treated	37	238	13.5%
Cluster3	Control	33	245	11.9%

Table C.14 : IST-3: Contingency Table of Clusters for gmm .

Figure C.5 : gmm on IST-3 Dataset: (a) The estimated cluster means and (b) ORs with $95\%$ CIs (log scale) when $K=3$ .

Panel of Figure 26 — Figure C.6 : cf on IST-3 Dataset.

Panel of Figure 27 — Figure C.7 : bart on IST-3 Dataset.

Panel of Figure 28 — Figure C.8 : DR-learner on IST-3 Dataset.

Panel of Figure 29 — Figure C.9 : R-learner on IST-3 Dataset.

Panel of Figure 30 — Figure C.10 : mob on IST-3 Dataset.

Clinical Decision-Making

Clinical decision-making studies how evidence, patient preferences, risks, and resource constraints can guide diagnosis, treatment, and care.

See all 128 papers

From Unsupervised Subgroups to Hypothetical State-Intervention Policies: An Evaluation of Selected Subgrouping Methods in Observational Health Data
Jul 29, 2026Vasundhara Acharya, Bulent YenerClinical Decision-Making
Rensselaer Polytechnic Institute, Troy, USA
CaSPECT: Discovering Causally Homogeneous Subgroups via Directed Spectral Clustering
Jul 3, 2026Arghya Pratihar, Shinjon Chakraborty, Swagatam DasCausal Effect Estimation Causal Structure Learning
1*Electronics and Communication Sciences Unit, Indian Statistical Institute, India. · 2Indian Statistical Institute, India.
Joint Causal Structure and Cluster Discovery Using Variational Inference
Aug 23, 2026Avni Rajpal, Anubhav Kumar, Rishabh Karnad +2Causal Discovery Causal Structure Learning
Department of Artificial Intelligence, Indian Institute of Technology Hyderabad, India · RIKEN Center for AI Project, Tokyo, Japan

Jul 29, 2026J/K move · Enter open · S save

From Unsupervised Subgroups to Hypothetical State-Intervention Policies: An Evaluation of Selected Subgrouping Methods in Observational Health Data

Vasundhara Acharya, Bulent Yener

Rensselaer Polytechnic Institute, Troy, USA

Clinical Decision-Making

Open paper AskPDF arXiv

Conventional subgroup analyses can yield unstable and difficult-to-interpret conclusions, especially in observational biomedical data where each individual is observed under only one exposure state, true individual treatment effects are unavailable, and causal structure is uncertain. We investigate whether subgroups constructed from pretreatment characteristics, without using exposure, outcome, or estimated treatment-effect information, can serve as interpretable units for budget-constrained policy prioritization. We propose a framework combining causal-discovery-informed covariate selection, discovery-evaluation sample splitting, inductive unsupervised clustering, uncertainty-aware subgroup selection, and held-out doubly robust policy evaluation. We compare K-means, hard, membership-weighted, and stochastic Fuzzy C-means, Bayesian Gaussian mixture models, and a supervised causal-forest-derived CATE-tree comparator. Policies are evaluated under a 70% budget for hypothetical obesity-to-non-obesity and elevated-to-lower-glucose state shifts in the PIMA Indians Diabetes dataset and for a lifetime-smoking-history contrast in NHANES. The highest estimated ungated utilities were 0.799 for the BMI policy using Bayesian GMM, 0.735 for the glucose policy using hard or membership-weighted FCM, and 0.775 for the smoking-history policy using K-means. All paired 95% confidence intervals for policy-risk differences included zero, and no comparison remained statistically significant after Holm adjustment. Bayesian pooling generally preserved ungated allocations, whereas Empirical Bernstein gating was more conservative. Policies with similar estimated utility could nevertheless prioritize different individuals. The findings should be interpreted as assumption-dependent decision-support evidence for hypothetical state contrasts rather than proof of intervention benefit.

Abstract

Clinical Decision-Making

Explore similar work

From Unsupervised Subgroups to Hypothetical State-Intervention Policies: An Evaluation of Selected Subgrouping Methods in Observational Health Data

CaSPECT: Discovering Causally Homogeneous Subgroups via Directed Spectral Clustering

Joint Causal Structure and Cluster Discovery Using Variational Inference