Efficient Experimentation for Estimation of Continuous and Discrete Conditional Treatment Effects

Authors: muhammed t razzak
panos  tigas
andrew  jesson
yarin  gal
uri  shalit
Venue: NeurIPS_cc_2024_Workshop_BDU
Published: 2024-09-05 08:01:15

Abstract: Accurately estimating personalized treatment effects often demands substantial data, incurring high costs across diverse applications such as personalized advertisement delivery and clinical trials. Existing methodologies employ deep models to estimate treatment effects in high-dimensional data, often relying on randomly selected experiments. We explore the potential of active learning techniques to enhance the efficiency of experimentation. Our focus centers on a relatively underexplored yet common scenario where each unit is subject to experimentation only once. We build upon the Bayesian active learning framework, to select units, and a treatment to apply to the unit, that maximize the information gain from each experiment. Our approach is flexible, accommodating both discrete and continuous treatment settings. Furthermore, we address the inefficiencies in batch experimentation by employing a greedy and a policy gradient-based optimization strategy.
We validate the effectiveness of our proposed method on synthetic and high-dimensional semi-synthetic datasets (based on IHDP and TCGA). Our results show significant improvements in experimentation efficiency over the baseline methods.

Main paper: 
 introduction: 
 Experimentation often incurs significant costs in terms of resources, time, and funding. This necessitates the development of more efficient and targeted experimental designs. Adaptive trial design is a promising paradigm for conducting experiments, allowing more efficient use of scarce resources (Foster et al., 2020) . Given a pool of samples and a set of (possibly continuous) interventions or treatments we wish to study, our goal is to construct an experiment that yields the most informative results about the effects of the interventions on a predefined outcome, while minimizing costs and resources used.The majority of experimental design literature has primarily focused on a setting in which one may conduct experiments on the same object multiple times, usually with differing actions (also known as treatments, or interventions) (Rainforth, 2017) . A somewhat less explored scenario is where one can actively select a unit from a pool of available units, and then select only a single action to apply to this unit. This scenario has been referred to as active learning for trials (Deng et al., 2011) or adaptive trials (Chen et al., 2021) . This scenario is relevant in clinical trials, animal studies, and any other scenario where the units are heterogeneous and can change after being subjected to the intervention, or can only be treated once for ethical or cost reasons.Our paper makes several contributions to the field of active learning for experimentation:1. We propose a new setting for active learning in which the objective is to efficiently select the most informative unit-treatment combinations for experimentation (Section 2). 2. We utilize deep learning methods enabling us to model units with high-dimensional covariates and explore a wide range of treatment options (binary, discrete, and continuous options). We select unit-treatment combinations with the highest Expected Information Gain by introducing both a greedy algorithm and novel gradient-based algorithm to optimize the joint mutual information between all experiments selected in each batch (Section 3). 3. We demonstrate the efficacy of our methods on synthetic and high-dimensional semisynthetic datasets, resulting in significant improvements in efficiency (Section 4). 
problem formulation: 
 Within the Neyman-Rubin causal model (Rubin, 1974) , each unit u has a potential outcome Y t (u) associated with every intervention t ∈ T , where T is the (possibly infinite) set of interventions under consideration. We are interested in efficiently learning the conditional average potential outcome(CAPO) function, E[Y t | X = x].In the regime of continuous interventions, this is often referred to as the (conditional) "dose response function."The expectation of the potential outcome given a set of covariates X = x describing u is identifiable from data where T is directly intervened on, assuming that X and T are causal parents of Y, that the observed outcome corresponds to the potential outcome of the assigned treatment, and that the treatment assigned to one unit, u, does not effect the outcome observed for any other unit, u ′ . Under these conditions, the CAPO takes the form of the expectationEQUATIONWe are specifically interested in a k-round, hybrid, "pool based" setting, consisting of a nonreplenishable pool of units, U = {u i } n i=1 ⊆ U, and a set of possible interventions T ⊆ T . This is a common setting, where experimenters will do sequential experiments either in batch or non-batch setting and have access to each unit once. A common example is preclinical trials, done in stages on animals to study safety, efficacy, and biological activity of a drug. Each unit, described by high-dimensional covariates, is unique and non-replenishable.We consider both discrete and continuous intervention spaces, T . At each round i ∈ [k], the experimenter uses a policy, π : U × T → S ⊆ X × T , to select a b-sized batch of unit-intervention pairs, s = {X(u j ) = x j , t j } b i=j . The experiments are then run and the experimenter observes the outcomes, {y j } b j=1 , for each. At the end of each round, the experimenter fits an estimator, µ(x, t; θ), of the CAPO, µ(x, t), to the data accumulated, D train = ∪ k i=1 {x j , t j , y j } b j=1 i . 
