Principal Stratification and Mediation

1 Introduction

In a typical A/B test, we analyze the relationship between a treatment and an outcome. In some instances, a third element—a mediator—in some sense sits between the treatment and outcome. We can place a collection of powerful methods within this framework, differing in what assumptions the methods make about principal strata and (in)direct effects.

The methods in this note elucidate scenarios like:

• Contamination, where the treatment is not administered perfectly,
• Dose-response, where the treatment effect depends on the mediator effect.
• Product funnel transitions, where the treatment pushes more people across a funnel stage, but may dilute subsequent funnel performance,
• Growth tactics, where there is uncertainty about whether the people acquired by the tactics behave the same way as the organic acquisitions, and
• Surrogates, where we try to learn about a long-term outcome based on short-term impact.

A variety of techniques can be used to explore these scenarios, including:

• Multiple Imputation: predicting counterfactuals allows us to estimate individual treatment effects. Multiple Imputation incorporates the uncertainty of these predictions.
• Regression Models: by fitting models for the mediator and the outcome, we can decompose the total effect into direct and indirect effects.
• Matching Methods: offer an alternative identification strategy to regression methods.

Examples are taken from both Product and Marketing Data Science.

2 Principal Strata

We assign units (henceforth we assume the units are people) either to a treatment or control group, at random. Let \( Z_i = 1 \) if we assign person \( i \) to treatment and 0 otherwise. We consider only binary treatments here.

The treatment an individual receives may influence their mediator value, but we assume it does not influence the mediator values of any other individual. (Rubin 1980) calls this the Stable Unit Treatment Value Assumption (SUTVA). Then we can represent the potential mediators for person \( i \) as \( \{ M_i(0), M_i(1) \},\) with \( M_i^\textrm{obs} = M_i(Z_i). \) The observed mediator value is random because it is a function of the random treatment assignment, but the potential mediators are covariates: fixed characteristics of the individual rather than random variables.

When the mediator is binary, \( M_i(z) \in \{ 0, 1 \} \) for \( z=0, 1. \) The potential mediators categorize people into four principal strata (Frangakis and Rubin 2002) as shown in Table  1. The principal stratification is Mutually Exclusive and Completely Exhaustive (MECE): everyone falls into one of these categories. The last column of the table shows the proportion of people in each category, with \( \pi_\textrm{nt} + \pi_\textrm{co} + \pi_\textrm{de} + \pi_\textrm{at} = 1. \)

To motivate the names of these strata, consider encouraging a patient to take a certain medicine. There are people who will take the medicine only if encouraged to do so, the Compliers. There are people who do not need encouragement and take the medicine with or without it, the Always Takers. There are those who will not take the medicine despite encouragement, the Never Takers. Finally, there are those who hate following directions and take the medicine only if not encouraged to do so, the Defiers.

Because we define the principal strata in terms of \( \{ M_i(0), M_i(1) \}, \) the principal stratum to which an individual belongs is a fixed characteristic of the individual, a covariate. Since we only observe one of the potential mediators, in general we do not know to which principal stratum an individual belongs. Many of the techniques we discuss here involve assumptions about which principal strata are possible based on knowledge of the treatment and mediator.

For example, in many scenarios it is reasonable to assume there are few or no Defiers: \( \pi_\textrm{de} \approx 0. \) In this case, random assignment allows us to estimate the other proportions. Since all observed instances of \( M_i^\textrm{obs} = 1 \) in the control group correspond to Always Takers, we observe the proportion of Always Takers in the group. Random assignment balances all covariates, including principal strata, so the proportion of Always Takers in the treatment group will be similar. All observed instances of \( M_i^\textrm{obs} = 0 \) in the treatment group correspond to Never Takers, so we observe the proportion of Never Takers in the group and can estimate the proportion in control. The proportion of Compliers is then \( \pi_\textrm{co} = 1 - \pi_\textrm{at} - \pi_\textrm{nt}. \)

Example 1: To illustrate how principal stratification applies in practice, consider running an ad campaign on Instragram. To measure the impact of the campaign, we may run a lift study. The ad platform holds back a random 10% of the audience from ad exposure. The remaining 90% form the Marketing group.

Despite the name, the ad platform cannot guarantee everyone in the Marketing group will see ads associated with our campaign. The platform has to select which ads to show when the person loads the page, but the ad corresponding to our campaign may be “below the fold,” and the person may not scroll down far enough to see it. In this example, \( Z_i \) indicates whether we assigned person \( i \) to the Marketing group, and \( M_i \) indicates whether the person scrolled down far enough to see the ad.

Table  2 summarizes the ad exposure in each group. In the Holdout, no one was exposed, and in the Marketing group, about 56% were. It is reasonable to assume there are no Defiers and no Always Takers in this example. Every exposed person in the Marketing group is a Complier and every unexposed person in the Marketing group is a Never Taker. There are exactly 5,000,000 Compliers and 4,000,000 Never Takers in the Marketing group, or around 56% and 44%. Although we don’t know which individuals in the Holdout group are Compliers, because of random assignment, we know around 56% of the Holdout group are Compliers, and around 44% are Never Takers.

3 Direct and Indirect Effects

