Optimal Individualized Treatment Rule For Combination Treatments Under Budget Constraints

In this section, we introduce the detailed modeling strategy for treatment encoder $\beta(\cdot)$. The treatment effects of combination treatments can be decoupled into two components: additive treatment effects, which is the sum of treatment effects from single treatments in combination; and interaction effects, which are the additional effects induced by the combinations of multiple treatments. Therefore, we formulate the treatment encoder as follows:

where $\beta_{0}(\mathbf{A})$ and $\beta_{1}(\mathbf{A})$ are additive and interactive treatment encoders, respectively. In particular, $\beta_{0}(\mathbf{A})$ is a linear function with respect to $\mathbf{A}$, where $\mathbf{W}=(\mathbf{W}_{1},\mathbf{W}_{2},...,\mathbf{W}_{K})\in\mathbb{R}^{r\times K}$ and $\mathbf{W}_{k}$ is the latent representation of the $k$th treatment. As a result, $\alpha(\mathbf{X})^{T}\beta_{0}(\mathbf{A})=\sum_{k:\{A^{k}=1\}}\mathbf{W}_{k}^{T}\alpha(\mathbf{X})$ are the additive treatment effects of the combination treatment $\mathbf{A}$. The constraints for $\beta_{1}(\cdot)$ ensures the identifiability of $\beta_{0}(\cdot)$ and $\beta_{1}(\cdot)$ such that any representation $\beta(\mathbf{A})$ can be uniquely decoupled into $\beta_{0}(\mathbf{A})$ and $\beta_{1}(\mathbf{A})$.

The interaction effects are challenging to estimate in combination treatments. A naive solution is to assume that interaction effects are ignorable, which leads the additive treatment encoder $\beta_{0}(\mathbf{A})$ to be saturated in estimating the treatment effects of combination treatments. However, interaction effects are widely perceived in many fields such as medicine (Stader et al., 2020; Li et al., 2018), psychology (Caspi et al., 2010), and public health (Braveman et al., 2011). Statistically, ignoring the interaction effects could lead to inconsistent estimation of the treatment effects (Zhao and Ding, 2023) and the ITR (Liang et al., 2018). Hence, it is critical to incorporate the interaction effects in estimating the ITR for combination treatments.

A straightforward approach to model the interactive treatment encoder $\beta_{1}(\mathbf{A})$ is similar to the additive treatment encoder $\beta_{0}(\mathbf{A})$, which we name as the treatment dictionary. Specifically, a matrix $\mathbf{V}=(\mathbf{V}_{1},\mathbf{V}_{2},...,\mathbf{V}_{|\mathcal{A}|})\in\mathbb{R}^{r\times|\mathcal{A}|}$ is a dictionary that stores the latent representations of each combination treatment so that $\beta_{1}(\mathbf{A})$ is defined as follows

where $\mathbf{e}_{\tilde{A}}$ is the one-hot encoding of the categorical representation of $\mathbf{A}$. Since the number of possible combination treatments $|\mathcal{A}|$ could grow exponentially as $K$ increases, the parameters of $\mathbf{V}$ could also explode. Even worse, each column $\mathbf{V}_{l}$ can be updated only if the associated treatment $\tilde{A}_{l}$ is observed. Given a limited sample size, each treatment could be only observed a few times in the combination treatment scenarios, which leads the estimation efficiency of $\mathbf{V}$ to be severely compromised. The same puzzle is also observed in other methods. For the Q-function in(4), the parameters in $\delta_{l}(\mathbf{X})$ can be updated only if $\tilde{A}_{l}$ is observed; In the Treatment-Agnostic Representation Network (TARNet) (Shalit et al., 2017) and the Dragonnet (Shi et al., 2019), each treatment is associated with an independent set of regression layers to estimate the treatment-specific treatment effects, which results in inefficiency estimation for combination treatment problems.

In order to overcome the above issue, we propose to utilize the feed-forward neural network (Goodfellow et al., 2016) to learn efficient latent representations in the $r$-dimensional space. Specifically, the interactive treatment encoder is defined as

