Bayesian Outcome Weighted Learning

Authors: sophia  yazzourh
nikki l b freeman
Venue: NeurIPS_cc_2024_Workshop_BDU
Published: 2024-08-25 18:39:07

Abstract: One of the primary goals of statistical precision medicine is to learn optimal individualized treatment rules (ITRs). The classification-based, or machine learning-based, approach to estimating optimal ITRs was first introduced in outcome-weighted learning (OWL).  OWL recasts the optimal ITR learning problem into a weighted classification problem, which can be solved using machine learning methods, e.g., support vector machines. In this paper, we introduce a Bayesian formulation of OWL. Starting from the OWL objective function, we generate a pseudo-likelihood which can be expressed as a scale mixture of normal distributions. A Gibbs sampling algorithm is developed to sample the posterior distribution of the parameters. In addition to providing a strategy for learning an optimal ITR, Bayesian OWL provides a natural, probabilistic approach to estimate uncertainty in ITR treatment recommendations themselves. We demonstrate the performance of our method through several simulation studies.

Main paper: 
 introduction: 
 Statistical precision medicine uses data to match patients to treatments to improve health outcomes [Kosorok and Laber, 2019] . This can be achieved through individualized treatment regimes (ITRs), functions mapping patient characteristics to treatment recommendations. One goal is to learn optimal ITRs that, when implemented, result in better outcomes on average in the target population compared to other strategies, such as a one-size-fits-all approach [Kosorok and Laber, 2019, Murphy, 2003] .Machine learning-based approaches for optimal ITR learning, such as outcome weighted learning (OWL) introduced by Zhao et al. [2012] , convert the statistical ITR-learning problem into a classification problem, bypassing the need to estimate ancillary functions like the conditional mean of the outcome and uncoupling the quality of the estimated ITR from the quality of the ancillary function estimates. In contrast, OWL enables direct learning the optimal ITR. Because of its flexibility, many extensions to OWL have been developed [Zhao et al., 2015 , Zhou et al., 2017 , Liu et al., 2018 , Fu et al., 2019 , Zhao et al., 2019] .Bayesian approaches for learning optimal ITRs have also been proposed, including Bayesian machine learning (BML) [Murray et al., 2018] which employs a variant of approximate dynamic programming and likelihood-based methods that model both the distribution of final and intermediate outcomes within the Bayesian framework [Thall et al., 2002 , 2007 , Arjas and Saarela, 2010 , Zajonc, 2012 , Xu et al., 2016 , Yu and Bondell, 2023 ].Here we focus on the classification-based perspective for learning optimal ITRs. While these approaches are powerful in terms of predictive power, they do not have a natural mechanism for inference or for quantifying the uncertainty in the learned ITR's treatment recommendations. Yet capturing this kind of uncertainty is crucial for healthcare decision-making. To address this, we introduce Bayesian outcome weighted learning (Bayesian OWL), the first optimal ITR-learning method to directly learn optimal ITRs from the classification perspective. By transforming the OWL framework into a probabilistic model, we generate a posterior distribution that enables inference and uncertainty quantification for the treatment recommendations themselves. 