Both treatment and mediator may influence the outcome. Assuming SUTVA applies to the outcome, let \( \{ Y_i(z, m) \} \) denote the set of potential outcomes for person \( i. \) This set will have one element for each treatment assignment and each mediator value. The treatment is binary; if the mediator is too, there will be four potential outcomes for each person. In general, \( Y_i(z, m) \) is a function with support on \( [0, 1] \times \mathcal{M}, \) where \( \mathcal{M} \) is the set of mediator values.

The “total” effect of treatment on the outcome for an individual is a comparison of these potential outcomes: \( \tau_i = Y_i(1, M_i(1)) - Y_i(0, M_i(0)). \) We can decompose this effect by adding and subtracting another potential outcome (VanderWeele 2015, § 2.16):

\begin{align*} \tau_i &= \{ Y_i(1, M_i(1)) - Y_i(1, M_i(0)) \} \\ &\phantom{=} \hspace{20pt} + \{ Y_i(1, M_i(0)) - Y_i(0, M_i(0)) \} \\ &=: \textrm{Indirect Effect}_i + \textrm{Direct Effect}_i. \end{align*}

The terminology is standard but I prefer to think of the indirect effect as being the part of the total effect explained by the mediator. The indirect effect is the change in outcome produced by changing the mediator from its control value to its treatment value, while holding the treatment itself fixed. The direct effect is the part of the total effect not explained by the mediator, the change in outcome produced by applying the treatment while holding the mediator fixed at its control value.

We define both direct and indirect effects in terms of a potential outcome impossible to realize, \( Y_i(1, M_i(0)). \) This potential outcome refers to the scenario where we assign the person to treatment, but somehow we keep their mediator at its control value. Without specifying a strategy for doing so, this potential outcome is at best ambiguous and at worst meaningless. The only way to keep the mediator at its control outcome is to assign the individual to the control group, in which case we observe \( Y_i(0, M_i(0)), \) not \( Y_i(1, M_i(0)). \) We never observe the latter quantity for anyone so (Frangakis and Rubin 2002) calls it a priori counterfactual.

For an individual with \( M_i(0) = M_i(1), \) the indirect effect is 0:

\begin{align*} \textrm{Indirect Effect}_i &= Y_i(1, M_i(1)) - Y_i(1, M_i(0)) \\ &= Y_i(1, M_i(0)) - Y_i(1, M_i(0)) \\ &= 0. \end{align*}

With binary mediators, any treatment effect among Never Takers or Always Takers must be a direct effect. The effect of treatment on the outcome among Compliers and Defiers may be a blend of direct and indirect effects.

We do not observe individual total effects (since we only observe one of the potential outcomes), but we may estimate average total effects across a population by comparing outcomes in treatment and control. Average direct and indirect effects are not identified even in an A/B test because we never observe the a priori counterfactual outcome, \( Y_i(1, M_i(0)). \) Still, there may be situations where we are willing to make assumptions about these effects. For example, the next sections assume the direct effect is 0, at least for certain groups. Since the direct effect is not identified, this is a tough assumption to validate, but alternative strategies try to estimate direct and indirect effects through regression models or matching strategies.

Example 1, cont’d: Table  3 shows the full set of potential mediators and outcomes for several individuals in the lift study introduced in § 2. The last columns show the total, direct, and indirect effects. In reality, we never observe the counterfactual mediators or outcomes and thus cannot calculate individual effects, but it is helpful here to clarify how effects are defined.

Alice is a Complier since \( M(0) = 0 \) and \( M(1) = 1. \) The total effect is

\begin{align*} Y_i(1, M_i(1)) - Y_i(0, M_i(0)) &= Y_i(1, 1) - Y_i(0, 0) \\ &= 1 - 0 = 1, \end{align*}

the direct effect is

\begin{align*} Y_i(1, M_i(0)) - Y_i(0, M_i(0)) &= Y_i(1, 0) - Y_i(0, 0) \\ &= 0 - 0 = 0, \end{align*}

and the indirect effect is

\begin{align*} Y_i(1, M_i(1)) - Y_i(1, M_i(0)) &= Y_i(1, 1) - Y_i(1, 0) \\ &= 1 - 0 = 1. \end{align*}

Suppose \( Y \) indicates a purchase. Alice’s purchase behavior is driven by the ad exposure: she only purchases if exposed.

Brian is also a Complier with a total effect of 1, but the direct effect is 1 and the indirect effect is 0. Brian’s purchase behavior is affected by the treatment assignment, but this is not explained by ad exposure. We wouldn’t expect to see this in real life: Brian wouldn’t know which experiment group we had assigned him, except via the ads shown, so a non-zero direct effect is implausible. Perhaps the treatment assignment had some side effect other than ad exposure.

Charlie is a Never Taker, and Charlie’s pattern of potential outcomes is identical to Alice’s. Nevertheless, Charlie’s total effect is zero:

\begin{align*} Y_i(1, M_i(1)) - Y_i(0, M_i(0)) &= Y_i(1, 0) - Y_i(0, 0) \\ &= 0 - 0 = 0. \end{align*}

The direct and indirect effects are also 0.

