Gaussian Processes for Individualized Continuous Treatment Rule Estimation
Pavel Shvechikov Evgeniy Riabenko
stat.ME



apalike

#1 1




1
 
 
 Gaussian Processes for Individualized Continuous Treatment Rule Estimation
    Pavel Shvechikov 

 and Evgeniy Riabenko The reported study was funded by RFBR, according to the research project No. 16-07-01192 А.

 Department of Computer Science, Higher School of Economics
    December 30, 2023
=====================================================================================================================================================================================================

 

0
 

 Gaussian Processes for Individualized Continuous Treatment Rule Estimation


  


Individualized treatment rule (ITR) recommends treatment on the basis of individual patient characteristics and the previous history of applied treatments and their outcomes. Despite the fact there are many ways to estimate ITR with binary treatment, algorithms for continuous treatment have only just started to emerge. We propose a novel approach to continuous ITR estimation based on explicit modelling of uncertainty in the subject’s outcome as well as direct estimation of the mean outcome using gaussian process regression. Our method incorporates two intuitively appealing properties – it is more inclined to give a treatment with the outcome of higher expected value and lower variance. Experiments show that this direct incorporation of the uncertainty into ITR estimation process allows to select better treatment than standard indirect approach that just models the average. Compared to the competitors (including OWL), the proposed method shows improved performance in terms of value function maximization, has better interpretability, and could be easier generalized to multiple interdependent continuous treatments setting.



Keywords: Dose Finding, Gaussian Process Regression, Personalized Medicine




1.45 









§ INTRODUCTION

 



Many clinician studies conclude that reactions to a fixed treatment could vary significantly among individuals in a population <cit.>. Such heterogeneity of treatment effects <cit.> raises a question about the applicability of this "one size fits all" approach to clinical practice and leads naturally to personalized medicine paradigm – the idea of healthcare improvement as a result of tailoring treatments on the basis of patient individual characteristics. 



 §.§ Treatment rules

One of the crucial concepts in personalized medicine research is individualized treatment rule (ITR) (aka individualized treatment regime) which is usually defined as a mapping from measured patient characteristics to recommended treatment <cit.>. An obvious goal of ITR construction is a maximization of patients' response to treatment averaged over all individuals in a population. This goal is sometimes referred to as expected reward in a population or value function. ITR estimation is a procedure which uses data from either observational or randomized clinician studies to fulfill this goal. 



Personalized medicine often comes in one of the two forms <cit.>: either identification of the subgroups of patients who benefit from a given treatment, or identification of the optimal treatment for a specific patient. The latter form (which we also adhere to in this work) could be further subdivided by the approach to ITR estimation:
0em 
	
  * Indirect approach usually consists of a two-stage procedure. First, modeling of either conditional mean outcomes or contrasts between mean outcomes takes place. Second, the ITR is derived as the rule assigning the treatments that maximize the model of the conditional mean outcome.
	
	
  * Direct approach approximates expected reward in a population with its sample estimate as a function of ITR. Maximization of this expected reward estimate with regard to parameters of ITR yields the desired treatment rule. 


Specific ITR estimation methods differ a lot depending on whether the treatment is discrete or continuous. In the latter case ITR is also called individualized dose rule (IDR) and is of a particular practical importance because of its usage in therapies balancing drug dose between insufficient (not causing any effect) and excessive (causing severe adverse effects). One example of such dose administering therapies is thrombosis prevention using a drug called warfarin, which is shown to be more effective when prescribed on the basis of environmental and genetic patient factors <cit.>. Another illustrative example is a treatment of inflammatory bowel disease with thiopurines <cit.>. 
These examples motivate a strong need in developing reliable and efficient methods for ITR estimation. 

There is a substantial amount of literature devoted to ITR estimation dealing with discrete (mostly binary) treatment using an indirect approach <cit.> as well as a direct one <cit.>. However, despite its importance, as far as we know, there are very few methods dealing with IDR estimation. 

IDR estimation poses additional challenges because the set of admissible treatments has a cardinality of the continuum. There are infinitely many treatments to choose from, only a few of which might be actually present in the data. These difficulties yield trivial generalization of discrete methods to continuous treatment cases either impossible or intractable. 

