Efficient Heterogeneous Treatment Effect Estimation With Multiple Experiments and Multiple Outcomes

Constructing personalized policies is increasingly of interest across many diverse communities ranging from internet companies to personalized medicine (Peysakhovich and Lada, 2016; Kent et al., 2018). Unfortunately, estimating heterogeneous effects, sometimes referred to as conditional average treatment effects (CATEs), requires much more statistical power than simply estimating the average treatment effect (Gelman, 2018). In this paper, we introduce the LR-learner, a method of combining multiple experiments and multiple outcome metrics to improve the precision of CATE estimation across all of them. Even if an analyst is only interested in estimating CATEs for a single experiment, adding data from past experiments allows a more efficient estimation.

At its core, the goal of heterogeneous estimation is to output an estimate of the treatment effect for a unit whose features are specified by the vector $x$. Under the potential outcomes framework (Rubin, 1974), we denote the potential outcomes under treatment and control as $Y_{1}$ and $Y_{0}$, and their conditional means as $\mu_{1}(x)=E[Y_{1}|X=x]$ and $\mu_{0}(x)=E[Y_{0}|X=x]$, respectively. Our estimand of interest is the

Much of the literature on experimental analysis focuses on the case of a single experiment with a single treatment and a single outcome (e.g. a medical trial of a drug with survival as an outcome). Here there is a large body of work on efficiently estimating both the average effect of this experiment as well as finding subgroups which are more or less affected by the treatment (Imbens and Rubin, 2015; Wager and Athey, 2018; Künzel et al., 2019).

In many environments of interest, however, we are not considering a single experiment/outcome in isolation. For example, the A/B testing that is common in modern internet companies is different than the classical case in a few ways. First, large companies can run hundreds of tests per day (bez, 2014; Xie and Aurisset, 2016), measuring and trading off multiple outcome measures. For example, a movie streaming company can conduct A/B tests that run the gamut from tweaking movie recommendation algorithm parameters, to testing landing page layout changes. Each of these experiments can be evaluated on a set of metrics ranging from user retention to server CPU usage. Decision makers then trade off between metrics to make a deployment or non-deployment decision. The effect of each intervention is likely to be heterogeneous across treated units. Second, unlike in many classical applications, the set of unit level features that analysts have access to is very high dimensional. Thus a key problem is one of representation learning: finding a good, lower dimensional, feature description for efficiently learning the heterogeneity of treatment effects.

Our key insight is that pooling information across multiple experiments and multiple outcomes and learning a single feature representation for treated units allows us to greatly improve the precision with which we estimate heterogeneous effects over standard methodologies. The potential to pool experiments is increasingly viable in the online A/B testing setting. In addition, the idea of multiple experiments with high dimensional feature spaces is becoming applicable to other domains e.g. pooling many clinical trials or experimental work in systems biology.

For the one experiment/one outcome case, a simple methodology for learning CATEs is the T-learner (Künzel et al., 2019), where we train a model, $\hat{\mu}_{1}(x)$, on data from units in the treatment group and train a model, $\hat{\mu}_{0}(x)$, on data from units in the control group. The estimated effect can be computed as $\hat{\tau}(x)=\hat{\mu}_{1}(x)-\hat{\mu}_{0}(x)$.

Here, we build on the T-learner and consider the case of multiple experiments and multiple outcomes. Every unit may be subject to multiple experiments, and all outcome metrics are observed. Our goal is to estimate the treatment effect for all units across all experiments. This task can be modeled as a tensor factorization problem and can be done efficiently when the tensor is low rank. We refer to our method as the Low Rank Learner, or LR-learner.

For each unit, the LR-learner learns a $d$-dimensional unit embedding using a small neural network over input unit level features. The LR-learner also learns an embedding for each experiment treatment and control. Finally, for each outcome, a linear operator transforms the experiment and unit embeddings so that the dot product of these transformed embeddings predicts the desired outcome. The neural network, experiment embeddings and operators are learned jointly using the observed data.

The LR-learner takes advantage of the low dimensional bottleneck of the embeddings and linear operators, which is what allows us to pool information across experiments and outcome types. In practice, true CATEs are not observed (since we do not observe any unit in both treatment and control), so we first evaluate the LR-learner on fully synthetic and semi-synthetic data. We show that pooling information across experiments and outcomes greatly improves efficiency relative to independent learning.