method: 
 This section outlines our proposed selection policies which aim to maximize the Expected Information Gain over a set of experiment designs in a batch of size b, denoted by ξ i=1:b = (x i , t i ).In contrast to traditional experimental design, we need to acquire units, x ∈ U, without replacement and select the treatment, t ∈ T , with replacement. Furthermore, while units are discrete and drawn from a pool set, treatments can be discrete or continuous.The problem is defined as selecting a batch of experiments that maximize the EIG objective:{(x, t) 1:b } * = arg max {(x,t) 1:b } EIG({(x, t) 1:b }) (2)where the utility set function is defined as the EIG objective EIG(•) ≜ I(Y ; Θ | •, D pool ).We propose several acquisition policies for both discrete and continuous treatment settings:Top-K. This strategy selects the top-k design with the highest estimated EIG per design. However, this does not maximize the EIG of the batch, as it does not take into account how other experiments in the batch affect the mutual information of each individual score. This can result in redundant experiments being conducted and results in inefficient experimentation (Foster et al., 2021; Kirsch et al., 2019) .Softmax Top-K. Kirsch et al. (2022) note this inefficiency and propose sampling the top-k experiments from a Softmax distribution, with temperature β, to the EIG scores. This selection policy has shown superior performance in terms of efficiency. It is hypothesized that this is due to the perturbation of the distribution taking into account that the mutual information changes when other experiments are selected.Greedy Optimization. When the set function we wish to optimize is known to be submodular and non-monotonic, then a greedy optimization-based algorithm can be used to maximize the set, which enjoys 1 -1 ϵ approximation guarantees (Nemhauser et al., 1978) . As shown in (Tigas et al., 2022; Kirsch et al., 2019) , Expected Information Gain over continuous outcomes (i.e. regression problems) is both submodular and non-monotonic thus it is suitable for applying a Greedy optimization-based algorithm.Policy-Gradient Based Optimization. To alleviate the approximation shortcomings of the greedy optimization-based algorithm, we notice that the objectives are differentiable with respect to the trial parameters (units and dosages), thus we can employ gradient-based methods to design experiments that maximize our objectives. We parametrize the policy π ϕ (X, T ) as joint distribution over units X and dosages T .We provide a more details on each method in Appendix B. 