Treatment assignment does not affect Charlie’s ad exposure and thus does not affect his purchase behavior. The indirect effect here is not just about the effect of ad exposure. Ad exposure does in fact cause Charlie to purchase, yet the indirect effect is 0. The indirect effect is the component of the total effect explained by ad exposure. Since Charlie is never exposed to ads, the total effect is 0, and so is the indirect effect. Two people in different principal strata can have identical sets of potential outcomes, yet have different total, direct, and indirect effects.

4 One-Sided Non-Compliance

In §§ 2– 3, we discussed our inability to identify the principal strata to which individuals belong, and to calculate direct and indirect effects, in the absence of a full set of potential mediators and outcomes. Despite this, it is sometimes possible to estimate effects in particular strata.

Example 1 explored a lift study with Marketing and Holdout groups. Despite the name, not everyone in the Marketing group was exposed to ads; however, no one in the Holdout was exposed to ads. This is an example of one-sided non-compliance. One-sided non-compliance occurs when only two principal strata occur: Compliers and either Never Takers or Always Takers. We assume Defiers do not exist with one-sided non-compliance.

In this example, it is reasonable to assume there are no Defiers, and also no Always Takers. This implies there are no individuals with \( M_i^\textrm{obs} = 1 \) in the control group, which is easy to check in Table  2. It also implies no one in the treatment group has counterfactual mediator value 1. Since we do not observe counterfactuals, the data cannot validate this assumption and we must instead rely on domain expertise. In Example 1, the implementation of the lift study supports this assumption, but in other cases it may not be obvious.

In this case we know \( M_i(0) = 0 \) for everyone. This tells us the counterfactual mediator for people in the Marketing group, and therefore tells us the principal stratum for everyone in the group. People in the Marketing group with \( M_i^\textrm{obs} = 1 \) are Compliers and people with \( M_i^\textrm{obs} = 0 \) are Never Takers. We don’t know the counterfactual mediators for people in the Holdout, so we don’t know the principal strata.

In § 3, we saw the indirect effect among Never Takers is 0. This is not an assumption; it is a consequence of the definitions of the principal strata and indirect effects. Any treatment effect among Never Takers is a direct effect.

A non-zero direct effect means switching someone from Holdout to Marketing, but without exposing the person to the ad, would somehow still affect their behavior. Since some invisible backend service controls the treatment assignment, the only way the assignment could influence behavior is by changing the ads seen. For the Never Takers, this mechanism does not apply, and a non-zero direct effect does not seem plausible. Since the direct and indirect effects are both zero for Never Takers, the total effect is also zero for this group, and any treatment effect can occur only among Compliers.

Comparing behaviors between the Marketing and Holdout groups tells us the average (total) effect of treatment on the outcome, \( \tau. \) This is an average of the total effects among Never Takers, \( \tau_\textrm{nt}, \) and Compliers, \( \tau_\textrm{co}: \)

\begin{align*} \tau &= \pi_\textrm{nt} \cdot \tau_\textrm{nt} + \pi_\textrm{co} \cdot \tau_\textrm{co} \\ &= \pi_\textrm{nt} \cdot 0 + \pi_\textrm{co} \cdot \tau_\textrm{co} \\ &= \pi_\textrm{co} \cdot \tau_\textrm{co} \end{align*}

Then the average total effect among the Compliers is:

\begin{equation}\label{eqn:wald} \tau_\textrm{co} = \frac{\tau}{\pi_\textrm{co}}, \end{equation}

where the numerator is the difference in outcome, Marketing minus Holdout, and the denominator is the fraction of people in the Marketing group who saw ads.

In this example, it is just as reasonable to assume the direct effect is 0 among the Compliers as among the Never Takers. Then Equation (\ref{eqn:wald}) gives not only the average total effect among Compliers, but also the average indirect effect: the average effect of ad exposure on the outcome, among everyone exposed.

When the proportion of Compliers is less than one, the average total effect, \( \tau, \) will be diluted compared to the effect just among the Compliers, \( \tau_\textrm{co}. \) Scaling up the average effect by the amount of dilution reverses this phenomenon.

Example 1, cont’d: Table  4 shows the purchases occurring in each experiment group of the lift study. The exposure pattern is the same as in Table  2. The table shows a conversion rate of 10% in the Holdout group and 12% in the Marketing group. The effect of treatment assignment is a 2 percentage point absolute lift in conversion rate, so \( \tau = 2\%. \) There are no Defiers or Always Takers, so this is a weighted average of the effects among Never Takers and Compliers. Treatment assignment does not influence ad exposure for Never Takers, and so we expect the effect on purchase behavior to be zero in this group. Since 56% of the campaign population are Compliers, the effect of assignment is a \( \tau_\textrm{co} = \tau / \pi_\textrm{co} = 2\% / 56\% = 3.6\% \) lift in conversion rate in this group. Since the direct effect is zero, this lift is the effect of ad exposure, not just treatment assignment.

In this example, accounting for ad exposure reveals the effect of ads to be \( \tau_\textrm{co} = 3.6\%. \) Had we not accounted for this, the effect \( \tau = 2\% \) would have been diluted in comparison. This dilution occurs because the Marketing group includes people with zero treatment effect, the people not exposed to ads.

