Locally Interpretable Individualized Treatment Rules for Black-Box Decision Models

Abstract: Individualized treatment rules (ITRs) aim to optimize healthcare by tailoring treatment decisions to patient-specific characteristics. Existing methods typically rely on either interpretable but inflexible models or highly flexible black-box approaches that sacrifice interpretability; moreover, most impose a single global decision rule across patients. We introduce the Locally Interpretable Individualized Treatment Rule (LI-ITR) method, which combines flexible machine learning models to accurately learn complex treatment outcomes with locally interpretable approximations to construct subject-specific treatment rules. LI-ITR employs variational autoencoders to generate realistic local synthetic samples and learns individualized decision rules through a mixture of interpretable experts. Simulation studies show that LI-ITR accurately recovers true subject-specific local coefficients and optimal treatment strategies. An application to precision side-effect management in breast cancer illustrates the necessity of flexible predictive modeling and highlights the practical utility of LI-ITR in estimating optimal treatment rules while providing transparent, clinically interpretable explanations.

• The paper introduces LI-ITR, a two-stage framework that combines VAE-based correlated perturbations with sparse mixture-of-experts models to produce subject-specific, interpretable treatment rules from black-box predictors.
• In simulations, LI-ITR kept treatment-effect coefficient bias below 0.10 and exceeded 99% optimal-treatment accuracy, outperforming LIME, OWL, and Q-learning while remaining robust to moderate model misspecification.
• In a breast cancer hepatotoxicity study, LI-ITR achieved the best policy value (−2.922), reduced local prediction error to 0.17 versus LIME’s 0.54, and revealed clinically meaningful differences in treatment drivers across patients.

Overview and motivation

Individualized treatment rules (ITRs) derived from flexible machine learning models typically face a trade-off: indirect methods such as Q-learning, A-learning, and G-estimation permit counterfactual reasoning but are sensitive to outcome-model misspecification, while direct methods such as outcome-weighted learning (OWL) are more robust but higher-variance. In both cases, the resulting decision rule is usually a single global rule applied uniformly across patients, even though treatment effects may vary in direction and magnitude across subpopulations — a phenomenon illustrated by the U-shaped relationship between CYP2D6 metabolizer status and breast cancer–specific mortality under tamoxifen (HR = 2.59 for poor metabolizers, HR = 4.52 for ultrarapid metabolizers) (2602.11520).

The paper introduces LI-ITR (Locally Interpretable Individualized Treatment Rule), a two-stage framework that (i) uses a trained black-box model to capture complex treatment-response structure, and (ii) constructs subject-specific, locally interpretable surrogate models from realistic perturbed synthetic samples. The approach addresses two known deficiencies of LIME-style local explanation: independent feature perturbation that destroys feature correlations, and single-surrogate local linear fits that fail when the neighborhood contains heterogeneous or discontinuous response structure (2602.11520).

Methodology

LI-ITR proceeds in four steps: train a flexible black-box outcome model; generate perturbed synthetic samples around the subject of interest; obtain black-box predictions on those samples; and fit an interpretable local model to recover subject-specific coefficients.

VAE-based perturbation. Rather than perturbing correlated features independently, the method trains a $\beta$-VAE on the covariates, maximizing the modified ELBO

$$\mathbb{E}_{q_{\boldsymbol{\phi}}(\mathbf{Z}\mid\mathbf{X})}[\log p_{\boldsymbol{\gamma}}(\mathbf{X}\mid\mathbf{Z})] - \beta \cdot \mathrm{KL}(q_{\boldsymbol{\phi}}(\mathbf{Z}\mid\mathbf{X}) \,\|\, p(\mathbf{Z})),$$

and then perturbs in latent space via $$\mathbf{Z}' = \mathbf{Z} + \min(1,\alpha)\odot\boldsymbol{\epsilon}'$$, $$\boldsymbol{\epsilon}'\sim\mathcal{N}(\mathbf{0},\mathbf{I})$$, decoding to synthetic feature vectors. Treatment is not passed through the VAE; instead, under the positivity assumption, $T'$ is drawn from a Bernoulli(0.5), reflecting the design's reliance on exchangeability of treatment within the local neighborhood.

