Maximin Learning of Individualized Treatment Effect on Multi-Domain Outcomes

Yuying Lu; Wenbo Fei; Yuanjia Wang; Molei Liu

cs.LGMar 28, 2026

Maximin Learning of Individualized Treatment Effect on Multi-Domain Outcomes

Authors: Yuying Lu, Wenbo Fei, Yuanjia Wang, Molei Liu

Organizations: Department of Biostatistics, Columbia Mailman School of Public Health. · Department of Biostatistics, Peking University Health Science Center. · Beijing International Center for Mathematical Research, Peking University.

Abstract

Precision mental health requires treatment decisions that account for heterogeneous symptoms reflecting multiple clinical domains. However, existing methods for estimating individualized treatment effects (ITE) rely on a single summary outcome or a specific set of observed symptoms or measures, which are sensitive to symptom selection and limit generalizability to unmeasured yet clinically relevant domains. We propose DRIFT, a new maximin framework for estimating robust ITEs from high-dimensional item-level data by leveraging latent factor representations and adversarial learning. DRIFT learns latent constructs via generalized factor analysis, then constructs an anchored on-target uncertainty set that extrapolates beyond the observed measures to approximate the broader hyper-population of potential outcomes. By optimizing worst-case performance over this uncertainty set, DRIFT yields ITEs that are robust to underrepresented or unmeasured domains. We further show that DRIFT is invariant to admissible reparameterizations of the latent factors and admits a closed-form maximin solution, with theoretical guarantees for identification and convergence. In analyses of a randomized controlled trial for major depressive disorder (EMBARC), DRIFT demonstrates superior performance and improved generalizability to external multi-domain outcomes, including side effects and self-reported symptoms not used during training.

Figures & tables

Panel of Figure 1 — Figure 1: Illustration of the DRIFT Framework and Concepts.

Figure 2: Boxplots of

\mathrm{ACC}_{\min}

(left) and

\mathrm{Cor}_{\min}

(right) for the predicted individual treatment effects (ITEs) over 100 replications. In each replication,

J=30

observed responses and

T=1000

external responses are sampled from the loading hyper-population with fixed concentration

\sigma_{v}=1

, while the deviation magnitude varies as

r\in\{0.6,1,1.5\}

. The subject sample size is fixed at

N=300

Figure 3: Boxplots of

\mathrm{ACC}_{\min}

(left) and

\mathrm{Cor}_{\min}

(right) for predicted ITEs over 100 replications. We fix sample size

N=300

and scalar

r=0.6

. The

J=30

observed-response loadings are generated with concentration at

\sigma_{v}\in\{2,4,8\}

, while

T=1000

external responses are sampled from the hyper-population with

\sigma_{v}=1

	CAST self-reported symptoms				FIBSER self-rated side effects
Method	1st	2nd	3rd	4th	Frequency	Intensity	Interference
Original	0.066	0.108	0.367	0.482	0.060	0.120	0.000
Factorized GEO	0.102	0.145	0.392	0.494	0.084	0.133	0.036
Obs Maximin	0.271	0.277	0.488	0.494	0.241	0.265	0.217
DRIFT	0.343	0.361	0.476	0.482	0.319	0.307	0.307

Table 1: Prediction accuracy (ACC) on external outcomes across items. For CAST, we list the four lowest-ranked items out of 17 (

1

st to

4

th): CAST

{7,2,17,8}

for Original GEO and Factorized GEO , CAST

{7,2,17,6}

for Obs Maximin and CAST

{2,7,6,15}

for DRIFT . Notably, CAST

{7,2,17,8}

are mania-related, while CAST

{6,15}

are anxiety-related items. For FIBSER, we report the ACC for each of the three items (Frequency, Intensity, and Interference). The best ACC and any comparable results within a 0.02 margin are highlighted in bold for each column.

Panel of Figure 4 — Figure 4: Loading Analysis: Latent Domain Interpretation and Geometric Decomposition

Explore similar work

Causal Risk Minimization for High-Dimensional Treatments
May 26, 2026Nikita Dhawan, Arnav Paruthi, Andrew Kim +3Causal Effect Estimation Causal Inference
University of Toronto, Vector Institute · University of Toronto · Vanguard
Learning Optimal Distributionally Robust Individualized Treatment Rules Integrating Multi-Source Data
Mar 5, 2026Wenhai Cui, Wen Su, Xingqiu ZhaoDistribution Shift Robustness Distributionally Robust Optimization
Department of Applied Mathematics, The Hong Kong Polytechnic University · Department of Biostatistics, City University of Hong Kong
Resolving the bias-precision paradox with stochastic causal representation learning for personalized medicine
May 7, 2026Peisong Zhang, Manqiang Peng, Yuxuan Wu +21Causal Effect Estimation Causal Representation Learning
School of Medicine, NUS, . · College of Design and Engineering, NUS, . · Aier Academy of Ophthalmology, Central South University. +14

May 26, 2026J/K move · Enter open · S save

Causal Risk Minimization for High-Dimensional Treatments

Nikita Dhawan, Arnav Paruthi, Andrew Kim, Lovedeep Gondara+2

University of Toronto, Vector Institute · University of Toronto · Vanguard

Causal Effect Estimation Causal Inference

Open paper AskPDF arXiv

Predicting the effect of interventions with many possible variations, e.g., therapeutic content that affects mental health outcomes or an earnings call transcript that drives movement in share price, is useful across several domains. However, classical causal estimators tend to assume that all possible interventions are observed, which is infeasible when interventions vary widely, for instance, in the space of all text strings. We adapt a well-known approach of recasting causal inference as a learning problem, to address high-dimensional treatment spaces. Specifically, under standard assumptions like no unobserved confounding, we show that causal error decomposes into a series of moment-balancing errors of increasing order, and design objectives that directly improve causal estimation. We also show how to project the effect of a high-dimensional treatment onto lower-dimensional treatment attributes, which allows a single model to answer several causal questions without additional attribute-specific training. We empirically evaluate our estimators in settings with high-dimensional continuous, discrete, and text treatments, the last of which used a semi-synthetic dataset of Amazon Reviews. Our experiments demonstrate the benefit of higher-order balance error optimization and competitive performance of projected causal estimates with attribute-specific estimators.

Abstract

Explore similar work

Causal Risk Minimization for High-Dimensional Treatments

Learning Optimal Distributionally Robust Individualized Treatment Rules Integrating Multi-Source Data

Resolving the bias-precision paradox with stochastic causal representation learning for personalized medicine