One of the recent works devoted to IDR estimation <cit.> follows a direct approach with the focus on policy interpretability, which is achieved through the use of decision trees with splits guided by complicated purity measure. This split criterion is based on the usual value function from discrete ITR with additional smoothing. The proposed smoothing is analogous to introducing a stochastic policy with probability mass concentrated around the deterministic policy assignments observed in the training data. An interesting peculiarity of the work is that only treatments observed in the training data are considered as a split candidates. 

Another recent approach to IDR estimation <cit.> is devoted to developing an extension of O-learning algorithm <cit.>. It resembles in spirit the work of <cit.> because of chosen direct approach to IDR estimation and the presence of smoothing in value function estimate. However, unlike in <cit.>, the smoothing is used only as an intermediate step required for optimization purposes, so that the annealing of the kernel bandwidth to zero gives precisely the same value function as is used for discrete ITR estimation. 

Specifically, <cit.> introduced a regression based approach inspired by the <cit.> and SVM; the proposed algorithm comprises global optimization task tackled by DC procedure which is known to depend on initialization and could guarantee only local convergence.


All of the aforementioned algorithms follow the direct approach to IDR estimation consisting of maximizing mean conditional reward over individuals in a population (aka value function), but do not deal with variance of the patients' potential outcome, which might be crucial for a real world application. 

Consider a simple example when patient is told to choose between two mutually exclusive options: 
0em 
	
  * Take a drug A with 0.99 chance of recovery (reward = 100) and 0.01 probability to acquire physical disability (reward = -1000). 
	
  * Take a drug B with 0.85 chance of recovery (reward = 100) and 0.15 probability not to experience any effect (reward = 0).


Mean reward in the first case is 89 while in the second one it is 85, suggesting that rational patient should always choose the first option. However, it is known that people are not rational and tend to averse losses <cit.>. This common behavior was also shown to be true in a health domain <cit.>. Therefore, it seems very natural to explicitly incorporate reasoning about mean and variance of a possible outcome into an ITR. Nevertheless, such reasoning is impossible when using direct approach. 



 §.§ Indirect methods



Our work aims to develop an indirect approach to ITR estimation. Indirect methods are usually criticized for potential overfitting during the first stage (fitting regression model to rewards) and for possible model misspecification.

In defense of indirect methods, we argue that the concerns about misspecification could be reduced to a great extent through the use of a flexible model that could fit the data well <cit.>. 

However, the more powerful model is, the more it is susceptible to overfitting. To cope with the overfitting problem on the level of model design, we have chosen to follow Bayesian paradigm, which allow to do natural regularization through the averaging over model parameters and is known to be free of overfitting problem <cit.>. Nevertheless, any Bayesian model becomes prone to overfitting as soon as optimization with regards to hyperparameters is involved (so-called type II maximum likelihood procedure). Careful treatment of overfitting requires either very small number of hyperparameters or fully Bayesian inference with Markov chain Monte Carlo methods <cit.>.





Another argument in favor of indirect methods is their capacity to estimate the uncertainty of the expected outcome. Such estimation is impossible when following direct approach but could be invaluable in practice, providing the means for balancing exploration and exploitation.


General task of ITR estimation is similar to the classic reinforcement learning (RL) setting <cit.>: patient characteristics, their treatments and responses are direct analogs of state, action and reward concepts from RL literature. Moreover, direct and indirect approaches to ITR estimation in RL community are called value-based and policy-based approaches. Recently, it was shown <cit.> that these two approaches share much more in common than it was believed before. To some extent this finding reconciles two different ITR estimation approaches and gives rise to an unexplored idea of their combination.


The contributions of this work are summarized as follows. 
0em 
	
  * We propose a new indirect approach to IDR estimation that benefits from separation of uncertainty due to finite sample estimation and irreducible variance of patient outcomes. 
	
  * We show how to reduce the hard task of continuous inversion to simple univariate global optimization problem which is well-known in active learning and Bayesian optimization literature. 
	
  * We evaluate the proposed method on various simulation studies, confirming its superior performance for randomized dose-finding clinical trials. 
 





§ PROPOSED METHOD

We assume that train data comes from randomized dose trial and every training record consists of three major components: 
0em 
	
  * patient covariates ∈̧ℝ^p;
	
  * treatment doze a ∈ [0,1], selected at random within safe bounded dose range, which w.l.o.g. is assumed to be equal to [0, 1] interval;
	
  * patient outcome r ∈ [r_lo, r_hi],   r_lo, r_hi∈ℝ.
	
