Extracted main text — title through conclusion, appendix excluded. This is what our citation measures are computed over, published so the extraction can be checked by eye.
46,350 characters · 22 sections · 29 citation commands
Personalized Assignment to One of Many Treatment Arms via Regularized and Clustered Joint Assignment Forests
We tackle the problem of learning an assignment policy that maps individual characteristics to one of potentially many treatment arms based on data from a randomized controlled trial. We consider a class of regularized forest-based algorithms that directly optimize for the outcome under assignment, propose a modification that also clusters treatment arms, document the performance of these approaches in a simulation exercise, and discuss extensions. Additionally, we show how having many treatment arms limits the prospects for best-arm identification and the performance of assignments based on separate arm-wise predictions, while achieving gains from better assignments remains feasible by optimizing for them directly.
Controlled trials in which treatment has been randomized across arms not only allow for the estimation of average effects, but they also provide an opportunity to learn which treatment works best for whom. A growing literature brings together methods from machine learning and causal inference to estimate heterogeneous treatment effects with flexible function forms athey2016recursive,athey2019generalized,wager2017estimation and by leveraging such information to optimize the assignment of treatments to individuals athey2020policy,Hitsch2018-bw,Sverdrup2020-ze. Other approaches directly optimize the utility of assignment of individuals to treatment Kallus2017-pr,Kitagawa2018-sp or the probability of assigning individuals to the best arm murphy2005generalization,zhou2018sequential. Yet many of these procedures focus on the case of a single treatment and control group zhou2018sequential,Sverdrup2020-ze,nie2021quasi,zhou2022offline,malearning.
In a number of clinical and behavior science settings, there might be a large suite of candidate interventions being tested for efficacy. Examples include psychological theory-informed nudges aimed at promoting gym visits or vaccine uptakes milkman2021mega, milkman2021megastudy or anti-depressants being considered for treatment of major depression ogawa2018efficacy. In such cases, existing methods that focus on settings with only a few treatment arms may not adapt well to the many-arms setting. For example, calculating personalized policies based on separate arm-wise treatment effect estimation may yield excessive variance in estimates and assignments.
In order to estimate personalized assignments with many treatment arms from a randomized trial, we focus on directly optimizing assignments using data across all arms jointly. Specifically, we consider a tree-based construction following the “personalization forest” proposed by Kallus2017-pr that pools information from all treatment arms and directly optimizes for the utility achieved under assignment Kitagawa2018-sp, in contrast to approaches that select models for the estimation of treatment effects and/or estimate outcomes of different arms separately. Like athey2019generalized, we estimate personalized treatment effects from the combination of trees in an honest way following athey2016recursive. We then leverage honest estimates from the training data to obtain an assignment rule.
To achieve better assignments for a large number of treatment arms, we augment our assignment forest in two ways. First, we propose a within-leaf regularization scheme that shrinks estimates towards leaf-wise averages. Second, we cluster treatment arms into groups based on the similarity of their estimated outcomes across units in the training sample. We then grow a regularized forest based on grouped treatment arms before recovering personalized assignments to one of the original treatment arms in a last step. In both cases, we pool information across arms, thus reducing the variance relative to estimating all arms separately. In its clustering version, we believe that our setup and approach is most similar to recent work in malearning, which proposes a supervised clustering approach via adaptive fusion in a parametric treatment-effect model, while we focus on unsupervised clustering in combination with a non-parametric regularized, honest, and jointly estimated random forest. Our approach also relates to banerjee2021selecting that pools treatments in order to select optimal nudges.
We compare the resulting algorithms to natural reference approaches in a simulation study. We show that direct optimization and regularization already lead to large improvements relative to benchmark methods that predict arm-wise outcomes separately, and to modest improvements over methods based on estimating separate heterogeneous treatment effects for all treatment arms relative to the control. We also document that adding clustering can improve assignment rules considerably even relative to the best non-clustered alternatives in our comparison group.
Our approach is motivated by the observation that a large number of arms makes it infeasible to consistently find the best arm and renders assignments based on separate arm-wise estimation inefficient. We make this claim precise in a theoretical illustration with many treatment arms where we show that best-arm identification becomes hard and assignment policies based on separate prediction of outcomes can perform poorly relative to joint assignment.
We introduce our setup in (ref) and describe the proposed algorithms in (ref). (ref) describes our simulation experiment and its results. We lay out our theoretical illustration in (ref). In (ref), we discuss extensions to more complex regularization and clustering schemes, as well as to non-experimental data. We conclude in (ref) by summarizing our findings and discussing important limitations and open questions for future research.
We consider data from a randomized experiment with $K+1$ treatment arms, and aim to estimate an assignment that maximizes out-of-sample outcomes for the same distribution of potential outcomes.
For treatments $k \in \{0,1,\ldots,K\}$, where we typically identify $k=0$ with the control arm, we let $Y^k \in \mathbbm{R}$ be the response of a given unit when assigned to treatment $k \in \{0,1,\ldots,K\}$. We also assume that there are features $X \in \mathcal{X}$ available that are not affected by treatment. Our goal is to find an assignment $a: \mathcal{X} \rightarrow \{0,1,\ldots,K\}$ such that the expected response
on new data drawn from the distribution of $(Y^0,\ldots,Y^K,X)$ is maximal, for the given assignment function $a$. (Hence, we assume that the potential outcomes $Y^k$ are in utility units.)
To find an assignment $a$, we assume that we have $n$ iid samples from an experiment available. In the experimental data, for each observation $i \in \mathcal{S} = \{1,\ldots,n\}$ treatment $T_i \in \{0,1,\ldots,K\}$ was randomly assigned independently of covariates (with for now fixed propensity scores $p^k, k \in \{0,1,\ldots,K\}$), and we observe data
From this data our goal is to estimate an assignment policy $a: \mathcal{X} \rightarrow \{0,1,\ldots,K\}$ that maximizes $U(a)$, where the optimal (but generally infeasible) assignment policy is given by $a^*(X) = \operatorname*{arg\,max}_a \operatorname{E}[Y^a|X]$.
We provide an algorithm that obtains an assignment $\hat{a}: \mathcal{X} \rightarrow \{0,1,\ldots,K\}$ for $\mathcal{X} = \mathbbm{R}^d$. One natural approach to constructing such an algorithm would be to estimate arm-wise outcomes $\operatorname{E}[Y^a|X]$ or treatment effects $\operatorname{E}[Y^a - Y^0|X]$ separately, and then assign an individual with characteristics $X$ to the arm with the highest estimated outcome or treatment effect. However, such an approach would yield excess variance, since each separate estimation would only use limited data. In addition, an algorithm optimized for the precise estimation of arm-wise outcomes or arm-wise treatments effects may not be suboptimal for the related, but different goal of finding an assignment that yields high utility.
Instead, we consider an approach that optimizes directly for an optimal assignment across all treatment arms and thereby pools all data. Specifically, from the training data we obtain joint assignment trees through recursive partitioning, which we then combine into a single joint assignment forest through bagging. Relative to the construction of the personalized forests from Kallus2017-pr, our implementation adapts honest estimation from athey2019generalized and specifically targets the challenge of many treatment arms by integrating regularization and clustering in the construction of trees and honest estimation of leaf-wise treatment effects.
For a given training sample of size $n$, we follow Kallus2017-pr in fitting a tree by recursively maximizing the regularized empirical analogue of the utility $U(\hat{a})$ for an assignment $\hat{a}$ determined by leaf-wise maximizers.
\paragraph{Regularized leaf-wise estimation and assignment for a given tree.}
Given a tree that partitions the covariate space $\mathbbm{R}^d$ into leaves, we calculate an assignment $a_\ell$ for leaf $\ell$ as the maximizer of the regularized arm-wise within-leaf averages $\hat{Y}^k_\ell$. Specifically, we write
for the arm-wise average outcomes and counts within a leaf in the training data. Here, $\overline{Y}^k_{\ell}$ is an unbiased estimator of $\mathrm{E}[Y^k|X \in \ell]$.
The arm-wise averages may be noisy, especially for those arms that have only a few observations in a given leaf. In a departure from the existing literature, we shrink arm-wise average outcomes $\overline{Y}^k_{\ell}$ towards the overall (weighted) average $\overline{Y}_\ell$ of leaf outcomes (which can be motivated by a homoscedastic Normal means model similar to (ref)) and set
and then choose the assignment $a_\ell = \operatorname*{arg\,max} \hat{Y}^k_\ell$. This regularization scheme reduces the probability that arms with a few draws that are spuriously high are chosen over arms with a high number of draws and a high true average. The regularization parameter $\lambda_1$ controls the amount of shrinkage; when $\lambda_1 = 0$, $\hat{Y}^k_\ell = \overline{Y}^k_{\ell}$, and we directly maximize the empirical outcome. We briefly discuss extensions that embrace heteroscedasticity and shrinking towards overall arm-wise averages in (ref) below.
\paragraph{Recursive splitting criterion.} For the leaf-wise assignments $a_\ell$, we recursively split a leaf $\ell_p$ into $\ell_p = \ell_1 \cup \ell_2$ along splits $x_j \leq c$ by maximizing
for $\hat{U}$ one of the estimators
of the utility achieved by leaf-wise assignment, where $ P^k = \frac{N^k}{N} = \frac{\sum_{T_i=k} 1}{N} $ and $N_\ell = \sum_{X_i \in \ell} 1$. In deciding on a split, we consider only splits for which:
If no such splits exist, we do not split the leaf $\ell_p$. This greedy algorithm yields a partition of $\mathbbm{R}^d$.
We obtain an assignment forest by bagging many trees as in Kallus2017-pr, and estimating honest estimates of the treatment-specific counterfactual outcomes on the training sample following wager2018estimation.
\paragraph{Bagging with treatment-arm randomization.} Instead of estimating a single tree, we obtain $M$ trees from the training sample by repeatedly drawing a bootstrap sample of size $\lceil\beta n\rceil$, without replacement, for $\beta \in (0,1)$ and repeating the above procedure on the chosen sample. We stratify each bootstrap sample by treatment arms to ensure that each draw represents the overall fraction of different treatments. When fitting the trees, we also at every split:
\paragraph{Honest estimation.} For every point $x \in \mathbbm{R}^d$ and every tree $m \in \{1,\ldots,M\}$, we follow athey2019generalized in obtaining honest estimates of the treatment-specific conditional potential outcome $E[Y^k|X]$, where we estimate expected outcomes only from data that the same tree was not fit on. Specifically, denote by $\mathcal{S}_m$ the bootstrap sample $m$ was fit on, and write $\ell_m(X) \subseteq \mathbbm{R}^d$ for the leaf that $x$ falls into. Then we let for all arms $k \in \{0,\ldots,K\}$
and set
Here, we allow the shrinkage parameter to differ between the construction of individual trees ($\lambda_1$) and the final, honest estimates of arm-specific outcomes ($\lambda_2$). This distinction allows us, for example, setting a lower $\lambda_2$ to avoid over-smoothing in constructing these final estimates (by choosing a lower $\lambda_2$), which are averaged over a large number of trees, while also choosing a higher $\lambda_1$ to avoid overfitting in the construction of individual trees.
\paragraph{Aggregation and assignment.}
Given tree-wise honest estimates $\hat{f}^k_m(X)$ at a new sample point $x \in \mathbbm{R}^d$, we estimate the conditional potential outcome $E[Y^k|X]$ by the average
We then obtain the assignment
\paragraph{Tuning parameters.} The parameters $\nu$ (minimal leaf size), $\lambda_1$ (within-leaf shrinkage when growing the tree), $\lambda_2$ (within-leaf shrinkage when estimating), $\varepsilon$ (minimal gain in objective), $\beta$ (fraction sampled for each tree), $\delta$ (fraction of covariates considered at each split), $\kappa$ (fraction of treatments considered at each split) are the tuning parameters of the assignment forest that control the complexity of the procedure, which can be chosen by cross-validation. The way of calculating $\hat{U}$ ($N$ vs $P$) can be seen as another tuning choice.
Some of the variation in outcomes is common across treatment arms. To the degree that this common variation can be predicted, we can reduce the variance in the evaluation of different treatment arms by subtracting such common variation. We therefore add a pre-processing step to our algorithm.
\paragraph{Arbitrary residualization.} Assume we had some fixed function $\bar{f}: \mathbbm{R}^d \rightarrow \mathbbm{R}$ available. Then ranking between assignments for potential outcomes $\tilde{Y}^k = Y^k - \bar{f}(X)$ are the same as in the original distribution, since
where $\mathrm{E}[\bar{f}(X)]$ does not vary with $a$. We can therefore estimate an assignment on data $(Y_i - \bar{f}(X_i),T_i,X_i)_{i=1}^n$, where we choose $\bar{f}$ to reduce the variance of the outcome.
\paragraph{Choices of baseline $\bar{f}$.} We consider three (oracle) choices for a function $\bar{f}: \mathbbm{R}^d \rightarrow \mathbbm{R}$ to reduce variation in estimating differential assignments:
\paragraph{Implementation.} Following e.g. Wager:2016dz,Wu:2017wj, we fit baseline prediction functions $\widehat{f}_{-i}$ that estimate $\bar{f}$ in the training dataset using cross-fitting to avoid biases from overfitting. We then run the forest algorithm on the residualized outcomes $Y_i - \widehat{f}_{-i}(X_i)$ where $\widehat{f}_{-i}$ does not use data from observation $i$. As our main implementation, we solve the weighted prediction problem $$\mathrm{E}\left[ \left( Y - f(X) \right)^2 / (p^T)^2 \right] \rightarrow \min_{f}$$ spiessoptimal using a random forest, which implements the third option (see (ref)). We do not residualize on the held-out data, since residualization shifts the absolute policy value of the assignment, which may be of interest. Our algorithm is summarized as (ref).
While simple regularization within arms reduces variation in model construction due to random outliers, it is not able to capture systematic similarity between specific arms. In this section, we therefore propose a simple clustering scheme that partitions the $K+1$ arms into $M+1$ groups $G$, where $\bigcup_{g \in G} g = \{0,\ldots,K\}$. Like Bonhomme2015-tf and Bonhomme2017-pd in their study of unobserved heterogeneity, we employ a $k$-means algorithm for clustering the $K+1$ arms into $M+1$ groups. Specifically, we expand the assignment forest from the previous section as follows:
Our proposed algorithm thus uses a clustering of treatment arms when constructing the assignment trees in Step (ref), but still provides arm-specific estimates from the resulting forest in Step (ref).
We believe that our approach is most similar to the recent clustering approach of malearning, which operates within the same framework, but considers a semi-parametric model with parametric treatment effects. Unlike our unsupervised clustering scheme, malearning leverages a fusion penalty term to obtain supervised clustering.
We conduct a simulation study to examine the performance of the algorithms described above.
For $X \sim \mathcal{N}(\mathbf{0}_d,\mathbbm{I}_d), \varepsilon \sim \mathcal{N}(0,1)$, $T \sim \mathcal{U}(\{0,1,\ldots,K\})$ with $d=3$ we generate outcomes by
In this $K+1$-treatment-arm setting (with choice of two parameters, the strength of treatment effects $\gamma$ and residual noise $\sigma$), the optimal treatment choices are $T=K$ (for $X_3 > 0$) and $T=1$ (otherwise), realizing an average outcome of $\gamma$. Average outcome for assigning to control is $-10$, for assigning to the global best is $0$, and for assigning randomly among one of the $K$ non-control treatments is also $0$. We run simulations for treatment effect strengths $\gamma \in \{10,20\}$ and noise level $\sigma \in \{10,20\}$ with $K \in \{9,29,49,99\}$ arms on samples of size $n = 5{,}000$.
We choose a simple setting with joint baseline and treatment-effect structure to clearly highlight the advantages of the joint assignment approach. Note, however, that treatment arms are not clustered and all have different effects on the outcome.
For our three-covariate simulation setup, we tune the following parameters: $\nu \in \{3,5\}$ (minimal leaf size), $\lambda_1 \in \{0,0.5,1\}$ (within-leaf shrinkage when growing the tree), $\lambda_2 \in \{0,0.5,1\}$ (within-leaf shrinkage when estimating), $\varepsilon \in \{0.5,1,2\}$ (minimal gain in objective), $\kappa \in \{0.5, 0.8, 1\}$ (fraction of treatments considered at each split), and choose the combination of these parameters for each arm-“noise” setting by cross-validation. We use a test set of size 10,000 for arriving at our final estimates. We note that it may make sense to use different regularization parameters across Steps (ref), (ref), and (ref) in (ref). In our simulations, we used the same regularization parameters ($\lambda_1$ and $\lambda_2$ before and after clustering) and found that works well.
We compare the performance of our regularized and clustered approach to two natural approaches to determining treatment assignments. First, we fit $K+1$ random forests separately for each control and treatment arm $k \in \{0,1,\dots, K\}$, and predict $\hat{Y}_k = \hat{f}_k(X)$ for each observation under consideration. The assignment rule is obtained as $\hat{a}(X) = \operatorname*{arg\,max} _{k \in \{0,1,\dots, K\}} \hat{f}_{k}(X)$. Second, we compare the performance of our approach to a multi-arm causal forest athey2019generalized,nie2021quasi that jointly estimates treatment effects $\hat{\tau}_k(X)$ for all $K$ treatment arms $k \in \{1,\ldots,K\}$. As an assignment rule for the multi-arm causal forest, we use $\hat{a}(X) = \operatorname*{arg\,max}_{k \in \{0,1,\dots, K\}} \hat{\tau}_{k}(X)$, where we set $\hat{\tau}_{0}(X) = 0$.
(ref) presents results for the optimal value function from 500 simulations for a 50-treatment arm design for the following three settings: a. $\gamma = 10, \sigma = 10$ (hereby referred to as regular setting), b. $\gamma = 20, \sigma = 10$ (”low noise” setting), and c. $\gamma = 10, \sigma = 20$ (“high noise setting”). (ref) presents results for the optimal value function from 500 simulations of the ”high noise” ($\gamma = 10, \sigma = 20$) setting for 10, 30, 50, and 100 treatment arms. We present corresponding results of the assignment rate in Figures (ref) and (ref) in the Appendix. Here, `Oracle Optimal Assignment' refers to following the optimal assignment rule, based on the ground truth in the simulation. `Random Assignment' refers to assigning units in every simulation randomly across treatment arms; `Global Best Assignment' refers to assigning units in every simulation the on-average best performing treatment.
We note that our joint assignment approaches (un-clustered and clustered) outperform the separate random forests approach in recovering the oracle outcome under optimal assignment for all $\gamma,\sigma$ settings, across all considered treatment arms (10, 30, 50, 100). The clustered DOF approach outperforms the un-clustered approach in recovering the oracle optimal outcome across all settings, making a strong case for treatment arm clustering under many arms. Compared to the multi-arm causal forest approach, our clustered approach outperforms it in recovering the optimal outcome across the “regular” and “high noise” settings, while delivering similar performance in a “low noise” setting (Figure (ref)). The resilience of our joint assignment approach compared to the multi-arm causal forest is particularly prominent in `high noise' settings (Figure (ref)).
Figures (ref) and (ref) in the Appendix demonstrate the performance of our clustered and unclustered approaches against the separate forest and multi-arm causal forest on the average (across 500 simulations) of the successful assignment rate to the oracle treatment assignment rule. In settings with fewer arms, we find our approach to perform well (Figure (ref)), but as the number of arms increases, both our unclustered and clustered joint assignment forest approaches perform worse than the “global best assignment” rule. While the algorithms perform better in “low noise” settings (for large arm settings), in “high noise” ones, it approaches assignment rates akin to random assignment (ref)).
\FloatBarrier
One motivation for our regularized joint optimization approach is that finding optimal treatment assignments becomes hard and separate estimation inefficient when there are many treatment arms. In this section, we theoretically illustrate these features.
Specifically, we consider different goals and procedures when assigning treatments in a small example. One natural goal when assigning treatments is maximize the probability of the chosen arm for a randomly chosen unit being the best arm for that unit. In (ref), we show in a high-dimensional Normal example that this goal of best-arm identification becomes hard when the number of arms increases, while assigning to maximize utility still can yield non-trivial solutions over the benchmark of random assignment. We therefore focus instead on the goal of optimizing for a treatment rule with high average outcomes.
A natural procedure to achieving good treatment assignments is to predict each treatment arm's outcome separately and then assign a given unit to the treatment with the highest predicted outcome. In (ref), we compare this method of assigning by estimating the outcomes of different treatment arms separately to optimizing directly for an optimal assignment in the same limiting regime, and show that except for edge cases the former performs strictly worse than the latter even in the limit. Together, the illustrations in this section motivate our focus in direct utility maximization over best-arm identification and the estimation of separate arms.
Comparing how hard goals are to achieve and how to best achieve them faces two related hurdles: First, procedures may not only differ in their target loss function, but also in the function class, regularization, and optimization method; and second, performance depends on the true distribution and there may not be a universally best solution, even holding e.g. the function class fixed. We overcome these challenges by a standard approach from statistical decision theory Wald1950-qc: we consider a distribution over the true state of the world, and compare the average performance for optimal (Bayes) solutions to each of the problems. This way we can abstract from the specific algorithm employed for each purpose and focus instead of the different optimization goals.
For tractability, we consider a simple homoscedastic baseline model with $\mathcal{X} = \{1,\ldots,N\}$ and
We assume that the means $\mu^k_j$ are themselves jointly distributed according to a multivariate Normal prior that is invariant to permutations in treatment arms and covariates identities, capturing the idea that these are ex-ante indistinguishable. This allows us to write $\mu^k_j = \alpha + \beta_j + \gamma^k + \delta^k_j$ with
where all draws are independent. We also assume that all cells $T=k, X=j$ have the same sample size $\frac{n}{(K+1) N}$. We study this general model in (ref); in this section, we assume throughout that $a,b,c,d, \sigma^2 > 0$.
For analyzing the case of many treatment arms in a high-dimensional setting, we consider the case where
and both the case of fixed (but potentially high) $K$ and the case where $K \rightarrow \infty$. This approximation regime represents a case where the number of covariate cells and treatment arms increases fast enough relative to the sample size for the estimation of the cell-wise means $\mu_j^k$ to be hard even in the limit. This is motivated by cases such as the mega-study in our application, where some of the many treatment arms have only a few hundred observations across all covariate values.
We compare two different goals when learning an assignment $\hat{a}: \mathcal{X} \rightarrow \mathbbm{R}$:
In our homoscedastic setting, the average outcomes for both goals are maximized by an assignment rule $\hat{a}$ that picks, for each cell, the arm with the highest posterior expectation. In (ref) we compare this assignment to the infeasible oracle $a^*(j) = \operatorname*{arg\,max}_k \mu_j^k$ (as an upper bound for performance) and the assignment $\underline{a}$ that assigns treatments randomly (as a lower bound for performance), and show:
With many treatment arms, finding the best arm hence becomes hard (even with an optimal procedure), while recovering a sizable fraction of the utility gain from personalization remains feasible. The reason is simple: with more and more arms, we may make selection mistakes in picking a similar arms with slightly lower mean; but that arm likely still has a higher-than-average mean. Indeed, in (ref) we show that the fraction of optimal utility recovered does not change with $K$ in our asymptotic approximation.
We now focus maximizing utility $\mathrm{E}[Y^{a(X)}]$, and consider two natural methods of achieving this goal:
While the former strategy is compelling in practice since it can be readily implemented by using separate prediction algorithms across treatment arms, we show in (ref) that this strategy is suboptimal in our high-dimensional asymptotic approximation:
When separate prediction algorithms only use data from a single arm, they may misattribute variation in the baseline that is common to all treatment arms with a given covariate value, and therefore overfit to individual cell outcomes. In a high-dimensional limit, this under-performance does not go away.
In this section, we discuss extensions to the baseline algorithm discussed in (ref).
The above algorithm assumes that treatments are assigned randomly with constant propensity score. But we can easily modify the algorithm for a known propensity score $p^k(X)$. If we also care about weighted outcomes
for some given positive weights $v(X)$, then two natural (unregularized) estimators of leaf-wise utility are
The latter estimator is unbiased for the utility of assigning the leaf accordingly by standard inverse-propensity weighting, yielding an overall unbiased estimate of the associated policy as in Hitsch2018-bw. We obtain the criteria in (ref) from $v \equiv 1$, $p^k \equiv P^k$.
In (ref), we estimate arm-wise average outcomes within leaves using regularization to avoid noise when there are only a few observations in a given arm. The regularization scheme implicitly assumes a homoscedastic Normal-means model with a Normal prior on leaf-specific arm-wise averages and an uninformative hyperprior (corresponding to an empirical-Bayes strategy). The tuning parameter $\lambda$ corresponds to the ratio of the unit-specific variance to the variance of the Normal prior. As a more complete treatment, we can also include shrinkage towards the overall arm-wise average across leaves. Further, we can estimate arm-wise variances of outcomes to refine shrinkage.
In (ref), we consider an ad-hoc $k$-means clustering scheme. We could instead consider a Bayesian model similar to (ref) that assumes that some arms are more similar to each other than others. Such a model would generalize the shrinkage scheme in (ref) by shrinking arms more towards those in the same group or with higher similarity. We could then estimate such a grouped shrinkage scheme with a $k$-means algorithm as in (ref) or with an estimation of arm similarity corresponding to covariances of arm-wise means in the Bayesian model.
Our current cluster scheme starts by clustering treatment arms into groups based on a simple, non-clustered run of the assignment-forest algorithm. Alternatively, we could iterate the prediction and clustering steps to refine the clustering, either starting with a non-clustered or a randomly clustered assignment. In addition, when clustering we could take into account whether an arm is likely to be chosen for treatment, and ensure that clusters capture similarity mainly for those observations-specific arms that are likely to affect assignment. Finally, we could maintain the clusters for assignment, assuming that within-group assignments are random.
In this article, we consider learning treatment assignments from experimental data with many treatment arms. We demonstrate the limits of estimating optimal treatment arms and recovering effective treatment assignment policies from separate arm-wise outcome predictions or treatment-effect estimates. As a feasible alternative, we provide a regularized tree-based algorithm that directly optimizes for treatment assignment, clusters treatment arms, and document its properties in a simulation study.
Our current analysis remains limited to experiments with known probabilities of assignment to different arms. When treatment assignment is endogenous and propensity scores are unknown, then their estimation poses additional challenges that are beyond the scope of this article. Similarly, we focus on existing experimental data, and do not consider optimal experimentation of the dynamic allocation to treatment arms.
The use of targeting rules, like those obtained from our algorithm, has the potential to improve utility through better allocation, but also comes with substantial risks when data is biased or personalization may reinforce or increase inequities. Adding fairness and equity constraints to the resulting treatment-assignment rules can be an important future extension.