When considering net rather than average effects, no dilution occurs. For example, when the outcome is a purchase, \( \tau \) and \( \tau_\textrm{co} \) represent lifts in conversion rates, among everyone and among Compliers, respectively. The net effect is the number of incremental purchases. Among everyone in the study, this is \( N \cdot \tau, \) where \( N \) is the total number of people assigned to treatment. Among the Compliers, this is \( N_\textrm{co} \cdot \tau_\textrm{co}, \) where \( N_\textrm{co} = N \cdot \pi_\textrm{co} \) is the total number of Compliers assigned to treatment. Whether we report the net effect among everyone, or among the Compliers, the result is the same:

\begin{align*} N_\textrm{co} \cdot \tau_\textrm{co} &= N \cdot \pi_\textrm{co} \cdot \frac{\tau}{\pi_\textrm{co}} \\ &= N \cdot \tau. \end{align*}

If the goal is to report net effects, we may ignore dilution.

In Example 1, there are 5,000,000 Compliers, so a 3.6 percentage point lift in conversion rate is 180,000 incremental purchases. Or we can calculate incremental purchases as a 2 percentage point lift among 9,000,000 people in the Marketing group. Either way, we get the same answer.

In many examples, the effect of the mediator is of more interest than the effect of the treatment. Although we did not assign the mediator at random, and there may be important differences between those who see ads and those who don’t, we can still estimate the effect of interest.

(Imbens and Rubin 2015, §§ 23–24) discuss inference for one-sided and two-sided noncompliance. The literature on Instrumental Variables calls our assumptions of zero direct effects exclusion restrictions.

5 Two-Sided Non-Compliance

Example 2: Our next case study measures the impact of a new product feature on an outcome like engagement or retention. Many tech companies roll out a new feature first to a small test group to measure its impact before releasing it to everyone. Other companies may question the value of such insights, or may be concerned about the drawbacks.

Groundbreaking features generate buzz, and a customer who learns about a new feature from friends in the roll-out group may be frustrated not to be able to use the feature themselves. This frustration is not only a bad customer experience; it also invalidates the Stable Unit Treatment Value Assumption discussed in § 2.

Instead we encourage some people, at random, to use the feature. No one is prevented from using the feature, but not everyone is encouraged, either. Depending on the nudge, this can provide an effective way of measuring the impact of the feature, without hurting anyone’s feelings.

Many mobile apps display a splash screen after an update, advertising new features. By raising awareness, the screen may entice people to try the new features, who otherwise would not have known about them. Consider showing the splash screen to some people and not others, at random.

If the splash screen succeeds, feature usage in the treatment group will be higher than in the control group. If the feature benefits engagement or retention, then we would expect these metrics to be higher in the treatment group than the control group, because feature usage is higher. Since the splash screen itself seems unlikely to affect these outcomes, we attribute any meaningful difference in outcome to differences in feature usage.

Here, \( Z_i \) indicates whether the person was shown the splash screen, and \( M_i \) corresponds to feature usage, which here we treat as binary. (In § 6, we discuss alternatives to binary indicators.) As in § 4, we have Never Takers who are not interested in the feature, even after seeing the splash screen. The splash screen entices some people to use the feature, who otherwise would not use it, and these people are the Compliers. We also have Always Takers, who are so enthusiastic about the feature they need no encouragement to use it. Finally, we may have Defiers who hate those splash screens and won’t use the feature if shown one, just to spite the app creators. While there may be a small number of people who so despise splash screens, we will assume they are neglible.

Table  5 shows the results. There are 10 million people split evenly across both experiment groups. As expected, feature usage and retention are both higher in the group shown the splash screen.

The 2 million people in the control group are all either Always Takers or Defiers. Since we assume Defiers are negligible, 40% of the control group are Always Takers. Since random assignment distributes Always Takers evenly across both groups, around 40% of the treatment group are also Always Takers, but we don’t know which individuals are. The other 100,000 people—2% of the treatment group—who used the feature are Compliers. Everyone else is a Never Taker. Thus \( \pi_\textrm{nt} = 58\%, \) \( \pi_\textrm{at} = 40\%, \) and \( \pi_\textrm{co} = 2\%. \)

The average effect of the splash screen was to increase retention by 0.04 percentage points. This is an average over Never Takers, Always Takers, and Compliers. As in § 3, it seems reasonable to assume this effect is concentrated among the Compliers. Equation (\ref{eqn:wald}) tells us the average effect of the splash screen among the Compliers was to increase retention by \( \tau_\textrm{co} = \tau / \pi_\textrm{co} = 0.04\% / 2\% = 2 \) percentage points. Assuming the direct effect is negligible, this is the average effect of feature usage on retention among Compliers.

This strategy does not offer insight regarding the effect of the feature on the Always Takers, the group most enthusiastic to use it. The effect of the splash screen on retention among Always Takers is zero, because it did not influence feature usage for this group. We never observe what happens for this group when they don’t use the feature, so there is no control to compare against. To measure the effect among Always Takers, we would need some strategy for preventing them from using the feature, but this is precisely the bad user experience we had hoped to avoid.

Only 2.4% of the people who used the feature are Compliers, so it may seem silly to focus on them. It doesn’t seem reasonable in this example to assume the effect of the feature for the Compliers is the same as for the Always Takers. The Always Takers are fundamentally more interested in the feature than the Compliers. It may be reasonable to assume the effect of the feature is higher among the Always Takers, but this is in no way supported by the data.