Finally, we evaluate the LR-learner on a real dataset of $53$ product experiments run at a large internet company, focusing on $5$ real key performance indicators used by teams to make product decisions. We show that there is substantial information in heterogeneous effects shared across experiments and metrics, and that the LR-learner outperforms the independent level T-learner. In addition, we discuss other real world advantages of the LR-learner including scalability and the ability to use the LR-learner as a feature extractor negating the need for constant retraining when new experiments are run.

There exist many methods for estimating CATEs on data from a single experiment (Peysakhovich and Lada, 2016; Powers et al., 2018; Wager and Athey, 2018; Künzel et al., 2019; Kennedy, 2020). Our main contribution over this large literature is the consideration of combining data from multiple experiments and multiple metrics into a single model. Although there are many methods that aim to directly estimate the treatment effect (Athey et al., 2018b), in this paper we consider only the Two-Model approach (Gutierrez and Gérardy, 2017), because it naturally extends to the multiple experiment case. Possible extensions of our ideas to other “meta-learners,” like the X-learner (Künzel et al., 2019) are an interesting future direction though beyond the scope of this paper.

There is also a lot of recent literature on matrix completion methods on panel data (Agarwal et al., 2021a, b; Athey et al., 2018a; Bai and Ng, 2021; Fernández-Val et al., 2021). They consider the setting in which units are observed over multiple time periods. In each period, a unit can be exposed to some binary treatment, and their outcome for that period is recorded. These papers assume some sort of sparsity structure for when units receive treatment or when outcomes are observed. Our methodology does not assume any sparsity structure, so is not reliant on panel data, but can also be applied to this setting.

The model architecture employed by the LR-learner is inspired by the RESCAL model used for tensor factorization (Lerer et al., 2019; Nickel et al., 2011). Typically, these methods are applied to construct low dimensional embeddings of objects such as social or knowledge graphs (Lerer et al., 2019). So, a conceptual contribution of our paper is to make a link between two concepts. First, we frame the causal inference problem as a missing data problem (Ding and Li, 2018). Second, to solve the missing data problem, we use efficient low dimensional methods to fill in unobserved tensor values. We note that unlike in the case of a matrix, where the singular value decomposition provides the “canonical” way to factorize, there are many ways to factorize a tensor (Welling and Weber, 2001). Exploring other tensor factorization methods is an interesting avenue for future work.

We consider treatment assigned at the level of an individual (aka. unit), and consider the presence of multiple outcome metrics. Let $n$ be the number of units, $m$ be the number of observed pre-treatment features, $K$ be the number of experiments, and $J$ be the number of observed outcome metrics. We denote by $Y_{ijkt}$ the potential outcome of unit $i$ on metric $j$ when assigned to experiment $k$ with treatment arm $t$. We denote $t=0$ to be the ‘control’ arm for any experiment.

We make the following assumptions:

The individual treatment effect (ITE) for individual $i$ for metric $j$ in experiment $k$ for treatment $t$ is denoted

In this framework, a typical ‘classical’ randomized trial (e.g. a medical trial) would have a single arm, a single outcome, and a feature vector for each unit. In this case, the ITE would simply be denoted by a vector with one value (i.e treatment effect estimate) per unit $i$.

If we move to the case of a standard trial with multiple outcome measures the ITE can be represented as a matrix with rows corresponding to individuals and columns corresponding to outcome measures. Finally, moving to the fully general multi-experiment, multi-treatment, multi-outcome case the ITE becomes a tensor. This brings us to our final key assumption:

Assuming the ITE tensor has low rank structure means that there is information about a given entry of interest in other entries. This is the core of the proposed methodology. In practice, we do not observe any ITEs since no individual is ever in both a treatment and a control simultaneously. Instead, we will focus on a relative of the ITE which instead featurizes each unit, $i$, using a vector, $x_{i}$. The CATE for treatment arm, $t$, in experiment, $k$, is defined as

We now discuss how to train the LR-learner as a tensor factorization. Our factorization has several components: We have a sample of units assigned to experiments/treatments and their features denoted by $\mathcal{D}=\{y_{ijkt},x_{i}\}$. We train a featurized tensor factorization model which includes

1. 1.

   A $d$ dimensional representation of each unit $i$, which is a function of $x_{i}$, called $v(x_{i})\in R^{d}$. This is expressed as a simple neural network which inputs $x_{i}$ and outputs a $d$ dimensional feature vector $v_{i}$.

2. 2.

   A $d$ dimensional representation of each experiment treatment arm $t$ and control $e^{t}_{k}\in R^{d}$ and $e^{0}_{k}\in R^{d}$.