setting: 
 We let A ∈ A = {-1, 1} denote the action, or treatment, and assume that observed treatments are assigned randomly as in a clinical trial with P (A = 1) = ρ known. Let X i = (X i,1 , . . . , X i,p ) ⊤ ∈ X denote the p-dimensional biomarker and prognostic information vector, and let R denote the outcome (bigger is better). We further assume that the reward can be rescaled so that R > 0. Then, the observed data is iid replicates of (A i , X i , R i ) for i = 1, . . . , n.An ITR is a function d that maps from patient features X to a recommended treatment in A. For a given ITR d, the value ofd is V (d) = E[R(d)],where R(d) is the reward we would observe if treatments were allocated according to rule d. An optimal ITR d opt satisfies V (d opt ) ≥ V (d) for all d ∈ D, where D is a class of ITRs. Our goal is to learn an optimal ITR d opt . Under the assumptions of causal consistency, the stable unit treatment value assumption, no unmeasured confounding, and positivity, V (d) can be identified from the observed data andV (d) = E X {E[R|A = d(x), X = x]}.If we let P denote the distribution of (X, A, R), and P d denote the distribution of (X, A, R) when A = d(X), then the reward we would expect if ITR d(X) were followed is given byEQUATIONZhao et al. [2012] showed that maximizing Equation ( 1) is equivalent to minimizing a weighted classification problem with the 0-1 loss function. By replacing the 0-1 loss with a confex surrogate loss, they formulated an objective function that can be efficiently minimized using machine learning techniques. To learn optimal ITRs, they proposed the OWL which minimizesQ OWL n (β) = 1 n n i=1 r i a i ρ + (1 -a i )/2 (1 -i h(x i , β)) + (2)where the empirical measure replaces the true measure, (z) + = max(z, 0) denotes the hinge loss function, and h(•) is the ITR parameterized by β. Song et al. [2015] introduced a penalized variant of OWL that included a regularization term p λ (β) for the ITR parameters. POWL minimizes the objective functionEQUATIONwhere p λ (β) is a penalty function and λ is a tuning parameter.Although the pure machine learning framework is powerful, it is limited in its ability to capture and model uncertainty as in a statistical framework. Polson and Scott [2011] bridged this gap between pure machine learning and statistical modeling for SVMs by showing how to cast SVM into a Bayesian framework. They considered the L α -norm regularized support vector classifier that chooses β to minimized α (β, ν) = n i=1 max(1 -r i x ⊤ i β, 0) + ν -α p j=1 |β j /σ j | α (4)where σ j is the standard deviation of the j-th element of x and ν is a tuning parameter. For this objective function, the learned classifier is a linear classifier. Polson and Scott [2011] shows that minimizing Equation ( 4) is equivalent to finding the mode of the pseudo-posterior distribution p(β|ν, α, r) ∝ exp(-d α (β, ν)) ∝ C α (ν)L(r|β)p(β|ν, α), where C α is a pseudo-posterior normalization constant. The main theoretical result from Polson and Scott [2011] is that the pseudolikelihood contribution L i (r i |β) is a location-scale mixture of normals (Polson and Scott [2011] , Theorem 1). 
