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.
18,612 characters · 13 sections · 6 citation commands
Shrinkage-Based Regressions with Many Related Treatments
\onehalfspacing
Observational causal models are increasingly popular in industry, often under the header of “double machine learning“ or “causal meta-learners” chern-dml, causal-ml-book. In practice, these methods are often used to study a single treatment or a small set of related treatments. However, many practitioners face a broad set of related but potentially sparse treatments and want to disentangle their individual impacts. Examples include estimating treatment effects of different marketing touchpoints, ordering different types of products, or signing up for different services. We refer to these as sub-treatments. This creates a dilemma: focusing on a single treatment typically yields precise estimates without the necessary granularity, while seperately estimating all sub-treatments typically results in estimates that are too imprecise to guide practical decision-making.
We propose a middle-ground between these extremes. Our core idea is to construct a “focal function” representing the single-treatment estimator, then use a specialized ridge regression to shrink additive interaction terms towards this component. This selective shrinkage improves the mean squared error (MSE) of causal estimates for each sub-treatment, leaves estimates essentially intact when a sub-treatment has sufficient data, and allows us to reconstruct the aggregate treatment effects for any degree of penalization without re-estimating the model. The use of closed-form ridge regression makes our method fast and scalable. The mechanics of this ridge regression also apply to estimating conditional average treatment effects, and to settings outside of causal meta-learners.
This work builds heavily on the theory of causal meta-learners panel-dml, chern-dml, causal-ml-book, dr-learner, x-learner, r-learner, and a smaller set of papers discussing their application in real-world settings cate-selection. It is closely related to shrinkage using hierarchical Bayesian models bda-3. We emphasize that our methods are primarily built for low-MSE estimation rather than rank-ordering treatments, for which partially linear models are sometimes inappropriate lal-rankings. However, to our knowledge the results below will also apply to propensity score reweighted models.
Section (ref) illustrates our model and its properties. Section (ref) demonstrates the models mechanics using simulated data, and shows how Wayfair has used this idea to unlock practical decision-making at a more granular level. Section (ref) concludes.
Our starting point is a partially linear single treatment model. For each unit $i$, we denote a scalar outcome $Y_i$ with a binary treatment $D_i$, a vector of confounders $X_i$ whose components are known but whose functional form $f$ is unknown, and an unobserved orthogonal error term $\epsilon_i$. This yields the structural equation:
Often, practitioners want to consider $K$ distinct treatments within a given business domain. As a running example, we consider the impact a potential customer being served a marketing touchpoint ($D_i$) when there are $K$ different types of marketing touchpoints (TV, search, social media, etc). We call these $K$ sub-treatments, and denote them $(D_{1i},...,D_{Ki})$. This motivates a regression with $K$ distinct treatment effects:
We propose a more general formulation that gives a middle ground between the single treatment case (which is coarse but precise) and the sub-treatment case (which is flexible but imprecise). Denoting $D'_{ki} \equiv max_k(D_{ki})$, we write:
We use the shorthand $D'_{ki}$ to save on notation, and to demonstrate that many of our theoretical results hold for a broader class of aggregation functions. We call this the focal function.
We operationalize this with a flexible functional form $f$ using causal meta-learners. Following the steps in r-learner, we assume that $E[\epsilon_i|X_i]=0 \forall X_i$ and apply a Robinson transformation:
Following convention, we denote residualized quantities with tildes:
Using formulation (ref) and converting to the estimating equation (ref) gives us relatively simple way to compromise between the precision of a single-treatment model and the generality of a sub-treatment model. In particular, we shrink the coefficients $\beta_k$ towards zero, but do not penalize $\beta_0$. This reduces the MSE of each sub-treatment estimate, but (we show below) lets us reconstruct the single-treatment estimator.
A simple and computationally efficient way to operationalize this is via ridge regression:
After tuning $\lambda$ using a hold-out sample, we can generate point estimates using a simple closed-form solution. Letting $X \equiv [\tilde{D}', \tilde{D}_{1}, \dots, \tilde{D}_{K}]$:
\footnotetext{This is the homoscedastic form of the sandwich covariance estimator. A heteroscedasticity-robust version replaces $\hat{\sigma}^2 X'X$ with $X' \hat{\Sigma} X$, where $\hat{\Sigma}$ is a diagonal matrix whose entries are the squared residuals $\hat{\varepsilon}_i^2$.}
In addition to giving us low-MSE estimates of heterogeneous treatment effects, our formulation lets us reconstruct the single-treatment estimator for any level of penalization. We show this in a few steps.
When ridge regularization is applied only to the sub-treatment parameters $\beta_k$ for $k > 0$, these coefficients shrink toward zero as $\lambda \to \infty$. Consequently, the ridge solution converges to a univariate projection of $Y$ onto $\tilde{D}'$, and the unpenalized parameter $\beta_0$ approaches:
Even when shrinkage is only partial, we can reconstruct an estimate of $\tau_0$ from our estimands. To see how, we expand the numerator of (ref):
Assuming the noise term is mean independent of $\tilde{D}'$, i.e., $\mathbb{E}[\tilde{D}' \epsilon] = 0$, we obtain the following closed-form expression:
To build intuition around this quantity, consider the case when $f(X) = 0$ (no selection bias) and $D' \equiv max_k(D_k)$. Then (ref) simplifies to the sum of conditional probabilities:
In ridge regression, the estimated coefficients $\beta_k^\lambda$ generally vary with the regularization parameter $\lambda$, affecting the contributions of the penalized variables $\tilde{D}_k$ to the fitted outcome. However, the quantity $\tau_0$ in Equation (ref), which corresponds to the projection of $Y$ onto the unpenalized regressor $\tilde{D}'$, remains invariant to $\lambda$.
To see why, recall that the ridge estimator solves the optimization problem in Equation (ref). Taking the derivative of the objective with respect to $\beta_0$ yields: \[ \frac{\partial}{\partial \beta_0} \mathbb{E}\left[ \left( Y - \beta_0 \tilde{D}' - \sum_{k=1}^K \beta_k \tilde{D}_k \right)^2 + \lambda \sum_{k=1}^K \beta_k^2 \right] = -2 \, \mathbb{E}\left[ (Y - \beta_0 \tilde{D}' - \sum_{k=1}^K \beta_k \tilde{D}_k) \tilde{D}' \right]. \] Setting this derivative to zero gives the following first-order condition. And we know ridge estimated coefficients $\beta_k^\lambda$ satisfies it: \[ \mathbb{E}\left[ (Y - \beta_0^\lambda \tilde{D}' - \sum_{k=1}^K \beta_k^\lambda \tilde{D}_k) \tilde{D}' \right] = 0. \] Letting $Y^\lambda = \beta_0^\lambda \tilde{D}' + \sum_{k=1}^K \beta_k^\lambda \tilde{D}_k$ denote the ridge-predicted outcome, this condition becomes: \[ \mathbb{E}\left[ (Y - Y^\lambda) \tilde{D}' \right] = 0, \quad \Rightarrow \quad \mathbb{E}\left[ Y^\lambda \tilde{D}' \right] = \mathbb{E}\left[ Y \tilde{D}' \right]. \] It follows that the projection coefficient of $Y$ onto $\tilde{D}'$ is: \[ \tau_0^\lambda = \frac{\mathbb{E}[Y^\lambda \tilde{D}']}{\mathbb{E}[\tilde{D}'^2]} = \frac{\mathbb{E}[Y \tilde{D}']}{\mathbb{E}[\tilde{D}'^2]} = \tau_0. \]
Thus, even though the individual coefficients $\beta_k^\lambda$ vary with $\lambda$, the combined projection $\tau_0^\lambda$ remains fixed. This implies that we can accurately reconstruct the single-treatment estimator regardless of the regularization strength applied to the sub-treatment terms.
To demonstrate the behavior of our method in a controlled setting, we simulate a data-generating process with six binary sub-treatments. These sub-treatments are unevenly distributed, with activation probabilities ranging from 5% to 20%, reflecting common real-world imbalances. The outcome depends on both a shared treatment effect, activated whenever any sub-treatment occurs, and additive effects from each sub-treatment.
Specifically, for each unit $i$, we draw $D_{1i}, \dots, D_{6i} \sim \text{Bernoulli}(p_k)$ independently, where:
We define the focal treatment $D'_i = \max_k(D_{ki})$ and generate outcomes according to: \[ Y_i = \beta_0 \cdot D'_i + \sum_{k=1}^6 \beta_k D_{ki}, \] with $\beta_0 = 5$ and sub-treatment coefficients $\boldsymbol{\beta_k} = [2, 2, 1, 1, -1, -1]$.\footnote{This simulation excludes confounders to keep the focus on the mechanics of the specialized ridge regression.}
Similar to the derivation of the single-treatment effect in Equation (ref), we can derive the expected impact of a specific sub-treatment as:
Figure (ref) illustrates how our estimator interpolates between high-variance heterogeneous treatment effects $\tau_k$ ($\lambda \approx 0$) and the shared average $\text{ATT}$ ($\lambda \rightarrow \infty$). Sparse sub-treatments shrink more aggressively, enabling stable and interpretable causal estimates even when sample sizes vary across treatments.
Figure (ref) confirms that the proposed estimator effectively trades off bias and variance across a range of regularization strengths.
We demonstrate our approach using real-world data from Wayfair, a large e-retailer that sells home goods ranging from small items (e.g., lamps) to large items (e.g., luxury beds). An important business decision in this context is whether certain product classes have an outsized long-term effect on customer spending. The shape of these effects is not obvious ex ante: it is possible that delivering small items reliably builds customer trust, or that delivering large, bulky items creates a more memorable positive experience.
We report anonymized heterogeneous sub-treatment effects from 53 product classes using our custom ridge regression method. The estimates are shown at three levels of regularization: no penalty ($\lambda = 0$), an optimally tuned penalty ($\lambda = 0.001$), and a large penalty ($\lambda = 100$), which approximates the projection of outcomes onto the focal treatment $\tilde{D}'$. As expected, increasing the penalty shrinks the class-specific estimates toward a shared baseline, similar to the behavior we observed in simulation.
Figure (ref) highlights a key feature of our method: the ability to communicate treatment effects in a way that is robust to data sparsity. Because product classes are sorted by prevalence, we can visually confirm that higher-volume categories (e.g., beds, rugs) show stable estimates across penalties, while sparser classes (e.g., pet, flooring) are more sensitive to regularization. This effect is not an artifact of the method but a direct reflection of sample size–driven variance, which our approach dampens systematically.
We also tune the shrinkage parameter using a cross validation procedure and find that our method provides much more stable and interpretable effect estimates for business stakeholders. While the individual sub-treatment estimates are not unbiased, they enable more confident decision-making in settings where data volume varies dramatically across categories. We can further reconstruct the unbiased single-treatment estimates we would have obtained from extreme penalization (green) without re-estimating the model.
In practical business applications, we often use causal meta-learners to estimate separate effects of many partially-overlapping sub-treatments. Instead of choosing between coarse models with few treatments or very noisy models with every treatment, we propose a model that applies selective shrinkage towards a so-called focal function (the max across sub-treatments). In addition to providing lower-MSE estimates for each sub-treatment, we can easily reconstruct the single-treatment estimator in an internally-consistent and computationally efficient way. We have found this to be a flexible and pragmatic approach, which helps stakeholders make key business decisions with an appropriate degree of confidence.