Note that we do not restrict the sign of outcomes but do require them to be bounded. We also introduce a shorthand notation _i = [_̧i,   a_i] for a concatenation of covariate vector and a treatment value. 


As was mentioned previously, indirect approach to IDR estimation requires building a model of the outcome value given patient covariates and treatment doze. 
One of the models meeting all criteria of being sufficiently flexible, robust to overfitting due to Bayesian nature, and able to make probabilistic inference is Gaussian process (GP) regression model <cit.>, which is the basis of this work. 

 




 §.§ Gaussian process framework
 
We proceed with minimal introduction to GP regression. For more thorough discussion one could refer to <cit.>. 

Gaussian process is a collection of random variables any finite number of which have joint Gaussian distribution. A GP is completely specified with its mean and covariance functions which we will denote as m() and k(, ') respectively. For clarity and without loss of generality <cit.> we will consider Gaussian process f with zero mean m() = 0 and denote f() ∼𝒢𝒫(0,  k(, ')). 

We assume that we were given a training dataset {(_i , y_i ) | i = 1,..., n} that consists of pairs of object description _i and noisy observation of function realization, i.e. y_i = f(_i) + ε, where the noise ε∼𝒩(0,σ_n^2 ). 


Having specified a joint Gaussian prior over the test output f_* and noisy train outputs y, the inference is carried out with conditioning on observed train outputs. This procedure allows to compute the mean and the variance of f_* posterior distribution given the test point _* using the following equations:

    f_*    ∼𝒩(f̅_*,  f_*) 
    f̅_*    = k_*(K + σ^2_n I)ıy
    f_*   = k(_*, _*) - k_*(K + σ^2_n I)ık_*,

where k_* = (k(_*, _1), ..., k(_*, _n)) and covariance matrix K has entries K_ij = k(_i, _j). To reduce clutter we will subsequently denote (K + σ^2_n I)ı as Λ. 

One of the main constitutes of GP flexibility is covariance function.
The choice of covariance function for particular application can be guided by Bayesian model selection <cit.>, however, for illustration purposes we have chosen probably the most commonly used squared exponential function with separate length-scales. We also extend this covariance function with the noise (aka nugget) term:

    k(_p, _q) 
    = σ_f^2 exp{ - ∑_i=1^p (x_pi - x_qi)^22θ_i} + σ_n^2	δ_pq,

where θ_i is called length-scale, x_pi is a p-th feature of i-th object and δ_pq is the Kronecker delta function. 

Note that all of the subsequent derivations could be replicated with other popular choices of covariance function (i.e. the Matérn class) and are not unique to squared exponential covariance function. 
















 §.§ Method formulation


Motivation of the proposed method is as follows. 
To solve the prescription problem for a new patient with characteristics _̧*, doctor, given the history of previously seen patients {(_̧i, a_i, r_i)}_i=1^n, have to simultaneously solve two conflicting tasks:
 0em 
	
  * Doctor has to be confident about potential dose effect and this confidence should be based on previous cases. 
	The more similar are patient characteristics _̧* to the ones from previous practice _̧i, the more doctor is sure that giving the same a_i to patient with _̧* will result in r_i. 
	
	
  * Doctor has to prescirbe the best treatment possible. The bigger is r_i, the more doctor wants a_* to be similar to a_i.


These two intuitive observations coupled with the notion of loss aversion in health domain <cit.> motivate prescription of the dose which maximizes probabilistic lower surface for the potential outcome. That is, in terms of mean and variance of posterior distribution over f_*

    a_* = _a   f̅_* - s √(f_* ),
 
where s is a scaling coefficient controlling the penalty for the uncertainty in the outcome. 

We argue that this simple heuristic, which hereinafter will be referred to as LCSL (lower confidence surface learning):
0em 
	
  * is a strong IDR competitive with the previously proposed approaches;
	
  * provides the means to account for loss aversion and uncertainty;
	
  * provides the means to perform adaptive trial design;
	
  * allows to explicitly address an exploration-exploitation tradeoff.













 §.§ Modelling correlation with exponential power family