Mixture-of-experts local modeling. Because a dispersed perturbed neighborhood may contain heterogeneous treatment-response regimes, the black-box predictions $\hat{Y}$ are modeled as a mixture of $K$ components, each a linear expert of the form

$$\mu_k(\mathbf{D}',\boldsymbol{\beta}_k) = \boldsymbol{\beta}_{k1}^{\mathrm{T}}H_0 + (\boldsymbol{\beta}_{k2}^{\mathrm{T}}H_1)\,T',$$

separating prognostic variables ($H_0$, which need not be interpretable) from prescriptive/tailoring variables ($H_1$). A neural-network gating network produces responsibilities $$\mathbb{E}_{q_{\boldsymbol{\phi}}(\mathbf{Z}\mid\mathbf{X})}[\log p_{\boldsymbol{\gamma}}(\mathbf{X}\mid\mathbf{Z})] - \beta \cdot \mathrm{KL}(q_{\boldsymbol{\phi}}(\mathbf{Z}\mid\mathbf{X}) \,\|\, p(\mathbf{Z})),$$0 via softmax over learned scores $$\mathbb{E}_{q_{\boldsymbol{\phi}}(\mathbf{Z}\mid\mathbf{X})}[\log p_{\boldsymbol{\gamma}}(\mathbf{X}\mid\mathbf{Z})] - \beta \cdot \mathrm{KL}(q_{\boldsymbol{\phi}}(\mathbf{Z}\mid\mathbf{X}) \,\|\, p(\mathbf{Z})),$$1, and all parameters are optimized jointly with AdamW. Two mechanisms enforce interpretable, stable assignments: a negative Shannon entropy penalty ($$\mathbb{E}_{q_{\boldsymbol{\phi}}(\mathbf{Z}\mid\mathbf{X})}[\log p_{\boldsymbol{\gamma}}(\mathbf{X}\mid\mathbf{Z})] - \beta \cdot \mathrm{KL}(q_{\boldsymbol{\phi}}(\mathbf{Z}\mid\mathbf{X}) \,\|\, p(\mathbf{Z})),$$2) encouraging sparse expert selection, and a hard argmax selection implemented through the Straight-Through Estimator so gradients flow through the softmax during backpropagation. The optimal treatment for a subject is read off directly from the selected expert's estimated coefficients.

A notable design choice is that the prognostic component may itself be left unspecified or modeled flexibly; interpretability is required only for the treatment-effect component. This is consistent with the paper's framing that only prescriptive coefficients enter the decision rule.

Simulation evidence

Simulations generate four correlated baseline covariates from a one-dimensional latent variable, treatment assignment via a logistic propensity depending on $$\mathbb{E}_{q_{\boldsymbol{\phi}}(\mathbf{Z}\mid\mathbf{X})}[\log p_{\boldsymbol{\gamma}}(\mathbf{X}\mid\mathbf{Z})] - \beta \cdot \mathrm{KL}(q_{\boldsymbol{\phi}}(\mathbf{Z}\mid\mathbf{X}) \,\|\, p(\mathbf{Z})),$$3, and a piecewise outcome in which main-effect coefficients are constant but treatment-interaction coefficients $$\mathbb{E}_{q_{\boldsymbol{\phi}}(\mathbf{Z}\mid\mathbf{X})}[\log p_{\boldsymbol{\gamma}}(\mathbf{X}\mid\mathbf{Z})] - \beta \cdot \mathrm{KL}(q_{\boldsymbol{\phi}}(\mathbf{Z}\mid\mathbf{X}) \,\|\, p(\mathbf{Z})),$$4 vary across regions defined by median splits of $$\mathbb{E}_{q_{\boldsymbol{\phi}}(\mathbf{Z}\mid\mathbf{X})}[\log p_{\boldsymbol{\gamma}}(\mathbf{X}\mid\mathbf{Z})] - \beta \cdot \mathrm{KL}(q_{\boldsymbol{\phi}}(\mathbf{Z}\mid\mathbf{X}) \,\|\, p(\mathbf{Z})),$$5 and $$\mathbb{E}_{q_{\boldsymbol{\phi}}(\mathbf{Z}\mid\mathbf{X})}[\log p_{\boldsymbol{\gamma}}(\mathbf{X}\mid\mathbf{Z})] - \beta \cdot \mathrm{KL}(q_{\boldsymbol{\phi}}(\mathbf{Z}\mid\mathbf{X}) \,\|\, p(\mathbf{Z})),$$6. Training sizes span $$\mathbb{E}_{q_{\boldsymbol{\phi}}(\mathbf{Z}\mid\mathbf{X})}[\log p_{\boldsymbol{\gamma}}(\mathbf{X}\mid\mathbf{Z})] - \beta \cdot \mathrm{KL}(q_{\boldsymbol{\phi}}(\mathbf{Z}\mid\mathbf{X}) \,\|\, p(\mathbf{Z})),$$7 to $$\mathbb{E}_{q_{\boldsymbol{\phi}}(\mathbf{Z}\mid\mathbf{X})}[\log p_{\boldsymbol{\gamma}}(\mathbf{X}\mid\mathbf{Z})] - \beta \cdot \mathrm{KL}(q_{\boldsymbol{\phi}}(\mathbf{Z}\mid\mathbf{X}) \,\|\, p(\mathbf{Z})),$$8, with synthetic sample sizes $$\mathbb{E}_{q_{\boldsymbol{\phi}}(\mathbf{Z}\mid\mathbf{X})}[\log p_{\boldsymbol{\gamma}}(\mathbf{X}\mid\mathbf{Z})] - \beta \cdot \mathrm{KL}(q_{\boldsymbol{\phi}}(\mathbf{Z}\mid\mathbf{X}) \,\|\, p(\mathbf{Z})),$$9 and $$\mathbf{Z}' = \mathbf{Z} + \min(1,\alpha)\odot\boldsymbol{\epsilon}'$$0.

Two results stand out. First, for prognostic coefficients, LI-ITR and LIME both achieve low bias. Second, for treatment-related coefficients — precisely those governing the ITR — LIME exhibits large bias (e.g., bias up to 1.47 for $$\mathbf{Z}' = \mathbf{Z} + \min(1,\alpha)\odot\boldsymbol{\epsilon}'$$1) that does not diminish with larger training or synthetic samples, whereas LI-ITR keeps bias below 0.10 in all settings with decreasing variability as sample sizes grow. This indicates that LIME's failure is structural (single-surrogate approximation over a heterogeneous neighborhood) rather than a sample-size artifact.

For optimal-treatment recovery on held-out data, LI-ITR achieves accuracy above 99% across all scenarios, essentially matching the black-box ITR itself, compared with roughly 88% for LIME and about 83% for OWL and Q-learning; causal forest performs well (97–99%) but offers no subject-level interpretable rule. Under deliberate misspecification of the prescriptive component (a quadratic term added to the data-generating process), LI-ITR remains above 0.98 accuracy, indicating robustness to moderate misspecification of the local expert form. The authors note this robustness is conditional: validity requires a sufficiently small neighborhood and a black-box response surface well approximated by a first-order Taylor expansion locally.

Application to endocrine therapy hepatotoxicity

The real-data application concerns precision side-effect management: selecting between tamoxifen and aromatase inhibitor (AI) therapy to minimize AST elevation among 1,819 women with hormone-receptor-positive breast cancer treated between 2016 and 2023 (445 tamoxifen, 1,374 AI). The endpoint is change in AST from pre-treatment (within eight weeks before initiation) to post-treatment (28–84 days after), with substantial inter-individual variability (mean change −1.0 U/L, SD 22.6).

A neural network predicting AST change achieves test $$\mathbf{Z}' = \mathbf{Z} + \min(1,\alpha)\odot\boldsymbol{\epsilon}'$$2 versus 0.78 for a linear model, supporting the need for flexible predictive modeling. A VAE trained on baseline laboratories and demographics yields latent variables with markedly weaker correlation structure than the original features, enabling more coherent local perturbations than direct feature-space perturbation. Local models were fit for 454 of 455 test subjects (one excluded due to extreme values producing out-of-range perturbations), with $$\mathbf{Z}' = \mathbf{Z} + \min(1,\alpha)\odot\boldsymbol{\epsilon}'$$3 regularization tuned by repeated cross-validation to control multicollinearity.

Policy evaluation via an inverse-probability-weighted value function (propensity scores from a probability forest) shows LI-ITR achieving the most favorable value (−2.922), ahead of the black-box ITR (−2.800), causal forest (−2.733), LIME (−2.413), OWL (−2.403), and Q-learning (−2.196). LI-ITR recommends AI for only 6% of patients versus 19% under Q-learning, while causal forest prescribes tamoxifen universally. Fidelity to the black-box predictor is substantially better than LIME: mean absolute prediction difference of 0.17 (SD 0.24) versus 0.54 (SD 0.61) for LIME — more than a threefold improvement attributable to the VAE-based perturbation plus gating-driven local learning.

Subject-level inspection confirms genuine heterogeneity in local structure. Two illustrative subjects both achieve high local fidelity (local $$\mathbf{Z}' = \mathbf{Z} + \min(1,\alpha)\odot\boldsymbol{\epsilon}'$$4 of 0.96 and 0.97) yet exhibit markedly different coefficient vectors: Subject 10 shows a large positive tamoxifen main effect dominated by laboratory interactions (strongly negative bilirubin coefficient, positive creatinine coefficient), whereas Subject 100 has a smaller tamoxifen effect with demographic variables (race, marital status) carrying larger weights and several laboratory interactions reversing sign. This demonstrates concretely why a single global interpretable rule would misrepresent the underlying black-box behavior.

Limitations and open questions

The authors identify several constraints. The method's validity rests on the local approximation assumption — neighborhoods small enough that the black-box surface is captured by a first-order expansion — and richer surrogates would be needed where higher-order local structure exists; adequacy can be checked with goodness-of-fit measures such as $$\mathbf{Z}' = \mathbf{Z} + \min(1,\alpha)\odot\boldsymbol{\epsilon}'$$5. Multicollinearity among predictors threatens both interpretability and numerical stability of the local models, mitigated here by cross-validated weight decay rather than eliminated. The standard VAE may be suboptimal when features are strongly correlated; the authors suggest TCVAE, which decomposes the KL term into mutual information, total correlation, and dimension-wise components, as a potentially better generator, though they do not evaluate it. Synthetic treatment assignment assumes positivity and uses Bernoulli(0.5) sampling, which presumes both regimes are well-defined across the covariate space — an assumption untestable from the data alone. Finally, the framework is presented for continuous outcomes in single-stage settings; extension to survival endpoints (e.g., Cox-PH risk scores) is noted as unexplored.

Conclusion

LI-ITR couples VAE-generated realistic perturbations with a sparsity-regularized, hard-selection mixture of linear experts to extract subject-specific, clinically interpretable treatment rules from arbitrary black-box outcome models. Simulations show accurate recovery of true local treatment-effect coefficients and near-perfect optimal-treatment identification (>99%), including under moderate misspecification, where LIME fails structurally. The breast cancer application demonstrates that these gains translate to practice: the best policy value among competing methods, threefold better local fidelity than LIME, and transparent per-patient explanations. The open questions — generator choice under strong feature correlation, local-model expressiveness, and extension beyond continuous single-stage outcomes — define the immediate agenda for this line of work.