our approach: 
 We follow the strategy employed by Polson and Scott [2011] to cast the OWL objective function into a probabilistic Bayesian learning framework. The conversion is not one-to-one since Polson and Scott [2011] constructed a Bayesian model for a standard SVM whereas the objective function for OWL Equation ( 2) is a weighted SVM problem. Throughout, we will assume R > 0. When this is not the case, a distance-preserving transformation of R from R to R + can be used. Assuming that h is linear, i.e., h(x i , β) = x ⊤ i β and following the strategy taken in Theorem 1 of Polson and Scott [2011] , the contribution of a single observation to the pseudo-likelihood is given byEQUATIONor in other words that L i (a i , λ i |r i , x i , β) is a scale mixture of Gaussians.In their formulation of Bayesian SVM, Polson and Scott [2011] use the exponential power prior for β, a prior that can be show to be equivalent to L1-regularization of the regression parameters.Regularization of the OWL parameters has been explored as in Song et al. [2015] . In this paper, we first construct our method as an analogy to the original formulation of OWL without penalization.We make this choice because (1) our primary aim is to develop a Bayesian classification-based ITR learning approach, and because (2) L1-regularization does not necessarily yield sparse rules (see the discussion in Section 4.1 of Polson and Scott [2011] ). However, regularization helps avoid overfitting, a common problem in machine learning. Thus, we also explore penalty priors for β, including the exponential power prior distribution and the spike-and-slab prior distribution Table 1 . Normal distribution p j=1 1 √ 2πσ 2 0 exp -1 2 (βj -µ0,j ) 2 σ 2 0, where µ 0 and σ 2 0 are hyperparameters of the Normal distribution Exponential power distributionp j=1 ω -1 2 j • exp -1 2ν 2 p j=1 β 2 j σ 2 j ωj • p j=1 p(ω j |α),where ω and α are hyperparameters and when α = 1 p(ω j |α) ∼ Exponential(2) Spike-and-slab distribution p j=1 (γ j N (0, ν 2 σ 2 j ) + (1 -γ j )δ 0 (β j ))π γj (1 -π) 1-γj where π and γ j are hyperparameters and δ 0 is the Dirac measureTo draw from the pseudo-posterior distribution, Polson and Scott [2011] employed two algorithms, an expectation-minimization (EM) approach and a Gibbs sampling approach. The approach we take is the latter. Although sampling the pseudo-posterior is likely to be more time intensive than estimation via the EM algorithm, the rationale for a fully Bayesian approach is to enable uncertainty quantification (Section 3.3). Because of conjugacy, derivation of the Gibbs sampling algorithms is straightforward and is not presented here (see Yazzourh and Freeman [2024] for details).Using the posterior predictive distribution, we can make treatment recommendations for a new patient and quantify our uncertainty in our recommendation. Let Θ = {β, λ} and ã denote the recommended treatment for a new patient with features x. Then p(ã = 1|x, X, r, a) = Θ p(ã = 1|x, X, r, a, β, λ)p(β, λ|x, r, a)dθ.For the class predictions, we can use the probit model which has the form p(a = 1|x) = Φ(x ⊤ β) where Φ is the cumulative distribution function of the standard normal distribution. Thus we can write the posterior predictive distribution asEQUATION)4 Simulation study 
