Neural Supervised Clustering for Learning Individualized Treatment Rules with Many Treatments

Daixi Huang; Yixin Yuan; Yifan Dai; Yufeng Liu

stat.MLOct 13, 2026

Neural Supervised Clustering for Learning Individualized Treatment Rules with Many Treatments

Authors: Daixi Huang, Yixin Yuan, Yifan Dai, Yufeng Liu

Organizations: University of Michigan · Yale University · University of North Carolina at Chapel Hill

Abstract

Learning an individualized treatment rule (ITR) with many treatments requires estimating treatment-specific effects from limited observations per arm. Treatments with similar conditional responses can share information, but linear contrast models may miss nonlinear heterogeneity. We propose nonlinear supervised clustering fusion (NSCF), a neural Q-learning framework with a shared nonlinear covariate representation and low-dimensional treatment embeddings. Pairwise fusion groups treatments with similar learned contrasts, yielding an ITR and an interpretable partition. An alternating algorithm combines neural training, embedding fusion, validation-based partition selection, and unpenalized grouped refitting. NSCF lowers regret relative to a strong neural reference in nonlinear exact-sharing simulations; perturbations characterize the limits of hard sharing. In a melanoma xenograft application, NSCF improves average policy value over the neural reference. These results support nonlinear representation learning with supervised treatment fusion when treatment effects exhibit shared structure.

Figures & tables

Figure 1: Exact sharing: eight methods and 50 paired seeds per scenario,

K=12

, two groups, training/validation sizes

1000/500

, separation 1.5. Diamonds denote means, box lines medians, whiskers 1.5 IQR, and circles observations beyond the whiskers; no seeds are excluded. Vertical scales differ across panels.

Figure 2: Sample-size comparison in nonlinear-hard:

K=12

, two groups, separation 1.0, validation size

n/2

. All eight methods are shown. The

n=200,600,1000

panels use 50 paired seeds;

n=400

uses 20. Vertical scales differ, with all outliers retained. Boxplot conventions follow Figure 1 .

Figure 3: Treatment-group complexity in nonlinear-hard, separation 1.5:

(G^{*},K,n)=(2,12,1000)

(3,12,1500)

, and

(5,15,3000)

. Each panel includes eight methods and 50 paired seeds. Sample size, architecture and sometimes

K

vary, so this is not a one-factor group-count ablation.

Figure 4: Recorded fusion paths for fixed seed 200, not selected by performance. Branch colors and bottom bars show selected groups. The dashed display cut lies between successive merge heights and reproduces the saved partition; it is not the selected

\lambda^{*}

, which is printed below each panel. Heights are

\log(1+1000\lambda)

. Exact recovery is 50/50, 49/50, and 12/50; the five-group example is not typical of recovery frequency.

Figure 5: Exact and very mild sharing: nonlinear-hard,

\epsilon=0

and 0.005,

K=12

, two prototypes, separation 1.5, and a

1000/500/10000

split. All eight methods and 50 paired seeds are retained. The exploratory 0.005 comparison gives a small mean advantage, with an interval including zero.

Figure 6: Melanoma PDX application. (a) Held-out reward (

-\texttt{RespScaled}

; higher is better): ten complete five-fold repeats of the same 32 models; diamonds denote means. (b) Recorded fusion path for fixed case006, selected cycle 4. Heights are evaluated

\lambda

values, not drug distances; tied merges imply no internal order. The dashed line marks the selected single residual-response group, not the intermediate four groups. Treatment means remain unfused; enc. denotes encorafenib. This is not a cross-fold consensus or drug-mechanism hierarchy.

Individualized Treatment Effect Estimation

Individualized treatment effect estimation predicts how a treatment’s effects vary across people to support personalized decisions.

See all 44 papers

Bayesian Supervised Causal Clustering
Mar 5, 2026Luwei Wang, Nazir Lone, Sohan SethClinical Decision-Making
School of Informatics, University of Edinburgh, Edinburgh, UK · Usher Institute, University of Edinburgh, Edinburgh, UK
Cohort Organized Learning: Clustering Through Agreement
Jun 19, 2026Finn Henry O'Shea, Maria Elena MonzaniUnsupervised Learning Clustering
SLAC National Accelerator Laboratory, 2575 Sand Hill Road, Menlo Park, CA 94025-7015, USA · Kavli Institute for Particle Astrophysics and Cosmology, Stanford University, Stanford, CA 94305-4085 USA
Algorithmically Aligned Neural Agglomerative Tree Construction
Oct 5, 2026Robert R Nerem, Pranav Singh, Cheyenne Ward +1Clustering Hierarchical Clustering
University of California San Diego · University of Washington

Mar 5, 2026J/K move · Enter open · S save

Bayesian Supervised Causal Clustering

Luwei Wang, Nazir Lone, Sohan Seth

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

Clinical Decision-Making

Open paper AskPDF arXiv

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.

Abstract

Individualized Treatment Effect Estimation

Explore similar work

Bayesian Supervised Causal Clustering

Cohort Organized Learning: Clustering Through Agreement

Algorithmically Aligned Neural Agglomerative Tree Construction