On the other hand, if the feature is impactful among Compliers, as it was in this example, we may consider investing in growth tactics to increase feature adoption. The target audience for growth tactics are people who will use the feature if nudged but not otherwise. This audience resembles Compliers more than Always Takers, so the average effect among Compliers is the quantity more relevant to growth marketers.

We may wish to investigate how Compliers and Always Takers differ in terms of observed covariates. We can do this through uplift modeling. First build a model that predicts the probability of using the feature among the treatment group, \( \pi_\textrm{co/at}(x) := \mathrm{Pr}\left\{ M = 1 \, \middle| \, Z=1, X=x \right\}. \) This is the probability a person with covariates \( X=x \) is either a Complier or an Always Taker. Then build a similar model on the control group: \( \pi_\textrm{at}(x) := \mathrm{Pr}\left\{ M = 1 \, \middle| \, Z=0, X=x \right\}. \) This is the probability a person with covariates \( X=x \) is an Always Taker. The uplift, \( \pi_\textrm{co}(x) := \pi_\textrm{co/at}(x) - \pi_\textrm{at}(x), \) is the probability a person with covariates \( X=x \) is a Complier. Finally, \[ \frac{\pi_\textrm{co}(x)}{\pi_\textrm{at}(x)} = \frac{\pi_\textrm{co/at}(x)}{\pi_\textrm{at}(x)} - 1 \] is the odds ratio a person with covariates \( X=x \) is a Complier rather than an Always Taker. Using a decision tree, we can identify the covariates that have the largest effect on this ratio. These covariates are the most important differences between Compliers and Always Takers. This model may also provide an audience for growth marketing, by identifying people most receptive to nudges.

Building an uplift model for retention estimates the effect of the splash screen, \( \tau(x), \) for a person with covariates \( X=x. \) This effect is an average over the Never Takers, Compliers, and Always Takers having those covariate values. The effect among the Compliers is \( \tau_\textrm{co}(x) = \tau(x) / \pi_\textrm{co}(x). \) Comparing the groups most and least likely to be Compliers, with the groups having the highest and lowest Complier treatment effects, provides insight about how the effect of the feature among Always Takers may differ from the Compliers.

6 Dose-Response Models

§§ 4– 5 highlighted the value of zero direct effects, when this assumption is reasonable. This value extends to the case of continuous mediators and outcomes.

As an example, consider testing the release of a new feature. We might be interested, not just in the impact of releasing the feature, but also whether people who use the feature more experience greater impact to an outcome like engagement or revenue. We might define the mediator as the amount of time spent interacting with the feature on the first day, and the outcome as some estimate of Lifetime Value (see also § 8). The treatment could be a feature gate (so that people in control are prevented from using the feature) or some growth tactic designed to increase feature usage.

In this case, the potential outcomes for each person are a function \( Y_i(z, m). \) The assumption of zero direct effect is \( Y_i(0, m) = Y_i(1, m) \) for each person \( i \) and all \( m. \) Technically, we need only a weaker assumption, \( Y_i(0, M_i(0)) = Y_i(1, M_i(0)), \) but it is hard to imagine a situation where the latter is plausible and the former isn’t.

Assuming \( Y_i \) is differentiable with respect to \( m, \) the mean value theorem guarantees a value \( m_i \) between \( M_i(0) \) and \( M_i(1) \) so that

\begin{align*} Y_i(1, M_i(1)) - Y_i(1, M_i(0)) &= Y_i^\prime(1, m_i) \cdot (M_i(1) - M_i(0)) \\ &=: \beta_i \cdot \tau_i^M, \end{align*}

where \( Y_i^\prime(1, m_i) \) denotes the derivative of \( Y_i \) with respect to the mediator, \( \beta_i = Y_i^\prime(1, m_i) \) and \( \tau_i^M \) is the effect of treatment on the mediator.

The total effect of treatment on the outcome is \( \tau_i^Y := Y_i(1, M_i(1)) - Y_i(0, M_i(0)). \) If the direct effect is zero, then

\begin{align*} \tau_i^Y &= Y_i(1, M_i(1)) - Y_i(1, M_i(0)) \\ &= \beta_i \cdot \tau_i^M. \end{align*}

It is unreasonable to expect \( \beta \) to be the same for everyone, but we could interpret the average value as the marginal impact of feature usage on the outcome. If this marginal impact were high, it would make sense to invest in increasing feature usage.

(Baiocchi et al. 2010) showed how to estimate this average marginal impact when the outcomes are not too fat-tailed. First pair individuals, assigning treatment to one person from each pair at random. Next estimate \( \tau^Y \) and \( \tau^M \) as the sample differences in outcome and mediator, respectively. Estimate \( \beta \) as \( \tau^Y / \tau^M. \) Pairing is not essential to calculate this estimate, but it is necessary for calculating confidence intervals. The result relies on a central limit theorem holding, which fat-tailed outcomes will violate.

7 Outcomes with Missing Data

In some instances, the mediator must be 1 to observe a meaningful outcome (Rubin 2006). For example, suppose the mediator indicates whether the person responded to a survey question, and the outcome corresponds to the response given. If the person does not respond, the outcome is missing, which is distinct from any observed response.