where $\mathcal{U}_{l}(\mathbf{x})=\mathbf{U}_{l}\mathbf{x}+\mathbf{b}_{l}$ is the linear operator with the weight matrix $\mathbf{U}_{l}\in\mathbb{R}^{r_{l}\times r_{l-1}}$ and the biases $\mathbf{b}_{l}$. The activation function is chosen as ReLU function $\sigma(\mathbf{x})=\max(\mathbf{x},0)$ in this paper. An illustration of the neural network interactive treatment encoder is shown in Figure 1. Note that all parameters in (9) are shared among all possible treatments, so all of the weight matrices and biases in (9) are updated regardless of the input treatment, which could improve the estimation efficiency, even though (9) may include more parameters than the treatment dictionary (8). As a result, the double encoder model with (9) not only guarantees a faster convergence rate (with respect to $K$) of the value function but also improves the empirical performance especially when $K$ is large or sample size $n$ is small, which will be shown in numerical studies and real data analysis. A direct comparison of the neural network interactive treatment encoder (9) and the treatment encoder (8), the additive model (4), TARNet (Shalit et al., 2017) and Dragonnet (Shi et al., 2019) are also shown in Figure 1.

Although the interactive treatment encoder (9) allows an efficient estimation, it is not guaranteed to represent up to $|\mathcal{A}|$ interaction effects. In the treatment dictionary (8), columns $\mathbf{V}_{l}$’s are free parameters to represent $|\mathcal{A}|$ treatments without any constraints. However, an “under-parameterized” neural network is not capable of representing $|\mathcal{A}|$ treatments in $r$-dimensional space. For example, if there are three treatments to be combined ($K=3$), and the treatment effects are sufficiently captured by one-dimensional $\alpha(\mathbf{X})$ with different coefficients ($r=1$). We use the following one-hidden layer neural network to represent the treatment in $\mathbb{R}$:

where $u_{2},b_{2},b_{1}\in\mathbb{R}$ are scalars and $\mathbf{U}_{1}\in\mathbb{R}^{1\times 3}$. In other words, the hidden layer only includes one node. In the following, we show that this neural network can only represent restricted interaction effects:

The proof of Proposition 1 is provided in the supplementary materials. Based on the above observation, it is critical to guarantee the representation power of $\beta_{1}(\mathbf{A})$ to incorporate flexible interaction effects. In the following, we establish a theoretical guarantee of the representation power of $\beta_{1}(\cdot)$ under a mild assumption on the widths of neural networks:

The above result is adapted from the recent work on the memorization capacity of neural networks (Yun et al., 2019; Bubeck et al., 2020). Theorem 2 shows that if there are $\Omega(2^{K/2}r^{1/2})$ hidden nodes in neural networks, then it is sufficient to represent all possible interaction effects in $\mathbb{R}^{r}$. However, obtaining the parameter set $\{\mathbf{U}_{l},\mathbf{b}_{l},l=1,2,3\}$ in Theorem 2 via the optimization algorithm is not guaranteed due to the non-convex loss surface of the neural networks. In practice, the neural network widths in Theorem 2 can be a guide, and choosing a wider network is recommended to achieve better empirical performance.

In summary, we propose to formulate the treatment encoder as two decoupled parts: the additive treatment encoder and the interactive encoder. We provide two options for the interactive treatment encoder: the treatment dictionary and the neural network, where the neural network can improve the asymptotic convergence rate and empirical performance with guaranteed representation power. In the numerical studies, we use the neural network interactive treatment encoder for our proposed method, and a comprehensive comparison between the treatment dictionary and the neural network is provided in the supplementary materials.

As we introduced in (6), the covariates encoder $\alpha(\cdot):\mathcal{X}\rightarrow\mathbb{R}^{r}$ constitutes the function bases of the treatment effects for all combination treatments. In other words, the treatment effects represented in (6) lie in the space spanned by $\alpha^{(1)}(\mathbf{X}),...,\alpha^{(r)}(\mathbf{X})$. Therefore, it is critical to consider a sufficiently large and flexible function space to accommodate the highly complex treatment effects and avoid possible model misspecification. In particular, we adopt three nonlinear or nonparametric models for covariates encoders: polynomial, B-Spline (Hastie et al., 2009), and neural network (Goodfellow et al., 2016).

First of all, we introduce the $\alpha(\mathbf{X})$ as a feed-forward neural network defined as follows:

$\mathcal{T}_{l}(\mathbf{x})=\mathbf{T}_{l}\mathbf{x}+\mathbf{c}_{l}$ is the linear operator with the weight matrix $\mathbf{T}_{l}\in\mathbb{R}^{r_{l}\times r_{l-1}}$ and the biases $\mathbf{c}_{l}$. The activation function is chosen as ReLU function $\sigma(\mathbf{x})=\max(\mathbf{x},0)$ in this paper. Note that the depth and the width of the covariates encoder $\alpha(\cdot)$ are not necessarily identical to those of the interactive treatment encoder $\beta_{1}(\cdot)$, and these are all tuning parameters to be determined through hyper-parameter tuning.

Even though neural networks achieve superior performance in many fields, their performance in small sample size problems, such as clinical trials or observational studies in medical research, is still deficient. In addition, neural networks lack interpretability due to the nature of their recursive composition; therefore, the adoption of neural networks in medical research is still under review. Here, we propose the polynomial and B-Spline covariates encoders to incorporate nonlinear treatment effects for better interpretation. For the polynomial covariates encoder, we first expand each covariate $x_{i}$ into a specific order of polynomials $(x_{i},x_{i}^{2},...,x_{i}^{d})$ where $d$ is a tuning parameter. Then we take the linear combinations of all polynomials as the output of the covariate encoders. Figure 2 provides an example of the polynomial covariate encoder with $d=3$. Similarly, as for the B-spline covariates encoder, we first expand each covariate into B-spline bases, where the number of knots and the spline degree are tuning parameters. Likewise, linear combinations of these B-spline bases are adopted as the output of the encoder. Although both polynomial and B-spline covariate encoders can accommodate interaction terms among the polynomial bases or B-spline bases for a better approximation for multivariate functions, exponentially increasing parameters need to be estimated as the dimension of covariates or the degree of bases increases. In the interest of computation feasibility, we do not consider interaction terms in the following discussion.

First of all, we propose the following doubly robust estimator for the treatment effects:

where $\hat{\mathbb{P}}(\mathbf{A}_{i}|\mathbf{X_{i}})$ is a working model of the propensity score specifying the probability of treatment assignment given pre-treatment covariates, and $\hat{m}(\mathbf{X}_{i})$ is a working model of treatment-free effects. The inverse probability weights given by the propensity scores balance the samples assigned to different combination treatments, assumed to be equal under the randomized clinical trial setting. By removing the treatment-free effects $m(\mathbf{x})$ from the responses before we estimate the treatment effects, the numerical stability can be improved and the estimator variance is reduced. This is also observed in (Zhou et al., 2017; Fu et al., 2016). Furthermore, the estimator in (14) is doubly-robust in that if either $\hat{\mathbb{P}}(\cdot|\cdot)$ or $\hat{m}(\cdot)$ is correctly-specified, $\hat{\alpha}(\cdot)^{T}\hat{\beta}(\cdot)$ is a consistent estimator of the treatment effects, and a detailed proof is provided in the supplemental material. This result extends the results in (Meng and Qiao, 2020) from binary and multiple treatments to combination treatments. Empirically, we minimize the sample average of the loss function (14) with additional penalties: for the additive treatment encoder, $L_{2}$ penalty is imposed to avoid overfitting; for the interactive treatment encoder, $L_{1}$ penalty is added since the interaction effects are usually sparse (Wu and Hamada, 2011).

In this work, the working model of the propensity score is obtained via the penalized multinomial logistic regression (Friedman et al., 2010) as a working model. Specifically, the multinomial logistic model is parameterized by $\gamma_{1},\gamma_{2},...,\gamma_{2^{K}}\in\mathbb{R}^{p}$:

The parameters $\gamma_{k}$’s can be estimated by maximizing the likelihood:

where the group Lasso (Meier et al., 2008) is used to penalize parameters across all treatment groups. A potential issue of propensity score estimation is that the estimated probability could be negligible when there are many possible treatments, which leads to unstable estimators for treatment effects. To alleviate the limitation on inverse probability weighting, we stabilize the propensity scores (Xu et al., 2010) by multiplying the frequency of the corresponding treatment to the weights.

For the estimation of the treatment-free effects $m(\cdot)$, we adopt a two-layer neural network:

where $\mathbf{w}^{2}_{m}\in\mathbb{R}^{h}$ and $\mathbf{W}_{m}^{1}\in\mathbb{R}^{h\times p}$ are weight matrices, and $\sigma(x)$ is the ReLu function. The width $h$ controls the complexity of this model. The weight matrices are estimated through minimizing:

Given working models $\hat{m}(\mathbf{x})$ and $\hat{\mathbb{P}}(\mathbf{a}|\mathbf{x})$ for treatment-free effects and propensity scores, we propose to optimize the double encoder alternatively. The detailed algorithm is listed in Algorithm 1. Specifically, we employ the Adam optimizer (Kingma and Ba, 2014) for optimizing each encoder: covariate encoder $\alpha(\mathbf{x})$, additive treatment encoder $\beta_{1}(\mathbf{a})$, and non-parametric treatment encoder $\beta_{2}(\mathbf{a})$. To stabilize the optimization during the iterations, we utilize the exponential scheduler (Patterson and Gibson, 2017) which decays the learning rate by a constant per epoch. In all of our numerical studies, we use 0.95 as a decaying constant for the exponential scheduler. To ensure the identifiability of treatment effects, we also require a constraint on the treatment encoder $\beta(\cdot)$ such that $\sum_{\mathbf{a}\in\mathcal{A}}\beta(\mathbf{a})=0$. To satisfy this constraint, we add an additional normalization layer before the output of $\beta(\cdot)$. The normalization layer subtracts the weighted mean vector where the weight is given by the reciprocal of the combination treatment occurrence in the batch. Since this operation only centers the outputs, the theoretical guarantee for $\beta(\cdot)$ in Section 4 still holds. Once our algorithm converges, we obtain the estimation for $\alpha(\cdot)$ and $\beta(\cdot)$, and also the estimated individualized treatment rule $\hat{d}(\cdot)$ by (3).

In addition, successful neural network training usually requires careful hyper-parameter tuning. The proposed double encoder model includes multiple hyper-parameters: network structure-related hyper-parameters (e.g., network depth $L_{\alpha}$ and $L_{\beta}$, network width $r_{\alpha}$ and $r_{\beta}$, encoder output dimension $r$), optimization-related hyper-parameters (e.g., additive treatment encoder penalty coefficients $\lambda_{a}$, interactive treatment encoder penalty coefficients $\lambda_{i}$, mini-batch size $B$, learning rate $\eta$, and training epochs $E$). These hyper-parameters induce an extremely large search space, which makes the grid search method (Yu and Zhu, 2020) practically infeasible. Instead, we randomly sample 50 hyper-parameter settings in each experiment over the pre-specified search space (detailed specification of hyper-parameter space is provided in the supplementary materials), and the best hyper-parameter setting is selected if it attains the largest value function on an independent validation set. Furthermore, due to the non-convexity of the loss function, the convergence of the algorithm also relies heavily on the parameter initialization. In the supplementary materials, we provide detailed analyses for numerical results under different parameter initializations.

In the following, we introduce our procedure for the budget-constrained individualized treatment rule for combination treatment (12) estimation. We use the plug-in estimates $\hat{\alpha}(\mathbf{x}_{i})^{T}\hat{\beta}(\mathbf{a}_{j})$ for $\delta_{ij}$, and calculate the cost for each combination treatment from the cost vector $\mathbf{c}$ by $c_{\mathbf{a}_{j}}=\mathbf{a}_{j}^{T}\mathbf{c}$. Then we apply the dynamic programming algorithm to solve (12) with plug-in $\delta_{ij}$’s. Although the multi-choice knapsack problem is a NP-hard problem, we can still solve it within pseudo-polynomial time (Kellerer et al., 2004). Specifically, we denote $\hat{Z}_{l}(b)$ as the optimal value $\frac{1}{n}\sum_{i=1}^{l}\sum_{j=1}^{|\mathcal{A}|}\hat{\delta}_{ij}d_{ij}$ for the first $l$ subjects with budget constraints $\frac{1}{n}\sum_{i=1}^{l}\sum_{j=1}^{|\mathcal{A}|}c_{\tilde{a}_{j}}d_{ij}\leq b$. Let $\hat{Z}_{l}(b)=-\infty$ if no solution exists and $\hat{Z}_{0}(b)=0$. We define the budget space as $\mathcal{B}=\{b:0\leq b\leq B\}$ including all possible average costs for $n$ subjects, where $0$ is the minimal cost if no treatment is applied to subjects, and the maximal cost is our specified budget $B$. Once the iterative algorithm ends, the optimal objective function is obtained as $\hat{Z}_{n}(B)$ and the optimal treatment assignment is the output $\{d_{ij}:i=1,...,n,j=1,...,|\mathcal{A}|\}$. The detailed algorithm is illustrated in Algorithm 2.