classification performance: 
 We conducted simulation studies to assess the classification performance of the proposed method, following Zhao et al. [2012] and Song et al. [2015] . We compared the performance of OWL, Bayesian OWL with normal priors for β, Bayesian OWL with exponential power prior for β, and Bayesian OWL with spike-and-slab prior for β. For each simulated patient, we generated a 10-dimensional vector of patient features, X 1 , . . . , X 10 , drawn independently and uniformly distributed on [-1, 1]. Treatment A was drawn from {-1, 1} independently of the prognostic variables with P(A = 1) = 1/2. The outcome variable R was normally distributed with mean A ) and standard deviation 1, where T 0 (X, A) was the interaction term between treatment and patient features. We examined two scenarios for the treatmentfeature interaction term: Scenario 1 : T 0 (A, X) = (X 1 + X 2 )A and Scenario 2 :Q 0 = 1 + 2X 1 + X 2 + 0.5X 3 + T 0 (X,T 0 (A, X) = 0.442(1 -X 1 -X 2 )ABoth scenarios 1 and 2 had linear decision boundaries determined by X 1 and X 2 . For scenario 1, the true optimal rule was given by 1(X 1 + X 2 > 0), while for scenario 2, it was 1(1 -X 1 -X 2 > 0). OWL was implemented with a linear kernel. For Bayesian OWL, Gibbs sampling was used to draw from the posterior distributions of the parameters 500 times. The first 150 draws were discarded as "burn-in" and point estimates of β were computed by taking the mean of the draws from the posterior distribution. Throughout, we set the hyperparameter ν = 0.8.For each scenario, we varied the training dataset from 100 to 200, 400 and 800 and tested on 1000 patients. For each training set size, we conducted 200 simulation runs. We evaluated classification performance using the misclassification rate, As expected, the classification performance improved among all the ITR learning methods evaluated as the sample size increased. However, OWL consistently outperformed Bayesian OWL in all sample sizes and in both scenarios. We hypothesize that, with additional hyperparameter tuning, the performance of Bayesian OWL can be improved. Ordinarily, one would be hesitant to propose a method that is dominated by an existing method. However, the dominance of OWL is with respect to the misclassification rate. OWL, even with 800 samples in our simulation, has a 10% misclassification rate, and there is no way to determine which of the 10% of the simulated patients are likely misclassified (given a non-optimal treatment recommendation). In contrast, Bayesian OWL yields the entire posterior distribution of the estimated optimal ITR and thereby allows for immediate uncertainty quantification of individual-level treatment recommendations. In essence, Bayesian OWL can inform us of which treatment recommendations it is less certain about whereas OWL cannot. We demonstrate this in Section 4.2.To highlight the utility of quantifying the uncertainty of individual-level treatment recommendations, we trained Bayesian OWL model using simulated data as in Scenario 1 and the exponential power prior. Then, we simulated another 1000 patients using the same generative approach, creating a fine grid for the key variables X 1 and X 2 , which are part of the true optimal rule. This allowed us to estimate uncertainty across the domain of the true optimal ITR for all combinations of X 1 and X 2 within [-1, 1] 2 .Figure 1 : Heatmap of uncertainty quantificationFigure 1 shows how Bayesian OWL quantifies uncertainty in its treatment recommendations.The true optimal ITR assigns treatment A = 1 to patients in the upper-right (where X 1 + X 2 > 0) and treatment A = -1 to those in the lowerleft. Using the posterior predictive distribution (Section 3.3), we calculated the uncertainty for each treatment recommendation individually for the 1000 patients in our test set. The heat map in Figure 1 visualizes this uncertainty across the ITR domain, with lighter colors (yellow and light green) indicating higher certainty (close to 1) and darker colors (purple and dark blue) indicating lower certainty (close to 0). As expected, Bayesian OWL shows greater certainty for patients with features far from the decision boundary and less certainty for those near it. Furthermore, in Figure 1 , misclassified individuals are shown as red points for those incorrectly recommended treatment -1 instead of 1 and orange points for those incorrectly recommended treatment 1 instead of -1. As expected, these misclassified patients are near the decision boundary, particularly in regions where the model has the greatest uncertainty (indicated by purple shading). 
discussion: 
 We introduce a Bayesian formulation of OWL, the first Bayesian strategy for directly learning an ITR from the classification perspective. Like OWL, our approach directly models the decision rule without relying on conditional mean models. It also allows us to quantify uncertainty in treatment recommendations at the individual level, unlike typical approaches focused on population-level uncertainty. This insight can enhance clinical study design by identifying patient types where treatment recommendations are confident and those needing further sampling.Our work has some limitations. We only consider linear rules, but nonlinear rules may sometimes be more appropriate or offer significant clinical improvements. Henao et al. [2014] proposed a Bayesian SVM for nonlinear decision boundaries, which could be adapted for this purpose. Additionally, we have not fully explored variable selection as in Song et al. [2015] , limiting its use in high-dimensional settings.The classification approach for optimal ITRs is powerful, but it struggles with inference and uncertainty quantification. Bayesian OWL addresses this by combining the strengths of direct learning with a probabilistic framework, expanding inferential potential and improving precision medicine evidence.Workshop on Bayesian Decision-making and Uncertainty, 38th Conference on Neural Information Processing Systems (NeurIPS 2024). 
 
 Appendix:  


Decision: This paper is borderline. The reviewers' actual text in general reads more positive than the scores, and one of the reviewers' own breakdown lists higher numbers. Reviewers in particular are concerned about lack of comparisons, but also note that the method is reasonably novel, which might be contributing to this. Overall my impression is this paper lands slightly above the decision boundary, and I recommend accepting this.

Reviews: review: title: title: it is unclear to me whether this work focuses on developing an approach or an application
main_review: main_review: **Summary**

This paper introduces a Bayesian approach to Outcome Weighted Learning (OWL) for learning optimal individualized treatment regimes (ITRs) with uncertainty quantification. The key contribution is the transformation of the OWL framework into a Bayesian probabilistic model, which allows for inference and individual-level uncertainty quantification.

**Strengths**
1. The manuscript introduces a novel Bayesian approach to OWL, which addresses key limitations in the field, such as the lack of uncertainty quantification in traditional machine learning methods. 
2. The proposed method directly learns optimal ITRs without needing to estimate conditional mean models, which provides an innovative contribution to Bayesian machine learning in healthcare applications.
3. The simulation study and discussion highlight the potential of this approach in precision medicine, with practical implications for improving treatment decisions.

**Weakness**
1. An abstract is absent in the manuscript.
   
2. While the paper claims to present a Bayesian learning method, it frequently shifts focus toward a specific application within individualized treatment regimes. This introduces terminology and concepts, such as biomarkers, prognostic information vectors, and cohorts, which may be unfamiliar to readers who are not experts in healthcare or precision medicine. These domain-specific terms can obscure the core contributions related to Bayesian optimization. The manuscript would benefit from a clearer focus on whether the work is primarily about developing a Bayesian learning framework or applying it to healthcare, with adequate explanations of these concepts for non-experts.

3. Figure 1 could be improved in terms of readability. The current use of lighter and darker colors is not easily distinguishable. Moreover, the labels, numbers, and markers are too small to read comfortably. 

**Questions**
1. Can the authors clarify the specific aspects of the Bayesian approach that enable uncertainty quantification at the individual level? 
2. How does the proposed method generalize beyond the healthcare domain?
title: title: Official review
main_review: main_review: **Summary**: This work presents a Bayesian approach to the outcome-weighted learning (OWL), a classification-based framework for estimating optimal individualized treatment regimes (ITR). The motivation behind this approach is to incorporate uncertainty in the estimation and prediction of optimal ITRs, which is crucial for decision-making and applications in personalized medicine.

**Evaluation** 
- **Quality (6/10)**: The proposal is of acceptable quality, but it needs refinement and expansion in key areas to elevate it to a higher scientific standard.
- **Clarity (6/10)**: Overall, the manuscript conveys its main ideas within the four-page limit. Yet, certain key points could benefit from further elaboration. For example, the discussion of the pseudo-posterior is somewhat lacking, particularly in clarifying that it is not strictly Bayesian, as it isn't derived from an actual data-likelihood. Additionally, some of the notation changes add confusion, such as using $r_i$ in the product $r_ix_i^\top\beta$ in section 2.3, but switching to $a_i$ in $a_ix_i^\top\beta$ at the start of section 3.
- **Originality (5/10)**:  Although the proposal of a pseudo-Bayesian OWL method for estimating ITR is somewhat original, it seems to me that the approach largely builds on and adapts concepts from Bayesian SVMs. I would have liked to see more methodological innovation addressing the unique challenges of learning optimal ITR specifically. 
- **Significance (5/10)**: I appreciate the motivation to introduce Bayesian-like and probabilistic modeling tools to the challenge of estimating ITR with proper uncertainty quantification. In my view, aside from the contributions of T. Murray, S. Murphy, and a few others, there is indeed a notable gap in the literature on this topic, primarily due to objections to likelihood-based inference in modern semiparametric causal inference. This work makes progress by proposing a pseudo-Bayesian approach. However, it falls short in some areas. For instance, it relies on data from RCTs with binary treatments, which are often scarce, small in size, and prone to selection bias. Extending the approach to observational settings, where $\rho$ is neither known nor fixed but must be estimated from a sufficient set of confounders, would increase its applicability to real-world healthcare scenarios 

**Pros**
- The problem is well motivated.

**Cons**
- The approach imposes significant constraints, such as focusing on linear optimal ITR, low-dimensional settings, and experimental data with randomized binary treatments, which limit its broader applicability.
- Lacks a clear benchmark comparison highlighting its strengths and weaknesses relative to Bayesian $Q$-learning for ITR.