experiments: 
 We evaluate our acquisition policies on synthetic and semi-synthetic datasets, and show significant improvements over the baseline.Datasets. We evaluate on 3 datasets: a synthetic dataset, and two semi-synthetic datasets based on the features/covariates in the IHDP (Hill, 2011; Shalit et al., 2017) and TCGA datasets (Cancer Genome Atlas Research Network et al., 2013) . For the synthetic dataset, suitable for discrete and continuous treatments, we develop a dose-response function based on the generalized dose response function developed in Taleb & West (2023) . The covariates of the units are 1-dimensional and normally distributed, with the treatment being either a continuous range or discrete values. For the two semi-synthetic datasets, we obtain covariates from two standard benchmark datasets in causal inference, IHDP and TCGA. The units in IHDP dataset contain 25 features, while the units in the TCGA dataset contain 4000 features. Further details are provided in Appendix C.Model. Our objectives rely on methods that are capable of modelling uncertainty and handling high-dimensional data modalities. For this we rely upon Deep Bayesian Neural Networks (BNNs). We utilize Deep Ensembles (Lakshminarayanan et al., 2017) which can be seen as approximate Bayesian Inference (Wilson & Izmailov, 2020) . We utilize a simple multi-layer perceptron S-learner (Rumelhart et al., 1986; Künzel et al., 2019) . We provide the hyperparameters in Appendix D.Metrics. The Mean Integrated Square Error (MISE) measures how well the model estimates the conditional average potential outcome (CAPO) across the entire dosage space:EQUATION)The metric is computed over a held-out test-set. This metric reflects our objective: to efficiently build an understanding of the entire unit-treatment-response function, through as few experiments as necessary. 1 (in A .1), the use of the EIG objective resulted in a more strategic distribution of the covariate space, with a higher density in regions where the outcome is more uncertain or variable. In the treatment space, we find a level of overlap that allows for more accurate estimation of the counterfactuals (Jesson et al., 2021) . This improved coverage and estimation leads to a reduction in MISE, as evidenced by the more accurate covariate-response curves shown in the figure.Synthetic Dataset Performance. The Greedy and Soft Top-K acquisition functions demonstrate strong performance on the discrete synthetic dataset. Notably, the Greedy method consistently achieves the lowest MISE values across various acquisition sizes and numbers of experiments conducted, making it the best-performing method for discrete treatments. When comparing the performance on the continuous synthetic dataset to the discrete synthetic dataset, we observe that the MISE values are generally lower for the discrete case. The policy gradient method begins to outperform the other acquisition functions as the acquisition size increases.Policy-gradient offers robust and scalable performance on high-dimensional data. We evaluate our acquisition functions' performance on two semi-synthetic datasets: IHDP (Table 3 and TCGA (Table 4 in Appendix A). These datasets are designed to test the efficiency of our methods in highdimensional settings, the primary focus of our work. The Greedy and Policy methods consistently achieve the lowest MISE values. For instance, on the IHDP dataset, Greedy achieves a MISE value of 0.1107 for the smallest acquisition size (4) when conducting 192 experiments, outperforming the random baseline by 27.1%. The Policy method performs well, particularly for larger acquisition sizes and a higher number of experiments conducted.The results on these datasets demonstrate that our methods and acquisition functions are scalable, versatile, and effective in real-world scenarios. The greedy and policy-gradient acquisition functions, in combination with deep learning for effective representation learning, provide an effective framework for efficient experimentation and estimation of conditional treatment effects. 
