Neural Supervised Clustering for Learning Individualized Treatment Rules with Many Treatments
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.
Related
- Supervised Learning Objectives for High-Dimensional Treatment Effects
- Online Individualized Treatment Selection under Partial Feedback: False Discovery Control and Minimax Optimality
- Nash Equilibrium Individualized Treatment Regimes under Network Interference
- Quantifying Precision in Measurement of Individualized, Evidence-Based Learning Across Sessions in Applied Behavior Analysis Therapy
- A Hybrid Spatial Statistical Learning Framework for Individualized Probability Estimation: Application to Multisite Autism Neuroimaging
BibTeX
@misc{huang2026neural,
title = {{Neural Supervised Clustering for Learning Individualized Treatment Rules with Many Treatments}},
author = {Daixi Huang and Yixin Yuan and Yifan Dai and Yufeng Liu},
howpublished = {arXiv},
year = {2026},
month = oct,
doi = {10.48550/arxiv.2610.19104},
eprint = {2610.19104},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2610.19104},
note = {CC BY-NC-SA 4.0},
}
RIS
TY - UNPB
TI - Neural Supervised Clustering for Learning Individualized Treatment Rules with Many Treatments
AU - Daixi Huang
AU - Yixin Yuan
AU - Yifan Dai
AU - Yufeng Liu
PY - 2026
DA - 2026/10/13
PB - arXiv
DO - 10.48550/arxiv.2610.19104
UR - https://arxiv.org/abs/2610.19104
AB - 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.
ER -