Or perhaps we have some product funnel, a sequence of steps a person must take in a particular order, such as adding something to their cart on an ecommerce site, checking out, and completing purchase. The mediator could indicate whether a person completed one stage in the funnel, and the outcome could indicate whether the person completed a later stage.

The mediator could indicate whether we have acquired a customer, and the outcome reflects their satisfaction, engagement, or retention. For a person who does not become a customer, it does not make sense to talk about their retention. In these scenarios, let \( \mathord{?} \) denote a missing or undefined outcome. Then \( Y_i(z, 0) = \mathord{?} \) for \( z=0,1. \)

Example 3: We will consider the impact of marketing on purchases of concert tickets. Marketing may cause some people to purchase tickets who otherwise wouldn’t, but it also may cause some people to buy better seats than they otherwise would. We want to understand not only the impact of marketing on the conversion rate, but also on average order value. In this case, the mediator indicates whether a person purchased, and the outcome is the order value. It may seem natural to represent the order value of a non-purchaser as zero, but technically a purchase with value 0 is distinct from a non-purchase, which we will represent with \( \mathord{?}. \)

Analyzing the impact of marketing on conversion rate is straightforward. Analyzing the impact on average order value is not, because the denominator—number of purchases—may be affected by the treatment. Table  6 shows an example. Marketing increased the conversion rate by 0.1 percentage points (absolute lift), but the average order value—calculated as total purchase value divided by number of purchases—is slightly lower in the Marketing group. Total revenue is about 8% higher in the Marketing group, so overall it’s clear marketing was beneficial, but could it be that marketing is causing people to spend less on tickets?

Analyzing this case study through the language of principal strata and (in)direct effects clarifies the situation. We again assume there are no Defiers—no one who is dissuaded from purchasing by marketing. There’s not much to say about the Never Takers. These are people who do not purchase regardless of their treatment assignment.

In many cases, it doesn’t make sense to talk about the effect of treatment on the order value among the Compliers. Causal estimands are comparisons of potential outcomes, and the potential outcome \( Y_i(0, M_i(0)) = \mathord{?} \) for Compliers. In the concert ticket example, we may think of the impact of marketing among Compliers as taking the average order value from $0 to the average value of \( Y_i(1, M_i(1)) \) among Compliers, \( \bar{Y}_\textrm{co}^1. \)

In other examples, the latter quantity is of primary interest. It could represent the average response among people who only responded because of the treatment, or the average satisfaction among customers acquired by the treatment. We may wish to know how this value compares to the analogous quantity for the Always Takers.

Estimating this quantity is nontrivial since we don’t know which individuals are Compliers. Consider the direct effect, \( Y_i(1, M_i(0)) - Y_i(0, M_i(0)) = \mathord{?} - \mathord{?}. \) Although it’s debatable whether it makes sense to subtract these undefined outcomes, in any real dataset we would coalesce them to some numeric value, such as by defining \[ Y_i^\prime(z, m) = \begin{cases} 0 & Y_i(z, m) = \mathord{?} \\ Y_i(z, m) & \textrm{otherwise.} \end{cases} \] This is especially natural in the concert ticket example. Then the direct effect would be 0. The average total effect (in terms of \( Y^\prime \)) among Compliers equals the average indirect effect equals \( \bar{Y}_\textrm{co}^1: \)

\begin{align*} & \frac{1}{N_\textrm{co}} \sum_{i \in \textrm{Compliers}} Y_i^\prime(1, M_i(1)) - Y_i^\prime(0, M_i(0)) \\ & = \frac{1}{N_\textrm{co}} \sum_{i \in \textrm{Compliers}} Y_i^\prime(1, M_i(1)) - Y_i^\prime(1, M_i(0)) \\ & = \frac{1}{N_\textrm{co}} \sum_{i \in \textrm{Compliers}} Y_i^\prime(1, M_i(1)) \\ & = \frac{1}{N_\textrm{co}} \sum_{i \in \textrm{Compliers}} Y_i(1, M_i(1)) \\ & = \bar{Y}_\textrm{co}^1, \end{align*}

where \( N_\textrm{co} \) is the number of compliers.

It makes sense to talk about the effect of treatment on the outcome among the Always Takers, since both potential outcomes are defined. In the concert ticket example, this is the effect of marketing on order value among people who purchased, but not because of marketing. Since the indirect effect is zero in this group, the total effect equals the direct effect.

Estimates of direct and indirect effects apply respectively to the two strata. The direct effect is non-zero only among Always Takers, and the indirect effect is non-zero only among Compliers. Estimating and interpreting the indirect and direct effects would be straightforward if we knew which individuals were Compliers and which were Always Takers.

Since this is a made up example, we can peek behind the curtain. Table  7 shows the average order value for each stratum, without and with marketing. The average order value among Compliers exposed to marketing is $40, quite a bit lower than the value for Always Takers, with or without marketing. The impact of marketing among Always Takers is to increase the average order value by $4, representing half the revenue impact of marketing.

After examining Table  6, it seemed marketing may reduce order value, but Table  7 makes it clear that marketing increases order value across the board. Order value among those who purchased because of marketing is lower than among the organic purchasers. The observed purchases in the Marketing group are a blend of Compliers and Always Takers. As a result, the average order value in the Marketing group is lower than in the Holdout group, because it’s dragged down by the low order value of the Compliers. Neglecting principal strata can easily lead to the wrong conclusions!