conclusion: 
 In this paper, we have proposed a novel framework for efficiently estimating treatment effects. We provided theoretical justification for our information-theoretic based objective functions, along with our combinatorial optimization. We demonstrated our methods on both synthetic and semi-synthetic datasets, showing significant gains in efficiency over the standard baseline. Future work could include extending our methods to handle larger batch sizes, and incorporating additional objectives for additional utility.  
 
 Appendix: A.1 Detailed Analysis of EIG Optimization Effects Table 3 shows the MISE for proposed acquisition functions on the IHDP dataset, for various acquisition sizes and number of experiments conducted.Table 3 : MISE for proposed acquisition functions on the IHDP dataset, for various acquisition sizes and number of experiments conducted.Table 4 shows the MISE for proposed acquisition functions on the TCGA dataset, for various acquisition sizes and number of experiments conducted. The results on both IHDP and TCGA datasets demonstrate the effectiveness of our proposed methods in high-dimensional settings. For the IHDP dataset, we observe that the Greedy method performs particularly well for smaller acquisition sizes, while the Policy method shows strengths in larger acquisition sizes and with more experiments conducted.On the TCGA dataset, which has a much higher dimensionality (4000 features), we see that both Soft Top-K and Policy methods outperform the random baseline consistently. The Policy method, in particular, shows robust performance across different acquisition sizes and number of experiments, highlighting its effectiveness in handling high-dimensional data.These additional results further support our main findings that our proposed methods, especially the Policy-gradient based approach, offer robust and scalable performance on high-dimensional data, providing significant improvements over random experimentation in estimating conditional treatment effects.This section provides more detailed descriptions of the algorithms used in our study.The Top-K algorithm selects the top-k designs with the highest estimated Expected Information Gain (EIG) per design. Here's a more detailed description of the algorithm:Algorithm 1: Top-K batch selection 1 for b = 1 to B do 2 x * , t * = arg max x∈D pool ,t∈T EIG({(x, t)}) 3 Add (x * , t * ) to the batch 4 D pool ← D pool \ {x * } ▷ Remove selected unit from pool 5 end B.2 Softmax Top-K AlgorithmThe Softmax Top-K algorithm introduces stochasticity to the selection process by sampling from a softmax distribution of the EIG scores: The Greedy Optimization algorithm iteratively selects the experiment that maximizes the marginal gain in EIG:Algorithm 3: 1 -1 ϵ Greedy batch selection 1 A ← ∅ 2 for update step c = 1 . . . B do 3 x * , t * = arg max x∈D pool ,t∈T EIG(A ∪ {(x, t)}) 4 A ← A ∪ {x * , t * } ▷ Select unit without replacement 5 D pool ← D pool -x *The Policy-Gradient Based Optimization algorithm uses gradient ascent to optimize a parameterized policy for selecting experiments:Algorithm 4: Policy-Gradient Based Optimization batch selection 1 Initialize policy parameters ϕ x for units and ϕ t for treatments2 for c = 1 to C do 3 Sample batch {(x, t) 1:B } ∼ π ϕ (X, T ) 4 Calculate EIG({(x, t) 1:B }) 5 Update ϕ x ← ϕ x + α x ∇ ϕx EIG({(x, t) 1:B }) 6 Update ϕ t ← ϕ t + α t ∇ ϕt EIG({(x, t) 1:B }) 7 end 8 Sample final batch {(x * , t * ) 1:B } ∼ π ϕ (X, T )The Expected Information Gain is calculated using a Monte Carlo estimation:EIG({(x, t) 1:B }) ≈ 1 M M m=1 log p(y m |x, t, θ m ) 1 N N n=1 p(y m |x, t, θ n ) (4)where θ m ∼ p(θ|D) are samples from the posterior distribution over model parameters, y m ∼ p(y|x, t, θ m ) are samples from the predictive distribution, and M and N are the number of Monte Carlo samples used.These algorithms form the core of our approach to efficient experimentation for estimating conditional treatment effects. Each algorithm offers different trade-offs between computational complexity and the quality of selected experiments, as discussed in the main text.The synthetic dataset is designed to be suitable for both discrete and continuous treatments. We develop a dose-response function based on the generalized dose response function introduced in Taleb & West (2023).The covariate x is generated from a normal distribution reflecting some underlying continuous variable:x ∼ N (70, 20)Treatments t are continuous and normally distributed:t ∼ N (15, 5)The outcome Y for each unit is computed through a response function that depends on both the covariate x and treatment t. The function is defined as:Y = min(f 1 , f 2 )wheref 1 = 1 1 + e -t+offset-5 f 2 = - 2 1 + e (-t+offset+2)•3 + 1 and offset = x -20 4We standardize the covariates, treatments, and output to enable the model to learn the function more easily. For the discrete version of this dataset, we simply choose two values for t, 10 and 20, and use the above dose-response function.The Infant Health and Development Program (IHDP) is a semi-synthetic dataset (Shalit et al., 2017) commonly used in literature to study the performance of causal effect estimation methods. Each unit is represented by 25 covariates describing different aspects of the infants and their mothers.The covariate x is a normalized 25-dimensional feature vector derived from the IHDP dataset. We produce a response function from it by projecting it into a 1-D space, using the parametric Uniform Manifold Approximation and Projection (UMAP) Sainburg et al. (2020) algorithm. We then normalize it to a normal distribution reflecting the following distribution:x ∼ N (70, 20)Treatments t are continuous and normally distributed:t ∼ N (15, 5)The outcome Y for each unit is computed using the same response function as in the synthetic dataset.This semi-synthetic dataset uses high-dimensional covariates seen by the model, but the underlying dose-response function is based on the 1-D projection.The TCGA dataset consists of gene expression measurements for cancer patients (Cancer Genome Atlas Research Network et al., 2013) . There are 9659 samples for which we used the measurements from the 4000 most variable genes. The gene expression data was log-normalized and each feature was scaled in the [0, 1] interval. For each patient, the features were scaled to have norm 1. We used the same version of the TCGA dataset as used by DRNet (Schwab et al., 2020) .The construction of this semi-synthetic dataset follows the construction of the IHDP Semi-Synthetic Dataset, except for the parametric UMAP algorithm reducing the dimensionality of the covariates from 4000 to 1, as an input to the dose-response function. Again, the model sees the high-dimensional covariates.For all experiments, we use Deep Ensembles (Lakshminarayanan et al., 2017) as our model, which can be seen as an approximate Bayesian inference method (Wilson & Izmailov, 2020) . The specific architectures and training details for each dataset are as follows:D.1 Synthetic Dataset• MLP with 2 hidden layers and 512 hidden units The policy-gradient method requires two additional hyperparameters: learning rates for variables that describe the unit and treatment. We performed a simple sweep over the parameters and found that learning rates of 0.01 and 0.001 for the unit and treatment, respectively, are relatively robust. We optimize these parameters using Adam.We used an internal cluster of 24 Nvidia Titan RTX cards for our experiments. The running time of individual experiments was relatively short, typically less than 4 hours each. However, an attempt to use the greedy algorithm on the TCGA dataset exceeded 14 hours before it was cancelled due to time constraints.In total, we estimate that across 4 datasets (synthetic discrete and continuous, IHDP, and TCGA), 4 methods, 3 acquisition sizes, and 10 seeds, approximately 480 experiments were conducted for this paper. With an average running time of 2 hours, this amounts to approximately 960 GPU-hours, along with associated CPU usage.This section provides a more detailed overview of related work in the fields of Bayesian Optimisation, Bayesian Optimal Experimental Design, Contextual Bandits, and Active Learning.Our setting is related to contextual Bayesian optimisation. Some of the more closely related methods include Profile Expected Improvement (Ginsbourger et al., 2014) , Multi-task Thompson Sampling (Char et al., 2019) , and conditional Bayesian optimization (Pearce et al., 2020) . All of these methods are, however, restricted to GPs and the criteria they use to choose optimal designs are not information-based. Our work differs from this family of methods, in that we aim to estimate the entire treatment-effect relationship, rather than merely the optimal treatment. Furthermore, Bayesian Optimisation typically operates in the setting in which one can repeat experiments on the unit of interest, whereas our work assumes that each unit can only be experimented on once.Our objective of selecting units/treatments bears the greatest similarity to experimental design. Experimental design aims to select the most informative experiments to conduct, given a limited budget. The difference in setting between experimental design and our work is that experimental design assumes that one can repeat experiments on the same unit, whereas our work assumes that each unit can only be experimented on once. (Foster et al., 2020) and MINEBED (Kleinegesse & Gutmann, 2020) , while the policy-based iDAD (Ivanova et al., 2021) applies to batch and adaptive settings. Our ability to handle discrete designs is another important distinction of our framework.Contextual bandits is another broad framework that our work is related to. An extensive line of research is focused on online linear bandits and discrete actions chosen using (variations of) UCB, Thompson sampling or ϵ-greedy strategy (Auer, 2002; Chu et al., 2011; Agrawal & Goyal, 2013; Han et al., 2020) . Krause & Ong (2011) instead model the reward as a GP defined over the context-action space and develop CGP-UCB. More recently, Zanette et al. (2021) proposed designing a batch of experiments offline to collect a good dataset from which to learn a policy.Deng et al. (2011) propose the use of Active Learning for recruiting patients to assign treatments that will reduce the uncertainty of an Individual Treatment Effect model. They focused on a small number of heterogeneous groups rather than individuals, with discrete treatments. As such, they tackle their setting through multi-armed bandits. While their objective remains the same as ours, in part our work could be seen as an extension of this work, with this work focusing on the individual level and operating in a setting with high-dimensional covariates. Atan et al. (2019) similarly focus on small heterogeneous groups with discrete treatments, and propose using Knowledge-Gradients to optimize the allocation process.Table 5 provides a comparison between various Adaptive Trial Methods and Active Learning approaches, highlighting the unique features of our proposed method.As shown in Table 5 , our framework is uniquely positioned to handle non-replaceable units, highdimensional covariates, both discrete and continuous treatments, and batch acquisition. This combination of features allows our approach to be more flexible and applicable to a wider range of experimental design scenarios compared to existing methods.Workshop on Bayesian Decision-making and Uncertainty, 38th Conference on Neural Information Processing Systems (NeurIPS 2024). 