3. 3.

   A $d\times d$ transformation matrix for each outcome $A_{j}$.

4. 4.

   The parameters are trained with an objective similar to RESCAL (Nickel et al., 2011)

   $$\min_{v,A,e}\sum_{(y_{ijkt},x_{i})\in\mathcal{D}}(y-v(x_{i})^{\prime}A_{j}e^{t}_{k})^{2}.$$

Given the trained model, the predicted potential outcome for a unit $i$ in experiment $j$ in arm $t$ for outcome $k$ is given by:

This means the estimate of the CATE is just a dot product between $v(x_{i})$ (transformed for the particular outcome) and the difference between the treatment and control embeddings. Formally:

The fact that the CATE is fully linear is useful in real world systems where new experiments are constantly run. We will explore this idea in detail in section 6.

So far, we have considered the case of a fixed dataset $\mathcal{D}$, and an analyst interested in heterogeneous effects in all or a subset of the experiments/outcomes. We can also consider a scenario where the analyst already has access to an LR-learner trained on a large set of existing experiments. When a new experiment is run, the analyst may want to estimate HTEs for the new experiment. Call this new data $\mathcal{D}_{\text{new}}$. For simplicity, consider the case where there is just a single outcome metric of interest, and potentially this outcome measure is not one that is considered in past experiments.

One potential solution is to add the new experiment to the data, call this combined dataset $\mathcal{D}_{c}$ and re-train the LR-learner on the complete dataset. However, when $\mathcal{D}$ is extremely large relative to $\mathcal{D}_{\text{new}}$, then the addition of the new data will not change the shared parameters across all experiments much. Thus, we can simply fix the parameters of the neural network, $v_{\mathcal{D}}(\cdot)$, to the ones of the model trained on the original data. In other words, we can find new experiment treatment/control vectors $e_{\text{new}}$ for the outcome in question by solving:

In other words, we simply construct an independent T-learner on the new experiment. However, instead of using the original feature space $x$ and possibly a complex model, we use the pre-trained LR-learner as a feature extractor and a linear regression as the underlying model. This reduces the complexity of the learning problem substantially.

This is very similar to the idea of fine-tuning supervised (or unsupervised) models in other areas of machine learning such as vision or NLP where an embedding layer of a model trained on a large dataset is used as the feature space for a simpler model in a new dataset (Mikolov et al., 2013; Mahajan et al., 2018).

In practice, if an LR-learner were to be deployed in an organization, we would suggest the following process: Retrain the full model on a pool of experiment data intermittently. Save the model, and in particular, save the feature extractor $v(\cdot)$. In daily analyst practice, when estimating heterogeneous treatment effects for a new experiment, reuse the saved $v(\cdot)$ to train linear independent T-learners.

In this paper, we devise the LR-learner, an approach for simultaneous CATE estimation across multiple experiments and multiple outcomes. Our tensor factorization approach assumes a low rank structure in the ITE tensor across experiments and outcome metrics. We show that there is indeed a low-rank structure for CATEs across experiments and outcomes, which allows the LR-learner to be $200\times$ more sample efficent compared to individual CATE models. We show using real data from Facebook that the LR-learner outperforms the Independent T-learner approach in both $\tau\textrm{-risk}$ and $\mu\textrm{-risk}$ evaluation metrics.

Our approach can be especially useful in domains such as internet companies where experiments are continuously being run, potentially across various unrelated applications, and outcome metrics are inexpensive to obtain. This is in contrast to many fields like medicine in which panel data is more common, and in which experiments may be even more related, but multiple outcome metrics can be difficult to obtain. One avenue for future work is to evaluate the LR-learner on panel data and compare it to other causal matrix completion methods (Agarwal et al., 2021a; Athey et al., 2018a).

We have focused purely on the case of A/B tests, but another fruitful direction is to combine the LR-learner architecture with datasets that are either wholly or partially observational (Shalit et al., 2017; Lada et al., 2019). The architecture naturally extends to estimating treatment assigment probability as well as outcomes, so a possible future direction is incorporating either X-learner (Künzel et al., 2019) or doubly-robust (Kennedy, 2020) methods to the LR-learner.

We establish the contribution of the LR-Learner as a method for estimating CATEs, and we hope that future work explores potential applications for the learned embeddings. The user, metric, and experiment embeddings have many potential downstream uses, such as finding or comparing similar users or experiments. The learned embeddings may also inform the design of future experiments.