In reality, we would never observe a table like Table  7, only Table  6. In § 9, we discuss a strategy for estimating which individuals are Compliers vs Always Takers, which may be practical if we can reliably predict purchase behavior.

In §§ 10– 11, we discuss strategies for estimating average direct and indirect effects. If \( \tau_\textrm{ie} \) is the average indirect effect (over everyone), then \( \tau_\textrm{ie} = \tau_\textrm{ie; co} \cdot \pi_\textrm{co}, \) where \( \tau_\textrm{ie; co} \) is the average indirect effect for Compliers. Then \( \tau_\textrm{ie; co} = \bar{Y}_\textrm{co}^1 = \tau_\textrm{ie} / \pi_\textrm{co}. \) If \( \tau_\textrm{de} \) is the average direct effect (over everyone), then \( \tau_\textrm{de} = \tau_\textrm{de; at} \cdot \pi_\textrm{at} \) and \( \tau_\textrm{de} / \pi_\textrm{at} \) is the average treatment effect among the Always Takers.

In our concert ticket example, the average treatment effect is to increase revenue per person—not per purchaser—from $1 to $1.08. This is easily calculated in any A/B test. The average treatment effect is 8¢. If we estimated the average indirect effect as 4¢, then the average indirect effect for Compliers is 4¢ / 0.001 = $40, since Compliers represent 0.1% of the population in this example. This is the average order value among Compliers in the Marketing group. The average direct effect by definition is the average total effect (8¢) minus the average indirect effect (4¢), or 4¢. The average direct effect among Always Takers is 4¢ / 0.01 = $4, since Always Takers are 1% of the population. The average order value for Always Takers is observed in the Holdout group: $100. The average order value for Always Takers in the Marketing group is $100 + $4 = $104. Thus, we can reconstruct Table  7 given estimates of average direct and indirect effects.

The discussion in this section also applies any time we know \( Y_i(0, 0) \) and \( Y_i(1, 0) \) for everyone. For example, there are scenarios where \( M_i = 0 \) implies \( Y_i = 0. \) Then the total effect for each Never Taker, \( Y_i(1, 0) - Y_i(0, 0), \) is known.

The direct effect for each Complier, \( Y_i(1, 0) - Y_i(0, 0), \) is also known. The total effect is known for each Complier in the treatment group:

\begin{align*} Y_i(1, M_i(1)) - Y_i(0, M_i(0)) &= Y_i(1, 1) - Y_i(0, 0) \\ &= Y_i^\textrm{obs} - Y_i(0, 0), \end{align*}

and by random assignment the average total effect in the control group equals the average total effect in the treatment group. The average indirect effect for Compliers equals the average total effect minus the average direct effect.

The average total effect for Always Takers may be found by removing the total effects from the Never Takers and Compliers:

\begin{align*} \tau &= \pi_\textrm{nt} \cdot \tau_\textrm{nt} + \pi_\textrm{co} \cdot \tau_\textrm{co} + \pi_\textrm{at} \cdot \tau_\textrm{at} \\ \tau_\textrm{at} &= \frac{1}{\pi_\textrm{at}} \cdot \left( \tau - \pi_\textrm{nt} \cdot \tau_\textrm{nt} + \pi_\textrm{co} \cdot \tau_\textrm{co} \right), \end{align*}

where \( \tau_g \) is the average total effect in stratum \( g. \)

8 Surrogate Indices

In this section we move beyond scalar, binary mediators and consider instead a vector of mediators believed to explain most of the effect of the treatment on the outcome. The potential outcomes for an individual are no longer a set of four values, but a function defined on the product space of treatment assignments and mediator values.

The decomposition into direct and indirect effects still makes sense in this context:

\begin{align*} \textrm{Total Effect}_i &= \{ Y_i(1, \vec{M}_i(1)) - Y_i(1, \vec{M}_i(0)) \} + \{ Y_i(1, \vec{M}_i(0)) - Y_i(0, \vec{M}_i(0)) \} \\ &=: \textrm{Indirect Effect}_i + \textrm{Direct Effect}_i. \end{align*}

If the direct effect is negligible, call \( \vec{M} \) a surrogate index (Athey et al. 2019). A surrogate index is a set of mediator variables that collectively capture the majority of the treatment’s effect on the outcome.

We seek a surrogate index because the outcome, \( Y, \) is in some sense expensive to observe. For example, \( Y \) might be a long-term outcome while \( \vec{M} \) is observable in the short-term. Or \( Y \) is lower signal-to-noise than \( \vec{M}, \) so it’s infeasible to measure the treatment effect on \( Y \) precisely, but it is straightforward to measure the effect on \( \vec{M}. \)

The idea then is to use estimates of the effect of treatment on the surrogate index to estimate the effect of treatment on the outcome:

\begin{align*} \tau_Y &:= \frac{1}{n} \sum_{i=1}^n Y_i(1, \vec{M}_i(1)) - Y_i(0, \vec{M}_i(0)) \\ &\approx \frac{1}{n} \sum_{i=1}^n Y_i(1, \vec{M}_i(1)) - Y_i(1, \vec{M}_i(0))\\ &\approx f\left( \frac{1}{n} \sum_{i=1}^n \vec{M}_i(1) - \vec{M}_i(0) \right) \\ &=: f(\tau_M). \end{align*}