Reviews: review: title: title: In this paper, the authors propose four approaches for experiment design. While it's an interesting contribution, I've identified several areas for improvement. I may have overlooked something, but I'm struggling to fully grasp the paper's primary contribution. I've included my comments below, and I hope the authors find them helpful.
main_review: main_review: In Section G.1, the authors state that Bayesian Optimization assumes repeated experiments. However, an alternative approach would be to remove units from the pool after use, similar to the algorithms in Sections B.1, B.2, and B.3. I believe Bayesian Optimization with Gaussian Processes could be applied to this problem unless I am missing something.

To enhance the readability of the plots in Appendix A.1, please include descriptive titles and labels. This will aid in visually differentiating the various densities and make the results more interpretable.

Throughout the paper, references to "our method" and "our dataset" can be confusing. Please explicitly state the specific method and dataset being used in each context. This will help readers understand the experiments and results more clearly.
 
The paper mentions that policy-based algorithms optimize "joint mutual information." However, the methodology section lacks details on how this is implemented. Please provide information on your approach that incorporates joint mutual information.

Based on the results in Tables 1-4, it appears that greedy and softmax methods often perform similarly or even better than policy-based methods. This raises the question of whether policy-based sampling is truly necessary.

The paper's primary contribution seems to be the development of a method to handle both continuous and discrete treatment options. This is achieved using Bayesian Neural Networks, which, from my understanding, treat discrete variables as a special case of continuous variables. Please clarify how your proposed methodology supports discrete and continuous treatment options.
title: title: Peer Review for Efficient Experimentation for Estimation of Continuous and Discrete Conditional Treatment Effects
main_review: main_review: ## 1. Summary
The paper proposes a framework aimed at improving the efficiency of experiments designed to estimate personalized treatment effects. The authors introduce an active learning approach to select the most informative unit-treatment combinations for experimentation in settings where each experimental unit can be treated only once. The authors employ Bayesian active learning techniques and introduce both greedy and policy gradient-based optimization strategies to enhance batch experimentation. The effectiveness of their approach is demonstrated through experiments on synthetic and semi-synthetic datasets.