Formula (<ref>) does not gives us a method to perform maximization over treatment domain. 
So we need to design a constructive way to carry out such an inversion.

Lets rewrite each of the terms in (<ref>) separately assuming that covariance function is squared exponential with separate length-scales. 
For the sake of brevity we omit length-scales in following derivations

    k(_p, _q) 
    = σ_f^2 exp{ -||_p - _q||^2_22} 
    = σ_f^2 exp{ -||_̧p - _̧q||^2_22}exp{ -||a_p - a_q||^2_22}

Lets denote

    K(_̧*) = σ_f^2 ·diag( exp{ -||_̧* - _̧1||^2_2/2},  … , exp{ -||_̧* - _̧n||^2_2/2})

Then one could write f̅_* in a form of a linear combination of exponential terms where α_i depends only on the train data and known covariates _̧* of test object:

    f̅_*    = [ exp{ -||a - a_1||^22}                ,  … , exp{ -||a - a_n||^22} ]K(_̧*)Λy
       = ∑_i α_i	exp{ -||a - a_i||^22},

where α_i = [K(_̧*)Λy]_i.

The same transformation could be done for f_*:

    f_*   = k(_*, _*) - k_*Λk_* 
       = k (_*, _*) - ∑_i,jγ_ijexp{ -||a - a_i||^2 + ||a - a_j||^2/2}

Here γ_ij = σ_f^4 Λ_ijexp{ -||_̧* - _̧i||^2 + ||_̧* - _̧j||^2/2}
and ∀  x_*   k( _*, _*) = σ_f^2, so neither β_ij nor k( x_*, x_*) depends on unknown a. 


To sum up, our optimization task could be simplistically formulated as follows

    max_a  ∑_i α_i	exp{ -(a - a_i)^2/2} - s √( k( _*, _*) - ∑_i,jγ_ijexp{ -(a - a_i)^2 + (a - a_j)^2/2})








 §.§ Global optimization


As is easily seen, the objective (<ref>) is a sum of unnormalized exponential functions with coefficients α_i, γ_ij of arbitrary sign. Thus, to allow for explicit maximization of this nonconvex objective, one need to employ global optimization methods. 

Even though the task (<ref>) is a very hard one, it is surely not a new one – problems of this type are inevitable in Bayesian optimization <cit.>, arising during maximization of the so-called acquisition function. 

There are many methods to solve global optimization problem of the form (<ref>), i.e. DIRECT <cit.>, which is derivative free and deterministic. In this work we adhere to a simpler yet efficient two-step approach to global optimization that is often implemented in Bayesian optimization frameworks <cit.>. First, seeds are generated – either uniformly at random, in a form of fixed step grid or as a Sobol' sequence <cit.>. Second, function is evaluated at each seed and (optionally) local optimization method is started with each seed as an initial point. Frequent choice of such local optimization algorithm is L-BFGS-B <cit.>. 







 §.§ Interpretability of the model


Flexibility of any model usually comes with the price of reduced interpretability which is often of particular interest for practical applications. 
However, as is evident from formula (<ref>), mean predictions of GP regression are additive with respect to covariance terms k_*. 
This means that, given test object _*, one could argue which objects from the train dataset influence mean prediction the most based on the components of covariance vector k_* and coefficients [Λ y]_i. The same reasoning holds for variance interpretation with the only difference of quadratic dependence replacing the linear one.


Additionally, the learned values of length-scales could serve as indicators of feature importances, giving rise to the so-called automatic relevance determination (ARD) effect. 


These simple model interpretations coupled with the recent work devoted to the tradeoff between accuracy and interpretability <cit.> along with possible interpretability increase due to additive GP modelling <cit.> allow to get even closer to interpretability level of linear and tree-based models. 






 §.§ Applicability to observational dose trials


The requirement of training data being obtained from randomized study is by no means necessary. Under usual assumptions in casual inference literature <cit.> the proposed method could be seen as a model of counterfactual mean outcome. These assumptions are as follows:
0em 
	
  * Conditional exchangeability (aka strong ignorability) – potential outcomes R^a are conditionally independent of A given covariates C, i.e. ∀ a,   R^a ⊥ A  C.
	
  * Positivity – there is nonzero probability of assigning any of the treatment levels to any of possible objects, i.e. ∀ C  p(C) > 0,  ∀ a    p(a  C) > 0.
	
  * Consistency – each observed outcome is precisely the one that would have been observed under given treatment level, i.e. R = R^A.