Although only noisy estimates of \( \tau_Y \) may be available from past experiments, they are unbiased, and by combining multiple studies, we may construct a model for \( f \) (Bibaut et al. 2023).

9 Multiple Imputation

Although we do not observe the principal stratum to which a person belongs, we may be able to predict the counterfactual mediator value, and thereby the principal strata. Imputed counterfactual mediator values allow us to probabilistically assign individuals to principal strata, such as Never Takers, Compliers, or Always Takers. This allows us to check assumptions not otherwise supported by the data, and offers alternative estimation strategies when those assumptions are violated.

In the cases considered so far, we assumed there were no Defiers. With estimates of principal strata, we could check this assumption. In the case of one- or two-sided non-compliance, we assumed no direct effects among Never Takers and Always Takers. Since the indirect effects are zero among these groups, a non-zero treatment effect among these groups contradicts the assumption of no direct effect.

We can estimate the average treatment effect among Compliers in this scenario, but if we doubt the zero direct effect assumption among Never Takers and Always Takers, we should also doubt it among Compliers. We should then not interpret the average treatment effect as the effect of the mediator. If we believe the direct effects are homogenous across strata, we can estimate the average indirect effect among Compliers by subtracting the average direct effect—estimated among the Never and Always Takers—from the average total effect among Compliers. There are just as many assumptions here as in § 4 and § 5, but perhaps these assumptions are reasonable in situations where the other set aren’t.

A model for the counterfactual mediator value would typically involve uncertainty. With a Bayesian posterior on the counterfactual mediator, we could use the method of Multiple Imputation to incorporate this uncertainty (Rubin 2004). The basic idea is to create a collection of imputed counterfactual mediator values, similar to Bootstrap replications (Tibshirani and Efron 1993). Each imputation in the collection implies a principal stratification, resulting in some quantity calculated (such as the average treatment effect among Compliers). The average of these quantities over the collection becomes the estimated value, and the method permits calculating confidence intervals reflecting the extra uncertainty associated with the imputation.

10 Regression Models

Although the direct and indirect effects are not identified by a randomized experiment, we can sometimes estimate them by a regression model. (VanderWeele 2015, § 2) fits two models:

\begin{align*} M &= \beta_0 + \beta_1 \cdot Z + \beta_2^T \cdot X \\ Y &= \theta_0 + \theta_1 \cdot Z + \theta_2 \cdot M + \theta_3^T \cdot X \end{align*}

and estimates the direct and indirect effects as \( \theta_1 \) and \( \beta_1 \cdot \theta_2, \) respectively. These estimates have a causal interpretation only if certain criteria hold. When we assign the treatment at random, we need \( X \) to contain all mediator-outcome confounders. The treatment may not affect these confounders. (VanderWeele 2015, § 3) discusses sensitivity analysis for unobserved confounding.

Armed with estimates for direct and indirect effects, we may explore the plausibility of the no-direct-effect assumption of § 4, § 5, and § 8. We may estimate the effect of treatment on the outcome among the Always Takers in § 7, or the average outcome among treated Compliers.

11 Matching Methods

Matching provides an alternative to regression models for estimating the indirect effect with binary mediators. Since the difference in outcome between treatment and control estimates the total effect, we may estimate the direct effect as the total effect minus the indirect effect.

The average indirect effect is the average effect of the mediator on the outcome for people in the treatment group. Since we do not assign the mediator at random, we need to adjust for confounding factors associated with both mediator and outcome. The treatment itself is one such factor.

To each person in the treatment group having \( M_i^\textrm{obs} = 1, \) we pair someone in the treatment group having \( M_i^\textrm{obs} = 0 \) and similar values of other observed covariates. Then we can do a paired t-test to estimate the effect of the mediator on the outcome. Systematic differences in outcomes within pairs cannot be a treatment effect, since the treatment is the same within each pair. This approach therefore allows us to estimate indirect effects leveraging the powerful techniques developed by Paul Rosenbaum and others over the last few decades. A good introduction to these methods is (Rosenbaum 2019).

12 Summary

Principal stratification and (in)direct effects allow us to articulate what quantities are most relevant in many studies. Estimating these quantities is another matter: without extra assumptions, we cannot estimate the stratum to which an individual belongs, nor can we estimate direct or indirect effects. Echoing (Tukey 1986), a combination of some data and a fervent desire for an answer is no guarantee we may find that answer in those data. Nevertheless, knowing what conclusions we may draw from what datasets can help identify what dataset would be able to give an acceptable answer.

Adopting a potential outcomes approach, not just for the outcome but also for the mediator, provides an effective framework for reasoning about non-compliance, missing data, and surrogacy. We may use techniques like Multiple Imputation, regression models, and matching methods to estimate the relevant quantities and capture the associated uncertainty.

When solving problems with this framework, ask:

• What is the treatment? What is the mediator? What is the outcome?
• Is it reasonable to say that certain principal strata cannot exist?
• Is it reasonable to predict the counterfactual mediator for anyone?
• Can we say anything about the direct effect for anyone?
• Is the mediator in some sense a prerequisite for the outcome?

13 References