## 2. Strengths and Weaknesses
### Originality
The originality of this paper is evidenced through the combination of active learning and treatment effect estimation, which is not widely represented in existing literature. While relevant literature is cited adequately, further elaboration on the unique contributions compared to previous studies could make it easier for readers to understand the original contribution of this paper better.

### Quality
The paper is of good quality, as the claims are both supported by clear theoretical formulation and empirical experiments. In addition, the paper provides tabled results and illustrative visuals to document the findings. 

However, there are three feedback/questions for authors' consideration -
1. Consider discussing the limitations of proposed methods. Are there any limitations of Greedy and Policy-Gradient based optimization, compared to Softmax Top-K (e.g., solve time, computational resources)?
2. Consider adding more explanation for the covariate-response curve
3. Is there a particular consideration for not including the results of Top-K method?

### Clarity
Overall, the paper is well-organized and clearly written. The structure allows readers to follow the logical progression from problem formulation to method implementation and results. Particularly, the paper summarized and discussed the key findings. Excellent work!

### Significance
While the results of this paper demonstrate improvements in experimentation efficiency, the paper could benefit from discussing more empirical significance of the findings.

It might be beneficial to add a section to discuss the potential real-world applications and implications. Which field/industry do you see this applied to?

## 3. Questions
1. What are the limitations of Greedy and Policy-Gradient based optimization, compared to Softmax Top-K (e.g., solve time, computational resources)?
2. Could the authors provide more insight into how the results would scale with larger datasets or different configurations of unit-treatment combinations?
3. How do the authors envision their methods being integrated into existing experimental designs in clinical trials or other fields?

## 4. Limitations
As discussed above, this paper could benefit from a more thorough discussion of 1) the limitations of proposed methods, 2) potential real-life applications and implications.

But all in all, the paper presents a compelling approach to improving experimentation efficiency for treatment effect estimation.