In case these assumptions hold true, structural model R^a  C = c of counterfactual outcome R^a is equal to conditional mean model R  C = c, A = a. This imply that the proposed method falls under umbrella of outcome regression causal inference models. 






§ NUMERICAL RESULTS


All experiments, unless stated otherwise, were conducted as follows. Train and test sample generation was repeated 50 times for each combination of scenario, train sample size, algorithm and variance penalty factor. Test sample size was fixed at the value of 1000. Rewards in train dataset were scaled to fit into [0,1]. For each train dataset we fit a Gaussian process regression that was allowed to internally restart optimization of marginal likelihood 10 times. The model maximizing marginal likelihood across these 10 restarts is chosen for subsequent evaluation on test data. Approximate model performance is obtained by plugging predicted optimal dose for test objects into the true underlying state-action value function Q(, A). Method quality is reported in terms of 𝒱̂(f), which is defined as approximate model performance averaged over all data generations. 

We have empirically observed that the choice of a particular global optimization method out of those described in section <ref> does not change performance much. Thus, we report all of the results using the simplest one – fixed step grid evaluation with number of grid point fixed at 50 for all of the simulation scenarios and settings. 

In this section we refer to our method as LCSL.X, where X specifies the value of variance penalty (<ref>) by the equation s = Φ^-1(X / 100), where Φ is normal CDF. 

All our experiments are built on top of Gaussian process framework <cit.>. 





 §.§ Simulation study


In this section we first present results of two simulation experiments which show that for simple scenarios our method works remarkably well.

We denote  for patient covariates matrix, i.e. = [_̧1, ..., _̧n], where _̧i is a column vector of i-th patient characteristics; A for a column vector of treatments, i.e. A = [a_1, ..., a_n]; C_i for an i-th column of  matrix.



Our first scenario simulates optimal treatment dependence on single covariate in a form of parabola (figure <ref>, left):

    R  ∼ 𝒩( Q_1(, A), 0.01),      A  ∼ 𝒰(0,1),     _̧i  ∼ 𝒰(0,1) 
    
    	Q_1(, A) = - 100 · (f^opt_1() - A)^2 
    
    	f^opt_1() = 4 · (C_1 - 0.5)^2




Our second scenario is a more complex one – it simulates optimal treatment assignment as a highly nonlinear function with abrupt change in the middle of covariate space and local periodicity (figure <ref>, right) :

    R  ∼ 𝒩( Q_2s(, A), 0.01),      A  ∼ 𝒰(0,1),     _̧i  ∼ 𝒰(0,1) 
    
    Q_2(, A) = - 100 · (f^opt_2() - A)^2 
    
    f^runge(x) = 	cos(3 π (4x - 0.3)) / (1 + 25 (4x - 0.55) ^2) 
    
    f^step(x) = 0.1 x · (10 + sin(20x) + sin(50x) ) - 1.3 
    
    f^opt_2() = (f^runge(C_1) + f^step(C_1) ·(C_1 > 0.5)) / 1.5 - 0.7






Because true f^opt() is nonlinear in both scenarios, we don't include linear OWL (L-OWL) from <cit.> into comparison. Kernel OWL (K-OWL) results are obtained by reusing the code from the supplementary material of <cit.>. 

Simulation results (table <ref>, <ref>) show that the proposed method works extremely well in both scenarios, showing unimprovable results for the first scenario. 


 




 







We proceed with results of LCSL on several more complex and higher dimensional simulations first introduced in <cit.>. In particular, we evaluate and compare our method on three scenarios, first two of which simulate randomized trial and the last one simulates observational trial. 




Scenario 3, randomized trial:


    R  ∼ 𝒩( Q_3(, A), 1),      A  ∼ 𝒰(0,2),     _̧i  ∼ 𝒰(-1,1)^30
    
    Q_3(, A) = 8 + 4 C_1 - 2C_2 - 2C_3 - 25 · (f^opt_3() - A)^2 
    
    f^opt_3() = 1 + 0.5 C_1 + 0.5 C_2



Scenario 4, randomized trial:


    R  ∼ 𝒩(Q_4(, A), 1),      A  ∼ 𝒰(0,2),     _̧i  ∼ 𝒰(-1, 1)^10
    
    Q_4(, A) = 8 + 4 cos(2π C_2) - 2 C_4 - 8 C^3_5 - 15 · |f^opt_4()-A|
     
    f^opt_4() = 0.6 ·(|C_1| ≥ 0.5) + C^2_4 + 0.5 log(|C_7|+1)





Scenario 5, observational trial:


    R  ∼𝒩(Q_5(, A), 1),      A ∼Trunc𝒩( f^opt_5(), 0, 2, 0.5),     _̧i  ∼ 𝒰(-1,1)^10
    
    Q_5(, A) ≡ Q_4(, A)
    
    
    
    f^opt_5() ≡ f^opt_4()




Results of randomized study simulation (tables <ref>, <ref>) allow to conclude that our method outperforms every other competitor (often by a large margin) in a randomized trial setting. Such superiority originates from the randomization of treatment assignment, which increases the accuracy of global response surface model and, as a result, improves quality of treatment recommendations. 



 
 







 




However, for simulation of observational study (table <ref>) we observe that out method yields superior results only when given enough data. 
Small samples do not allow to properly model response surface because all of the points are already concentrated in near optimal regions, giving insufficient information about treatment dependence on covariates.
This effect could be mitigated to some extent by large variance penalty (figure 2, last column). 

LCSL performance dependence on the variance penalty coefficient, denoted by s in objective (<ref>), needs additional elaboration. We have varied s on standard normal quantile grid (from 0.5 to 0.99 quantile with a 0.01 step) and observed a sustainable relationship between penalty value and method performance (figure <ref>). 

The first interesting effect, equally present in all scenarios, is that strong variance penalty always improves upon mean prediction given very small train data (positive trend on every plot in the first row of figure <ref>). 

The second observation is that increasing train sample size in randomized studies reduces the optimal variance penalty (first two columns in figure <ref>). This observation could be explained by increase of model accuracy and, as a consequence, diminished need in variance penalty. Note that for more complex setup (scenario 2) optimal penalty shrinks slower than for a simpler one (scenario 1). However, we also report a consistent small improvement in average 𝒱̂(f) irrespective of train data size for small values of penalty (figure <ref>). This effect can be viewed as a regularization, preventing the model of too assertive extrapolation. 

The last noteworthy remark concerns observational study (last column of figure <ref>). In such setting we observe that large variance penalty increases performance irrespective of train sample size. 
That is because data generation policy (A ∼𝒩( f^opt_4() , 0.5)) is by design close to optimal. Thus, in observational studies the value of variance penalty should be treated as a proxy of prior knowledge about the degree of data generation policy optimality. 




We conclude that different setups need to treat variance penalty differently. Randomized studies benefit most when treating variance penalty as a regularization. In contrast, observational studies should treat variance penalty as a means for explicit reasoning about degree of data generation policy optimality. 






 







§ CONCLUSIONS
 
We have presented an indirect approach to IDR estimation which is based on Gaussian process regression. 

In spite of obvious difficulties of indirect IDR estimation, we have shown that it is possible to reduce such an intractable problem to a univariate global optimization setting which is ubiquitous in Bayesian optimization and active learning.

We have discussed the advantages and disadvantages of the proposed method including explicit probability reasoning about potential outcome and computational challenges arising as a consequence of using flexible nonparametric model. 

We have also reported various simulation results which not only show that our method is highly competitive with other algorithms of ITR estimation, but uncover unexpected variance penalty interpretations. 



§ FUTURE WORK



As was previously noted, ITR estimation shares many similarities with policy search problem arising in reinforcement learning (RL). This suggests researches dealing with ITR and DTR (dynamic treatment regime) estimation could get inspiration from already existing methods in reinforcement learning world.


For example, it would be very interesting to see an adaptation of reward weighted regression <cit.> to direct DTR estimation. Also, an extension of the proposed LCSL method to multi-stage treatment assignments would be of a great interest. Such extension could probably inherit some of the useful ideas from data-efficient RL algorithm called PILCO <cit.>. 


Diffusion of statistical science and reinforcement learning may bring benefits to all parties and thus is highly desirable.




