Experimentation has many special designs, each warranting a specialized type of model. Managing all of these can be difficult. Keeping the implementation sane is harder. In this post, we show that four features of experimental design, in any combination, are covered by one model, which we can fit efficiently. Let’s start with the motivation.

A streaming service wants to know whether a new playback setting helps. Accounts are shared by a household, and each account has up to five profiles. The setting changes how the app behaves on the household’s devices, so the service randomizes at the account level: it shows the account a prompt that invites its profiles to turn the setting on. Each profile then decides for itself. The service is cautious, so it ramps the prompt over four days, from 5% of new accounts to 50%. It measures each profile’s hours watched and its rebuffers per play.

This is a cluster randomized, encouragement experiment, with propensities that change by cohort, and with a ratio metric. In previous posts we analyzed several of these features separately: ratio metrics with a bivariate regression and the delta method, local average treatment effects with an instrument, and CUPED from aggregated data. This post claims that one model covers every combination of these features, and that the model can be fit exactly from a compressed dataset whose size does not grow with the number of rows.

The data in this post is synthetic and the scenario is educational. It does not describe a production system.

The Running Example

Each account $a$ enters the experiment on $\text{day}_a \in \{1, 2, 3, 4\}$, and the accounts that enter on the same day form a cohort. An account is prompted, $Z_a = 1$, with probability $e_{\text{day}}$, which ramps up from day to day. Each profile $p$ of the account has covariates measured before the experiment: hours per day $x_1$, tenure in years $x_2$, whether it is a kids profile, and its main device. A prompted profile turns the setting on, $T_{ap} = 1$, more often if it watches a lot, and less often if it is a kids profile. Take-up also depends on an engagement level that we do not observe, which also raises the profile’s hours. Profiles of the same account share network quality, so their outcomes are correlated.

Code

simulate_accounts <- function(n_accounts = 20000, propensity = c(0.05, 0.10, 0.25, 0.50), force_Z = NULL,
                              full_compliance = FALSE, one_profile = FALSE) {
  accounts <- tibble(
    account = seq_len(n_accounts),
    day = sample.int(length(propensity), n_accounts, replace = TRUE, prob = rev(seq_along(propensity))),
    n_profiles = if (one_profile) 1L else pmin(1L + rpois(n_accounts, 1.2), 5L),
    quality = rnorm(n_accounts, 0, 0.5),                       # shared by the account's profiles
    engagement = rnorm(n_accounts, 0, 0.5)                     # account part of engagement
  ) |>
    mutate(Z = rbinom(n(), 1, propensity[day]), Z = if (is.null(force_Z)) Z else force_Z)
  accounts[rep(accounts$account, accounts$n_profiles), ] |>
    mutate(
      profile = sequence(accounts$n_profiles),
      kids = rbinom(n(), 1, if_else(profile == 1, 0.05, 0.3)),   # the first profile is rarely a kids profile
      device = sample(c("tv", "mobile", "web"), n(), replace = TRUE, prob = c(0.5, 0.35, 0.15)),
      engagement_p = engagement + rnorm(n(), 0, 0.8),          # the profile's engagement, not observed
      x1 = rlnorm(n(), 0.7 - 0.12 * day - 0.6 * kids + 0.3 * engagement_p, 0.6),
      x2 = rgamma(n(), 2, 1) * (1.2 - 0.15 * day),
      complier = rbinom(n(), 1, plogis(-0.3 + 0.5 * log(x1) + 0.6 * engagement_p - 0.8 * kids)),
      T = if (full_compliance) Z else Z * complier,
      y = 0.5 + 0.8 * x1 + 0.1 * x2 + 0.5 * engagement_p + quality +                 # hours per day
        T * (0.3 + 0.15 * x1 - 0.2 * kids) + rnorm(n(), 0, 0.5 + 0.3 * x1),
      D = 1 + rpois(n(), exp(1 + 0.4 * log(x1) + 0.3 * engagement_p + 0.2 * (device == "tv")) * (1 + 0.15 * T)),      # plays
      N = rpois(n(), D * 0.08 * exp(-0.4 * T * (1 + 0.5 * (device == "mobile")) + 0.4 * quality - 0.2 * engagement_p))  # rebuffers
    )
}

propensity <- c(0.05, 0.10, 0.25, 0.50)
set.seed(2026)
profiles <- simulate_accounts()

Each row of the data is one profile. Table 1 shows the rows of two accounts: a prompted account in which some profiles turned the setting on and some did not, and an account that was not prompted.

Code

example_accounts <- c(
  profiles |>
    group_by(account) |>
    filter(first(Z) == 1, n() >= 3, any(T == 1), any(T == 0)) |>
    pull(account) |>
    first(),
  profiles |>
    group_by(account) |>
    filter(first(Z) == 0, n() >= 2) |>
    pull(account) |>
    first())
profiles |>
  filter(account %in% example_accounts) |>
  select(account, profile, day, Z, T, kids, device, x1, x2, y, D, N) |>
  kable(digits = 2)

Table 1: The rows of two accounts. T is the treatment received, y the hours per day, D the plays and N the rebuffers.

account	profile	day	Z	T	kids	device	x1	x2	y	D	N
1	1	2	0	0	0	web	1.89	2.32	0.05	6	2
1	2	2	0	0	1	mobile	0.87	1.24	0.38	2	0
15	1	1	1	1	0	web	1.67	1.56	3.02	7	0
15	2	1	1	0	1	mobile	2.21	1.66	1.96	5	1
15	3	1	1	0	0	mobile	0.70	2.37	0.43	5	1
15	4	1	1	0	1	tv	0.48	1.22	0.17	2	0
15	5	1	1	0	1	tv	5.61	3.48	4.87	6	0

Table 2 summarizes the experiment by cohort.

Code

profiles |>
  group_by(day) |>
  summarise(e_day = propensity[first(day)], accounts = n_distinct(account), profiles = n(),
            prompted = mean(Z), takeup = mean(T[Z == 1]), kids = mean(kids), mean_x1 = mean(x1)) |>
  kable(digits = 3, col.names = c("Day", "$e_{\\text{day}}$", "Accounts", "Profiles", "Share prompted", "Take-up", "Share kids", "Mean $x_1$"))

Table 2: One simulated experiment, by cohort. Take-up is the share of prompted profiles that turned the setting on.

Day	$e_{\text{day}}$	Accounts	Profiles	Share prompted	Take-up	Share kids	Mean $x_1$
1	0.05	8051	17547	0.051	0.444	0.188	2.038
2	0.10	5955	13142	0.109	0.427	0.194	1.796
3	0.25	3969	8816	0.245	0.442	0.183	1.601
4	0.50	2025	4484	0.500	0.418	0.178	1.400

Why not add up each account’s profiles first, and analyze one row per account? Because the treatment, the covariates and the question all live at the profile level. Within a prompted account, some profiles turned the setting on and some did not, so an account total mixes treated and untreated profiles. A kids profile and an adult profile watch differently and respond differently, and an account total mixes those too. And the business question is about profiles: how much does the setting change a profile’s hours, and its rebuffers per play, when the profile actually uses it? So we model each profile as a row, $y_{ap}$, and handle the correlation between profiles of the same account in the standard errors.

Because we wrote the simulation, we can compute the true effects. We simulate the same accounts twice, once with every account prompted and once with none, and look at the profiles that turn the setting on when prompted, the compliers. Our two targets are the effect of the setting on compliers’ hours, $\tau_{LATE}$, and on compliers’ rebuffers per play,

\[\tau_R = \frac{E[N(1) \mid C]}{E[D(1) \mid C]} - \frac{E[N(0) \mid C]}{E[D(0) \mid C]},\]

where $N(t)$ and $D(t)$ are a profile’s rebuffers and plays with the setting on, $t = 1$, or off, $t = 0$, and $C$ denotes the compliers.

Code

truth_pair <- function(seed, n_accounts = 1e6) {
  set.seed(seed); p1 <- simulate_accounts(n_accounts, force_Z = 1)
  set.seed(seed); p0 <- simulate_accounts(n_accounts, force_Z = 0)
  C <- p1$complier == 1                                      # the same profiles comply in both worlds
  c(tau_LATE = mean(p1$y[C] - p0$y[C]),
    tau_R = sum(p1$N[C]) / sum(p1$D[C]) - sum(p0$N[C]) / sum(p0$D[C]))
}
truth <- colMeans(do.call(rbind, map(1:2, truth_pair)))
truth

   tau_LATE       tau_R 
 0.61334648 -0.02724864

One Universal Model

The design

The experiment has four features, and each one changes the analysis. Keeping track of the complexity is hard, but we will collapse these into a single, universal model.

Clusters. We randomize accounts, and profiles of the same account are correlated, so the standard errors must be cluster robust.
Encouragement. We randomize the prompt $Z_a$, but the treatment received $T_{ap}$ is chosen by the profile, so $Z$ is an instrument for $T$.
A ramp. The probability of a prompt depends on the cohort, so each account gets the inverse propensity weight $w_a = 1/\hat e_{\text{day}}$ if prompted and $w_a = 1/(1 - \hat e_{\text{day}})$ if not, where $\hat e_{\text{day}}$ is the share of the accounts entering on that day that were prompted.
A ratio metric. Rebuffers per play needs a numerator and a denominator, modeled jointly.

We rely on three assumptions.

(A1) Accounts are independent of each other. Profiles of the same account may be correlated in any way.

(A2) Within a cohort, the prompt is randomized with probability $e_{\text{day}}$, independently of everything else.

(A3) The prompt affects a profile’s outcomes only through $T$ (exclusion), and no profile turns the setting on only when it is not prompted (monotonicity). Monotonicity holds trivially here, because an unprompted profile cannot turn the setting on.

The model

Two stage least squares with weights and cluster robust standard errors handles all four features at once. It is the universal model, and is formulated as

\[\begin{align} y_{ap} &= \underbrace{\beta_{00} + (x_{ap} - x_{ref})^\top \beta_{0x}}_{\text{baseline line}} + T_{ap} \underbrace{\left(\beta_{10} + (x_{ap} - x_{ref})^\top \beta_{1x}\right)}_{\text{treatment line}} + \varepsilon_{ap} && \text{second stage} \\ T_{ap} &= \gamma_{00} + (x_{ap} - x_{ref})^\top \gamma_{0x} + Z_a \left(\gamma_{10} + (x_{ap} - x_{ref})^\top \gamma_{1x}\right) + \nu_{ap} && \text{first stage for } T \\ T_{ap} (x_{apj} - x_{ref,j}) &= \gamma^{(j)}_{00} + (x_{ap} - x_{ref})^\top \gamma^{(j)}_{0x} + Z_a \left(\gamma^{(j)}_{10} + (x_{ap} - x_{ref})^\top \gamma^{(j)}_{1x}\right) + \nu^{(j)}_{ap} && \text{first stage for interaction } j = 1, \dots, k \end{align}\]

Here $x_{ap}$ is the vector of the profile’s $k = 5$ covariates: hours per day, tenure, the kids indicator, and the device as two indicators, mobile and web, so that tv is the baseline. $x_{apj}$ is its $j$-th entry, and $x_{ref}$ is a fixed reference point (see Centering the covariates).

The second stage describes a profile’s outcome with two lines in its covariates. The baseline line is the outcome without the setting, and the treatment line is the effect of the setting. $\beta_{10}$ is the main effect of the treatment, the effect at $x = x_{ref}$, and $\beta_{1x}$ lets the effect change with each covariate. Every covariate interacts with the treatment, but the covariates do not interact with each other.

The first stage handles the encouragement. Profiles choose whether to turn the setting on, so profiles with $T = 1$ differ from profiles with $T = 0$ in ways other than the setting. The treatment $T$ and its $k$ interactions $T (x_j - x_{ref,j})$ are therefore endogenous, and each gets its own first stage equation. Each equation has the same form as the second stage, with the randomized prompt $Z_a$ in place of $T$: it predicts the endogenous regressor from the covariates, the prompt, and the prompt’s interactions with the covariates. The second stage is then fit with these predictions in place of $T$ and its interactions.

Three more choices complete the model. Each account carries its inverse propensity weight $w_a$, which handles the ramp. The standard errors add up the profiles of an account before squaring, which makes them cluster robust. And one fit can hold several outcome columns: hours $y$, and for the ratio metric, rebuffers $N$, plays $D$, and two more columns, $TN$ and $TD$ (see What the model estimates).

What the model estimates

When applying our universal model to this cluster randomized encouragement experiment, the model will be estimating the Local Average Treatment Effect (LATE). However, when modeling other experiments, this universal model can also model an average treatment effect, and an intent to treat effect.

The treatment line describes the compliers. The prompt changes the treatment only for compliers, the profiles that turn the setting on when prompted. Every other profile has the same $T$ with or without the prompt, so the first stage, and the treatment line it feeds, can only learn from compliers. The treatment line is a fit to the compliers’ average effect at each value of the covariates, $\tau_C(x)$, where each $x$ counts in proportion to how many compliers have it. The appendix derives this from the model’s equations.

The LATE is the treatment line at the compliers’ average covariates. Because the line is fit to the compliers, its average over the compliers is their average effect, the LATE. A line’s average is its value at the average input, so

\[\boxed{\tau_{LATE} = \beta_{10} + (\bar x_C - x_{ref})^\top \beta_{1x}}\]

where $\bar x_C$ is the mean of the covariates among compliers. This holds even when $\tau_C(x)$ is not a line. We never observe who is a complier, but we can estimate their mean covariates with a ratio of two intent to treat effects. The product $T \cdot x$ is nonzero only for profiles that took the treatment, so under (A3)

\[\bar x_C = \frac{E[T \cdot x \mid Z = 1] - E[T \cdot x \mid Z = 0]}{E[T \mid Z = 1] - E[T \mid Z = 0]}.\]

In our experiment the compliers watch more than the average profile, so reading the line at the average profile would answer a different question than the LATE.

The ratio needs the compliers’ own levels. Write $\ell_N$ for the treatment line of the outcome $N$ read at $\bar x_C$, and similarly for other outcomes. $\ell_N$ is the compliers’ effect on rebuffers, but $\tau_R$ needs the compliers’ levels, $E[N(1) \mid C]$ and $E[N(0) \mid C]$. The baseline line does not give them: it describes all profiles, and compliers differ from the rest in ways that $x$ does not capture. The column $TN$ does. It equals $N$ for profiles that turned the setting on and 0 otherwise, so its treatment line is fit to the compliers’ rebuffers with the setting on, and $\ell_{TN} = E[N(1) \mid C]$. Subtracting the effect gives $E[N(0) \mid C] = \ell_{TN} - \ell_N$. The same holds for plays, so

\[\tau_R = \frac{\ell_{TN}}{\ell_{TD}} - \frac{\ell_{TN} - \ell_N}{\ell_{TD} - \ell_D}.\]

These are the two effects of our running example. The same fitted model answers other questions too.

The effect for a type of profile. Reading the treatment line at other covariates gives the effect for profiles like those, for example kids profiles that watch one hour a day.
How the effect varies. The slopes $\beta_{1x}$ are the heterogeneity in the effect, for example how much smaller it is for kids profiles.
Simpler designs. Each feature of our experiment can be switched off, and switching it off gives the textbook analysis of a simpler design:
- No encouragement: when every prompted profile takes the treatment, $T = Z$, the first stage predicts $T$ perfectly, and the model is weighted least squares.
- No clusters: when each account has one profile, the cluster robust standard errors are the heteroskedasticity robust standard errors.
- No ramp: when the propensity is the same every day, the weights are a constant and cancel.
- No ratio: when every row has $D = 1$, the ratio reduces to the LATE on $N$.
- No covariates: the model is a difference in means, or a Wald ratio with encouragement.

The switches touch different parts of the model, the first stage, the standard errors, the weights and the outcome columns, so they combine freely. A purely randomized AB test with one row per user and a mean metric is the model with every switch off; our experiment is the model with every switch on. The appendix proves each reduction, and Table 3 in Checking every design tests all 16 combinations against textbook code.

How the model is fit

In this section we describe how to train the model from scratch using first principles. It is possible to use off the shelf software like iv_robust from estimatr, but we write out all the steps and linear algebra here in order to be prepared for how to derive the compression formulas.

Centering the covariates

We center each covariate at a fixed reference point $x_{ref}$. Centering works like demeaning the columns, except that $x_{ref}$ is a constant we choose before the experiment rather than the sample mean. Any choice gives the same fitted model; it only changes which point the coefficients $\beta_{00}$ and $\beta_{10}$ describe. We use a constant because the compressed version of the model, later in the post, is built in one pass over the data, before the sample mean is known. Here $x_{ref}$ is 2 hours per day, 2 years of tenure, an adult profile and a tv. From here on, $x_{ap}$ denotes the centered covariates, $x_{ap} - x_{ref}$, so the reference point is $x = 0$. With an intercept in front, the profile rows of the prompted account from Table 1 become

Code

x_cols <- c("x1", "x2", "kids", "mobile", "web")
x_ref <- c(x1 = 2, x2 = 2, kids = 0, mobile = 0, web = 0)

prepare <- function(profiles, ramp = TRUE) {
  e_hat <- profiles |>
    distinct(account, day, Z) |>
    group_by(day) |>
    summarise(e_hat = mean(Z))
  profiles |>
    left_join(e_hat, by = "day") |>
    mutate(w = if (ramp) if_else(Z == 1, 1 / e_hat, 1 / (1 - e_hat)) else 1,     # inverse propensity weight
           mobile = as.numeric(device == "mobile"), web = as.numeric(device == "web"),
           TN = T * N, TD = T * D)
}

centered <- function(df) {
  r <- cbind(1, sweep(as.matrix(df[x_cols]), 2, x_ref))
  colnames(r) <- c("1", paste0(x_cols, "_c"))
  r
}

df <- prepare(profiles)
example <- df |>
  filter(account == example_accounts[1])
centered(example)

     1       x1_c       x2_c kids_c mobile_c web_c
[1,] 1 -0.3299705 -0.4370072      0        0     1
[2,] 1  0.2090540 -0.3439723      1        1     0
[3,] 1 -1.2956183  0.3744773      0        1     0
[4,] 1 -1.5190796 -0.7822066      1        0     0
[5,] 1  3.6062466  1.4780119      1        0     0

The two stages as matrices

With centered covariates, the second stage is $y_{ap} = m_{2,ap}^\top \beta + \varepsilon_{ap}$, with regressors $m_{2,ap} = (1, \; x_{ap}, \; T_{ap}, \; T_{ap} x_{ap})$ and coefficients $\beta = (\beta_0, \beta_1)$, where $\beta_0 = (\beta_{00}, \beta_{0x})$ is the baseline line and $\beta_1 = (\beta_{10}, \beta_{1x})$ is the treatment line. The first stage predicts the endogenous regressors $T_{ap}$ and $T_{ap} x_{ap}$ from the instruments $m_{1,ap} = (1, \; x_{ap}, \; Z_a, \; Z_a x_{ap})$,

\[(T_{ap}, \; T_{ap} x_{ap}) = \Gamma^\top m_{1,ap} + \nu_{ap},\]

which stacks the $k + 1$ first stage equations of the model, one for $T_{ap}$ and one for each of the $k$ entries of $T_{ap} x_{ap}$. $\Gamma$ has one column of first stage coefficients per equation, the $\gamma$ and the $\gamma^{(j)}$, and $\nu_{ap}$ stacks the first stage errors. Stacking the rows gives the matrices $M_1$ and $M_2$. Here are their rows for the prompted account. They differ only in their second half, where $M_1$ multiplies by the prompt $Z$ and $M_2$ by the treatment received $T$.

Code

design <- function(df) {
  r <- centered(df)
  list(M_1 = cbind(r, `colnames<-`(df$Z * r, paste0("Z*", colnames(r)))),     # first stage: instruments
       M_2 = cbind(r, `colnames<-`(df$T * r, paste0("T*", colnames(r)))))     # second stage: regressors
}
round(design(example)$M_1, 2)

     1  x1_c  x2_c kids_c mobile_c web_c Z*1 Z*x1_c Z*x2_c Z*kids_c Z*mobile_c
[1,] 1 -0.33 -0.44      0        0     1   1  -0.33  -0.44        0          0
[2,] 1  0.21 -0.34      1        1     0   1   0.21  -0.34        1          1
[3,] 1 -1.30  0.37      0        1     0   1  -1.30   0.37        0          1
[4,] 1 -1.52 -0.78      1        0     0   1  -1.52  -0.78        1          0
[5,] 1  3.61  1.48      1        0     0   1   3.61   1.48        1          0
     Z*web_c
[1,]       1
[2,]       0
[3,]       0
[4,]       0
[5,]       0

Code

round(design(example)$M_2, 2)

     1  x1_c  x2_c kids_c mobile_c web_c T*1 T*x1_c T*x2_c T*kids_c T*mobile_c
[1,] 1 -0.33 -0.44      0        0     1   1  -0.33  -0.44        0          0
[2,] 1  0.21 -0.34      1        1     0   0   0.00   0.00        0          0
[3,] 1 -1.30  0.37      0        1     0   0   0.00   0.00        0          0
[4,] 1 -1.52 -0.78      1        0     0   0   0.00   0.00        0          0
[5,] 1  3.61  1.48      1        0     0   0   0.00   0.00        0          0
     T*web_c
[1,]       1
[2,]       0
[3,]       0
[4,]       0
[5,]       0

With the weights in $W = \text{diag}(w)$, two stage least squares replaces $M_2$ with its first stage prediction $\hat M_2 = M_1 (M_1^\top W M_1)^{-1} M_1^\top W M_2$ and regresses $y$ on it with weights $W$. $M_1$ and $M_2$ have the same number of columns, one instrument for each regressor, so the estimate simplifies to

\[\boxed{\hat\beta = (M_1^\top W M_2)^{-1} M_1^\top W y.}\]

The matrix $M_1^\top W M_2$ is the weighted cross product of the instruments with the regressors, and it plays the role that $X^\top X$ plays in least squares.

The residuals use the treatment actually received, $\hat\varepsilon_{ap} = y_{ap} - m_{2,ap}^\top \hat\beta$. The cluster robust covariance of $\hat\beta$ is the sandwich

\[\boxed{\widehat{\text{Cov}}(\hat\beta) = (M_1^\top W M_2)^{-1} \left[\sum_a \left(\sum_p w_a \, m_{1,ap} \, \hat\varepsilon_{ap}\right) \left(\sum_p w_a \, m_{1,ap} \, \hat\varepsilon_{ap}\right)^\top \right] (M_1^\top W M_2)^{-\top}.}\]

The inner sum adds up the residuals of an account’s profiles, each multiplied by its instruments. Profiles are added within an account before the outer product is taken, so their correlation enters the covariance. That is what makes it robust to clustering.

The outcome can be a matrix $Y$ with one column per metric. All columns share $M_1$, $M_2$ and $M_1^\top W M_2$, so the coefficients form a matrix $\hat B$ with $2(k+1)$ rows and one column per metric. Stacking the inner sums of all the metrics into one vector per account gives the covariance between the coefficients of different metrics, which the ratio needs.

From coefficients to effects

Each effect is a function of $\hat B$. Reading a treatment line at $\bar x_C$ is a dot product with the delta vector $\delta = (0, \dots, 0, \; 1, \; \bar x_C - x_{ref})$, which has zeros for the baseline line. Here $\bar x_C$ is the compliers’ mean of the original covariates, so $\bar x_C - x_{ref}$ is that mean in centered units. Then $\ell_N = \delta^\top \hat\beta_N$ for the column $\hat\beta_N$ of $\hat B$ that belongs to $N$, and the LATE is $\delta^\top \hat\beta_y$. The variance of each effect follows from the delta method, with the gradient of the effect with respect to $\hat B$ and the joint covariance of all the coefficients. We never pick coefficients by hand: we only write the effect as a function of $\hat B$.

fit_rows fits the model on the profile rows. The outcome matrix holds hours $y$ and the four columns for rebuffers per play.

Code

y_cols <- c("y", "N", "D", "TN", "TD")
d <- length(x_cols) + 1                              # intercept and k covariates
o <- length(y_cols)                                  # number of outcome columns

fit_rows <- function(df) {
  m <- design(df); Y <- as.matrix(df[y_cols])
  M1_W_M2 <- crossprod(m$M_1 * df$w, m$M_2)                                   # M_1' W M_2
  B <- solve(M1_W_M2, crossprod(m$M_1 * df$w, Y))
  eps <- Y - m$M_2 %*% B                                                # second stage residuals
  inner <- do.call(cbind, map(seq_len(o), \(l) rowsum(m$M_1 * (df$w * eps[, l]), df$account)))   # one row per account
  M1_W_M2_inv <- kronecker(diag(o), solve(M1_W_M2))
  list(B = B, V = M1_W_M2_inv %*% crossprod(inner) %*% t(M1_W_M2_inv))
}

complier_mean <- function(df) {
  if (all(df$T == df$Z)) return(colSums(df$w * df[x_cols]) / sum(df$w))
  arm <- df |>
    group_by(Z) |>
    summarise(across(all_of(x_cols), \(v) sum(w * T * v) / sum(w)), T = sum(w * T) / sum(w))
  unlist(arm[arm$Z == 1, x_cols] - arm[arm$Z == 0, x_cols]) / (arm$T[arm$Z == 1] - arm$T[arm$Z == 0])
}

effects builds $\delta$, writes each effect as a function of $\hat B$, and computes the gradient with complex step differentiation: adding an imaginary step $i\kappa$ to one coefficient and reading the imaginary part of the effect gives the derivative to machine precision, because nothing is subtracted.

Code

gradient <- function(f, theta, kappa = 1e-20) {
  map_dbl(seq_along(theta), \(i) Im(f(theta + 1i * kappa * (seq_along(theta) == i))) / kappa)
}

effect_functions <- function(x_C, x_ref, cols) {
  delta <- c(rep(0, length(x_ref) + 1), 1, x_C - x_ref)               # reads a treatment line at x_C
  line <- \(B, col) sum(delta * B[, col])
  effect_list <- list(tau_LATE = \(B) line(B, "y"))
  if ("TN" %in% cols) effect_list$tau_R <- \(B) line(B, "TN") / line(B, "TD") - (line(B, "TN") - line(B, "N")) / (line(B, "TD") - line(B, "D"))
  effect_list
}

effects <- function(fit, x_C, x_ref) {
  effect_list <- effect_functions(x_C, x_ref, colnames(fit$B))
  theta <- as.vector(fit$B)
  map_dfr(names(effect_list), \(nm) {
    f <- \(th) effect_list[[nm]](matrix(th, nrow(fit$B), dimnames = list(NULL, colnames(fit$B))))
    grad <- gradient(f, theta)
    tibble(effect = nm, estimate = Re(f(theta)), se = sqrt(drop(t(grad) %*% fit$V %*% grad)))
  })
}

fit <- fit_rows(df)
x_C <- complier_mean(df)
effects(fit, x_C, x_ref) |>
  mutate(truth = truth[effect])

# A tibble: 2 × 4
  effect   estimate      se   truth
  <chr>       <dbl>   <dbl>   <dbl>
1 tau_LATE   0.664  0.0572   0.613 
2 tau_R     -0.0255 0.00410 -0.0272

Both estimates are within 0.9 standard errors of the truth.

Checking every design

Table 3 checks the claim that the model covers every design. For each of the 16 combinations of the switches, with all covariates included, we simulate an experiment with those features and fit it twice: with fit_rows, and with the textbook model for that design from the estimatr package, either lm_robust or iv_robust, with or without weights and clusters. For a mean metric we use hours, and the textbook model reports the estimate and its standard error directly. For a ratio metric, estimatr fits one outcome at a time and does not report the covariance between outcomes, so we compare the estimates only. The appendix checks the ratio’s standard errors by simulation.

Code

textbook <- function(df, x_C, clustered, encouraged, outcome, covs = x_cols) {
  df_c <- df |>
    mutate(across(all_of(covs), \(v) v - x_C[cur_column()]))     # coefficient on T is the line at x_C
  rhs <- paste0("(", paste(covs, collapse = " + "), ")")
  f <- as.formula(if (encouraged) paste(outcome, "~", rhs, "* T |", rhs, "* Z") else paste(outcome, "~", rhs, "* T"))
  estimator <- if (encouraged) iv_robust else lm_robust
  fit <- if (clustered) estimator(f, data = df_c, weights = w, clusters = account, se_type = "CR0")
         else estimator(f, data = df_c, weights = w, se_type = "HC0")
  c(estimate = fit$coefficients[["T"]], se = fit$std.error[["T"]])
}

check_design <- function(encouraged, clustered, ramp, metric) {
  set.seed(7)
  df <- simulate_accounts(propensity = if (ramp) propensity else rep(0.3, 4), full_compliance = !encouraged,
                          one_profile = !clustered) |>
    prepare(ramp = ramp)
  x_C <- complier_mean(df)
  fit_u <- effects(fit_rows(df), x_C, x_ref)
  if (metric == "mean") {
    tb <- textbook(df, x_C, clustered, encouraged, "y")
    tibble(universal = fit_u$estimate[1], textbook = tb[["estimate"]], universal_se = fit_u$se[1], textbook_se = tb[["se"]])
  } else {
    l <- map(set_names(c("N", "D", "TN", "TD")), \(col) textbook(df, x_C, clustered, encouraged, col)[["estimate"]])
    tibble(universal = fit_u$estimate[2], textbook = l$TN / l$TD - (l$TN - l$N) / (l$TD - l$D),
           universal_se = fit_u$se[2], textbook_se = NA_real_)
  }
}

designs <- expand_grid(encouraged = c(FALSE, TRUE), clustered = c(FALSE, TRUE),
                       ramp = c(FALSE, TRUE), metric = c("mean", "ratio"))
design_tab <- designs |>
  mutate(res = pmap(designs, check_design)) |>
  unnest(res)

Code

design_tab |>
  mutate(across(c(encouraged, clustered, ramp), \(v) if_else(v, "yes", "no"))) |>
  kable(digits = 6, col.names = c("Encouragement", "Clusters", "Ramp", "Metric", "Estimate (universal)",
                                  "Estimate (textbook)", "SE (universal)", "SE (textbook)"))

Table 3: Sixteen designs, each fit with the universal model and with the textbook model for that design. For ratio metrics the textbook fits do not report a joint covariance, so only the estimates are compared.

Encouragement	Clusters	Ramp	Metric	Estimate (universal)	Estimate (textbook)	SE (universal)	SE (textbook)
no	no	no	mean	0.557413	0.557413	0.020775	0.020775
no	no	no	ratio	-0.029766	-0.029766	0.001641	NA
no	no	yes	mean	0.581682	0.581682	0.035907	0.035907
no	no	yes	ratio	-0.028284	-0.028284	0.002633	NA
no	yes	no	mean	0.541539	0.541539	0.015395	0.015395
no	yes	no	ratio	-0.029436	-0.029436	0.001141	NA
no	yes	yes	mean	0.534708	0.534708	0.026310	0.026310
no	yes	yes	ratio	-0.031397	-0.031397	0.001748	NA
yes	no	no	mean	0.591617	0.591617	0.044127	0.044127
yes	no	no	ratio	-0.022754	-0.022754	0.003326	NA
yes	no	yes	mean	0.633877	0.633877	0.077722	0.077722
yes	no	yes	ratio	-0.023077	-0.023077	0.005889	NA
yes	yes	no	mean	0.624210	0.624210	0.034929	0.034929
yes	yes	no	ratio	-0.023948	-0.023948	0.002495	NA
yes	yes	yes	mean	0.610768	0.610768	0.057555	0.057555
yes	yes	yes	ratio	-0.021721	-0.021721	0.004115	NA

Every estimate agrees, and so does every standard error that the textbook model reports. The first row is a purely randomized AB test with one row per user and a mean metric, analyzed with Lin’s regression and heteroskedasticity robust errors. The last row is our experiment.

Compressing the Universal Model

The model is universal, but fitting it on the profile rows is expensive when there are hundreds of millions of them and dozens of metrics. We now derive, from first principles, a compressed dataset that still gives the exact estimate and the exact covariance. We start with the math, then build the compressed dataset by hand for two simple cases, and then for the full model.

What the model needs

The estimate needs $M_1^\top W M_2$ and $M_1^\top W Y$, which are sums over profile rows. Within an account, the instruments are $m_{1,ap} = (1, x_{ap}, Z_a, Z_a x_{ap})$, and $Z_a$ and $w_a$ are the same for every profile, so these sums only need three blocks per account:

\[S_a = \sum_p \begin{pmatrix} 1 \\ x_{ap} \end{pmatrix} \begin{pmatrix} 1 & x_{ap}^\top \end{pmatrix}, \qquad S^T_a = \sum_p T_{ap} \begin{pmatrix} 1 \\ x_{ap} \end{pmatrix} \begin{pmatrix} 1 & x_{ap}^\top \end{pmatrix}, \qquad S^Y_a = \sum_p \begin{pmatrix} 1 \\ x_{ap} \end{pmatrix} y_{ap}^\top.\]

$S_a$ holds the account’s number of profiles, the sums of its covariates, and the sums of the products of every pair of covariates. $S^T_a$ holds the same sums over the profiles that took the treatment. $S^Y_a$ holds the sums of the outcomes, and of the outcomes times each covariate; here $y_{ap}$ is the row of $Y$, all outcome columns of the profile. Because $Z_a$ and $w_a$ are constant within an account, each block of $M_1^\top W M_2$ is a weighted sum of these blocks over accounts. For example, the block that pairs the instruments $(Z_a, Z_a x_{ap})$ with the regressors $(T_{ap}, T_{ap} x_{ap})$ is $\sum_a w_a Z_a S^T_a$.

To see what these blocks hold, take the two accounts of Table 1, with one covariate, the pre-period hours $x_1$, and the outcome hours $y$. Account 15 was prompted, and only its first profile turned the setting on. Account 1 was not prompted.

Code

two_accounts <- df |>
  filter(account %in% example_accounts) |>
  select(account, profile, Z, T, x1, y)
two_accounts

# A tibble: 7 × 6
  account profile     Z     T    x1      y
    <int>   <int> <int> <int> <dbl>  <dbl>
1       1       1     0     0 1.89  0.0541
2       1       2     0     0 0.867 0.377 
3      15       1     1     1 1.67  3.02  
4      15       2     1     0 2.21  1.96  
5      15       3     1     0 0.704 0.426 
6      15       4     1     0 0.481 0.173 
7      15       5     1     0 5.61  4.87

With the reference point $x_{ref} = 2$, each profile row gets an intercept and the centered covariate $x_1 - 2$, which we store as two columns, one and x_c.

Code

two_accounts_c <- two_accounts |>
  mutate(one = 1, x_c = x1 - 2)
two_accounts_c |>
  select(account, profile, T, one, x_c, y)

# A tibble: 7 × 6
  account profile     T   one    x_c      y
    <int>   <int> <int> <dbl>  <dbl>  <dbl>
1       1       1     0     1 -0.114 0.0541
2       1       2     0     1 -1.13  0.377 
3      15       1     1     1 -0.330 3.02  
4      15       2     0     1  0.209 1.96  
5      15       3     0     1 -1.30  0.426 
6      15       4     0     1 -1.52  0.173 
7      15       5     0     1  3.61  4.87

The three blocks are sums over the profile rows of an account, so dplyr computes them with one group_by and one summarise. Each block is a small matrix, so we compute one column per entry. $S_a$ and $S^T_a$ are symmetric, so their off-diagonal entries only need to be computed once.

Code

two_sums <- two_accounts_c |>
  group_by(account, Z) |>
  summarise(S_11 = sum(one * one), S_12 = sum(one * x_c), S_22 = sum(x_c * x_c),                   # S_a
            ST_11 = sum(T * one * one), ST_12 = sum(T * one * x_c), ST_22 = sum(T * x_c * x_c),    # S^T_a
            SY_1 = sum(one * y), SY_2 = sum(x_c * y),                                            # S^Y_a
            .groups = "drop")
two_sums

# A tibble: 2 × 10
  account     Z  S_11   S_12  S_22 ST_11  ST_12 ST_22   SY_1   SY_2
    <int> <int> <dbl>  <dbl> <dbl> <dbl>  <dbl> <dbl>  <dbl>  <dbl>
1       1     0     2 -1.25   1.30     0  0     0      0.431 -0.433
2      15     1     5  0.671 17.1      1 -0.330 0.109 10.4   16.1

Arranged as matrices, the blocks of the prompted account are

Code

prompted <- filter(two_sums, account == example_accounts[1])
with(prompted, list(
  S = matrix(c(S_11, S_12, S_12, S_22), 2),
  S_T = matrix(c(ST_11, ST_12, ST_12, ST_22), 2),
  S_Y = matrix(c(SY_1, SY_2), 2)))

$S
          [,1]       [,2]
[1,] 5.0000000  0.6706322
[2,] 0.6706322 17.1438281

$S_T
           [,1]       [,2]
[1,]  1.0000000 -0.3299705
[2,] -0.3299705  0.1088805

$S_Y
         [,1]
[1,] 10.43645
[2,] 16.14321

Each entry is a sum we can read off the rows above. $S_a$ counts the account’s 5 profiles in its first entry, adds up the centered covariate, 0.67, and adds up its squares, 17.14. $S^T_a$ is the same sum over the profiles that took the treatment, which is only the first profile, with an intercept of 1 and a centered covariate of -0.33. $S^Y_a$ adds up the hours, 10.44, and the hours times the centered covariate, 16.14. The other account was not prompted, so none of its profiles took the treatment, and its $S^T_a$ is zero.

The covariance needs the inner sums $\sum_p w_a m_{1,ap} \hat\varepsilon_{ap}$. Split $\hat B$ into its baseline and treatment blocks, $\hat B_0$ and $\hat B_1$. The residuals of account $a$, each multiplied by its intercept and covariates, add up to

\[E_a = \sum_p \begin{pmatrix} 1 \\ x_{ap} \end{pmatrix} \hat\varepsilon_{ap}^\top = S^Y_a - S_a \hat B_0 - S^T_a \hat B_1,\]

and the inner sum is $w_a E_a$ for the instruments $(1, x_{ap})$ and $w_a Z_a E_a$ for the instruments $(Z_a, Z_a x_{ap})$. So $E_a$ only depends on the account’s three blocks. We will collect them into an important intermediate vector and call that the account summary, $u_a$, which has

\[\boxed{u_a = \begin{pmatrix} \text{vec}(S^Y_a) \\ \text{vec}(S_a) \\ \text{vec}(S^T_a) \end{pmatrix} = \begin{pmatrix} \text{vec}\left(\sum_p \begin{pmatrix} 1 \\ x_{ap} \end{pmatrix} y_{ap}^\top\right) \\[4pt] \text{vec}\left(\sum_p \begin{pmatrix} 1 \\ x_{ap} \end{pmatrix} \begin{pmatrix} 1 & x_{ap}^\top \end{pmatrix}\right) \\[4pt] \text{vec}\left(\sum_p T_{ap} \begin{pmatrix} 1 \\ x_{ap} \end{pmatrix} \begin{pmatrix} 1 & x_{ap}^\top \end{pmatrix}\right) \end{pmatrix}}\]

Then $\text{vec}(E_a) = L(\hat B) \, u_a$ for a matrix $L(\hat B)$ that depends only on the coefficients, and

\[\sum_a \text{vec}(E_a) \, \text{vec}(E_a)^\top = L(\hat B) \left(\sum_a u_a u_a^\top\right) L(\hat B)^\top.\]

For our two accounts, each account summary $u_a$ stacks the two entries of $S^Y_a$, then the four entries of $S_a$, then the four entries of $S^T_a$, each matrix read column by column. The off-diagonal entries appear twice. That gives 10 numbers per account:

Code

two_u <- with(two_sums, cbind(SY_1, SY_2, S_11, S_12, S_12, S_22, ST_11, ST_12, ST_12, ST_22))
rownames(two_u) <- paste("account", two_sums$account)
two_u

                 SY_1       SY_2 S_11       S_12       S_12      S_22 ST_11
account 1   0.4313195 -0.4334005    2 -1.2470686 -1.2470686  1.295958     0
account 15 10.4364508 16.1432144    5  0.6706322  0.6706322 17.143828     1
                ST_12      ST_12     ST_22
account 1   0.0000000  0.0000000 0.0000000
account 15 -0.3299705 -0.3299705 0.1088805

Accounts with the same assignment and cohort share $Z_a$ and $w_a$, so the whole model needs only three sums per stratum $h = (Z, \text{day})$:

The number of accounts in stratum $h$, $G_h$.
The first moments $\sum_{a \in h} u_a$, which give the estimate.
The second moments $\sum_{a \in h} u_a u_a^\top$, which give the covariance.

These are conditionally sufficient statistics: given the stratum, they carry everything the model can learn from the data.

From here on we use three names for the data at each step.

A profile row is a row of the original, uncompressed data: one profile of one account, with its day, prompt, treatment, covariates and outcomes, as in Table 1. Our experiment has 43,989 profile rows.
An account summary is the vector $u_a$ of sums over the profile rows of one account. It is an intermediate: we build one per account in a single pass over the profile rows, add it into its stratum’s sums, and discard it.
A compressed row holds the sums of one stratum. It has three fields: the count $G_h$, the first moments $\sum_{a \in h} u_a$, which is a vector, and the second moments $\sum_{a \in h} u_a u_a^\top$, which is a matrix. The compressed dataset has one compressed row per stratum, 8 in our experiment, whatever the number of profile rows.

The functions below turn account summaries into the compressed dataset, and the compressed dataset into the fit. They work for any number of covariates $k$, with $d = k + 1$ for the intercept and the $k$ covariates, and any number of outcome columns $o$.

Code

compress_summaries <- function(accounts, u) {          # accounts: one row per account with Z, day, w; u: one row of u_a per account
  strata <- group_by(accounts, Z, day, w)
  h <- group_indices(strata)
  group_keys(strata) |>
    mutate(G = tabulate(h),
           first_moments = map(seq_len(n()), \(k) colSums(u[h == k, , drop = FALSE])),       # sum of u_a
           second_moments = map(seq_len(n()), \(k) crossprod(u[h == k, , drop = FALSE])))    # sum of u_a u_a^T
}

blocks <- function(s, d, o) {                        # split u (or a sum of u) into S^Y, S and S^T
  list(S_Y = matrix(s[1:(d * o)], d), S = matrix(s[d * o + 1:(d * d)], d), S_T = matrix(s[d * o + d * d + 1:(d * d)], d))
}

fit_first_moments <- function(compressed, d, outcomes) {            # the estimate needs only sum u_a
  o <- length(outcomes)
  b <- map(compressed$first_moments, blocks, d = d, o = o)
  M1_W_M2 <- reduce(map(seq_along(b), \(k) compressed$w[k] * rbind(cbind(b[[k]]$S, b[[k]]$S_T), compressed$Z[k] * cbind(b[[k]]$S, b[[k]]$S_T))), `+`)
  M1_W_Y <- reduce(map(seq_along(b), \(k) compressed$w[k] * rbind(b[[k]]$S_Y, compressed$Z[k] * b[[k]]$S_Y)), `+`)
  B <- solve(M1_W_M2, M1_W_Y); colnames(B) <- outcomes
  list(M1_W_M2 = M1_W_M2, B = B)
}

L_matrix <- function(B, d) {                         # vec(E_a) = L u_a
  o <- ncol(B)
  cbind(diag(d * o), -kronecker(t(B[1:d, , drop = FALSE]), diag(d)), -kronecker(t(B[d + 1:d, , drop = FALSE]), diag(d)))
}
stack_instruments <- function(Z, d, o) kronecker(diag(o), rbind(diag(d), Z * diag(d)))   # vec(E_a) -> vec(inner sum) / w_a

fit_compressed <- function(compressed, d, outcomes) {                   # the covariance also needs sum u_a u_a^T
  o <- length(outcomes)
  fit <- fit_first_moments(compressed, d, outcomes)
  L <- L_matrix(fit$B, d)
  meat <- reduce(map(seq_len(nrow(compressed)), \(k) {
    P <- stack_instruments(compressed$Z[k], d, o)
    compressed$w[k]^2 * P %*% L %*% compressed$second_moments[[k]] %*% t(L) %*% t(P)
  }), `+`)
  M1_W_M2_inv <- kronecker(diag(o), solve(fit$M1_W_M2))
  c(fit, list(V = M1_W_M2_inv %*% meat %*% t(M1_W_M2_inv)))
}

complier_mean_compressed <- function(compressed, d, o, x_ref) {        # Wald ratio of T (x - x_ref) over T, from sum u_a
  arm <- \(z_value) reduce(map(which(compressed$Z == z_value), \(k) {
    b <- blocks(compressed$first_moments[[k]], d, o)
    compressed$w[k] * c(b$S_T[, 1], b$S[1, 1])            # weighted sums of T, T (x - x_ref), and of 1
  }), `+`)
  share <- \(v) v[1:d] / v[d + 1]
  diff <- share(arm(1)) - share(arm(0))
  set_names(diff[-1] / diff[1] + x_ref, names(x_ref))
}

One continuous covariate

We start with a single covariate, the pre-period hours $x_1$, and the outcome hours. This is the setting of CUPED. Here $x_{ap} = x_{1,ap} - 2$ is the centered pre-period hours, so the three blocks of an account are

\[S^Y_a = \begin{pmatrix} \sum_p y \\ \sum_p x y \end{pmatrix}, \qquad S_a = \begin{pmatrix} n_a & \sum_p x \\ \sum_p x & \sum_p x^2 \end{pmatrix}, \qquad S^T_a = \begin{pmatrix} \sum_p T & \sum_p T x \\ \sum_p T x & \sum_p T x^2 \end{pmatrix},\]

where $n_a$ is the number of profiles. With dplyr, each account becomes one row of sums.

Code

x_ref_1 <- c(x1 = 2)
cuped_summaries <- df |>
  mutate(x = x1 - x_ref_1[["x1"]]) |>
  group_by(account, Z, day, w) |>
  summarise(sum_y = sum(y), sum_xy = sum(x * y),                                  # S^Y_a
            n = n(), sum_x = sum(x), sum_xx = sum(x^2),                           # S_a
            sum_T = sum(T), sum_Tx = sum(T * x), sum_Txx = sum(T * x^2),          # S^T_a
            .groups = "drop")
cuped_summaries |>
  filter(account %in% example_accounts)

# A tibble: 2 × 12
  account     Z   day     w  sum_y sum_xy     n  sum_x sum_xx sum_T sum_Tx
    <int> <int> <int> <dbl>  <dbl>  <dbl> <int>  <dbl>  <dbl> <int>  <dbl>
1       1     0     2  1.12  0.431 -0.433     2 -1.25    1.30     0  0    
2      15     1     1 19.5  10.4   16.1       5  0.671  17.1      1 -0.330
# ℹ 1 more variable: sum_Txx <dbl>

The rows of our two accounts are the sums we computed by hand above, now for every account. Stacking the blocks gives $u_a$, whose columns repeat the off-diagonal entries of $S_a$ and $S^T_a$. Then each compressed row keeps the count, the first moments and the second moments of one stratum. Here are the first moments of the 8 compressed rows, written as the totals of each column:

Code

u_cuped <- with(cuped_summaries, cbind(sum_y, sum_xy, n, sum_x, sum_x, sum_xx, sum_T, sum_Tx, sum_Tx, sum_Txx))
compressed_cuped <- compress_summaries(cuped_summaries, u_cuped)
cuped_summaries |>
  group_by(Z, day, w) |>
  summarise(G = n(), across(sum_y:sum_Txx, sum), .groups = "drop")

# A tibble: 8 × 12
      Z   day     w     G  sum_y  sum_xy     n   sum_x sum_xx sum_T sum_Tx
  <int> <int> <dbl> <int>  <dbl>   <dbl> <int>   <dbl>  <dbl> <int>  <dbl>
1     0     1  1.05  7639 39224. 42908.  16646   645.  45107.     0    0  
2     0     2  1.12  5321 24686. 16001.  11713 -2423.  23961.     0    0  
3     0     3  1.33  2988 12836.  4665.   6654 -2654.  11597.     0    0  
4     0     4  2.00  1014  3886.   541.   2240 -1329.   3751.     0    0  
5     1     1 19.5    412  2449.  3410.    901    21.4  2749.   400  227. 
6     1     2  9.39   634  3387.  2509.   1429  -257.   3039.   610  204. 
7     1     3  4.05   981  4686.  1776.   2162  -863.   3711.   955  -65.5
8     1     4  2.00  1011  4376.    85.4  2244 -1363.   3427.   939 -283. 
# ℹ 1 more variable: sum_Txx <dbl>

These first moments are the familiar CUPED aggregates, $\sum y$, $\sum x$, $\sum x^2$ and $\sum xy$, kept separately for the treated profiles and for each stratum. The second moments are new. Each compressed row also holds the $10 \times 10$ matrix $\sum_a u_a u_a^\top$, whose entries are sums over accounts of products of account totals, such as $\sum_a \left(\sum_p y\right)\left(\sum_p x y\right)$. They carry the correlation between the profiles of an account, which the cluster robust covariance needs. The first stratum’s matrix starts like this:

Code

round(compressed_cuped$second_moments[[1]][1:4, 1:4])

        sum_y   sum_xy      n  sum_x
sum_y  333228   734454 105397  48149
sum_xy 734454 12236353 108343 445425
n      105397   108343  44826    879
sum_x   48149   445425    879  48253

From the compressed dataset we fit the model and compare it with iv_robust on all the profile rows, using only $x_1$ as a covariate.

Code

fit_cuped <- fit_compressed(compressed_cuped, d = 2, outcomes = "y")
x_C_cuped <- complier_mean_compressed(compressed_cuped, d = 2, o = 1, x_ref_1)
effects_cuped <- effects(fit_cuped, x_C_cuped, x_ref_1)
textbook_cuped <- textbook(df, x_C_cuped, clustered = TRUE, encouraged = TRUE, outcome = "y", covs = "x1")
tibble(method = c("compressed dataset", "iv_robust on profile rows"),
       LATE = c(effects_cuped$estimate, textbook_cuped[["estimate"]]),
       se = c(effects_cuped$se, textbook_cuped[["se"]])) |>
  kable(digits = 10)

method	LATE	se
compressed dataset	0.6695098	0.05721581
iv_robust on profile rows	0.6695098	0.05721581

One categorical covariate

Next, a single categorical covariate, the device. With tv as the baseline, the covariates are the indicators $\text{mobile}_{ap}$ and $\text{web}_{ap}$, so $(1, x_{ap}) = (1, \text{mobile}_{ap}, \text{web}_{ap})$. The indicators make the products simple: $\text{mobile}^2 = \text{mobile}$, and $\text{mobile} \cdot \text{web} = 0$ because a profile has one device. So $S_a$ and $S^T_a$ only hold counts of profiles by device, and $S^Y_a$ holds sums of hours by device:

\[S_a = \begin{pmatrix} n_a & n_{a,\text{mobile}} & n_{a,\text{web}} \\ n_{a,\text{mobile}} & n_{a,\text{mobile}} & 0 \\ n_{a,\text{web}} & 0 & n_{a,\text{web}} \end{pmatrix}.\]

Code

device_summaries <- df |>
  group_by(account, Z, day, w) |>
  summarise(sum_y = sum(y), sum_y_mobile = sum(mobile * y), sum_y_web = sum(web * y),        # S^Y_a
            n = n(), n_mobile = sum(mobile), n_web = sum(web),                                # S_a
            n_T = sum(T), n_T_mobile = sum(T * mobile), n_T_web = sum(T * web),               # S^T_a
            .groups = "drop")
device_summaries |>
  filter(account %in% example_accounts)

# A tibble: 2 × 13
  account     Z   day     w  sum_y sum_y_mobile sum_y_web     n n_mobile n_web
    <int> <int> <int> <dbl>  <dbl>        <dbl>     <dbl> <int>    <dbl> <dbl>
1       1     0     2  1.12  0.431        0.377    0.0541     2        1     1
2      15     1     1 19.5  10.4          2.38     3.02       5        2     1
# ℹ 3 more variables: n_T <int>, n_T_mobile <dbl>, n_T_web <dbl>

Code

x_ref_2 <- c(mobile = 0, web = 0)
u_device <- with(device_summaries, cbind(sum_y, sum_y_mobile, sum_y_web,
                                    n, n_mobile, n_web, n_mobile, n_mobile, 0, n_web, 0, n_web,
                                    n_T, n_T_mobile, n_T_web, n_T_mobile, n_T_mobile, 0, n_T_web, 0, n_T_web))
compressed_device <- compress_summaries(device_summaries, u_device)
fit_device <- fit_compressed(compressed_device, d = 3, outcomes = "y")
x_C_device <- complier_mean_compressed(compressed_device, d = 3, o = 1, x_ref_2)
effects_device <- effects(fit_device, x_C_device, x_ref_2)
textbook_device <- textbook(df, x_C_device, clustered = TRUE, encouraged = TRUE, outcome = "y", covs = c("mobile", "web"))
tibble(method = c("compressed dataset", "iv_robust on profile rows"),
       LATE = c(effects_device$estimate, textbook_device[["estimate"]]),
       se = c(effects_device$se, textbook_device[["se"]])) |>
  kable(digits = 10)

method	LATE	se
compressed dataset	0.6721917	0.09203501
iv_robust on profile rows	0.6721917	0.09203501

Categorical covariates compress even further when each account has one profile row. Then $u_a$ is fixed by the account’s device and treatment, up to its outcome $y$, so the sum of $u_a u_a^\top$ over a cell of identical profile rows only needs the cell’s count, $\sum y$ and $\sum y^2$. For the first profile of each account, the compressed dataset has one compressed row per cell:

Code

df |>
  filter(profile == 1) |>
  group_by(Z, day, device, T) |>
  summarise(n = n(), sum_y = sum(y), sum_y2 = sum(y^2), .groups = "drop") |>
  head(6)

# A tibble: 6 × 7
      Z   day device     T     n sum_y sum_y2
  <int> <int> <chr>  <int> <int> <dbl>  <dbl>
1     0     1 mobile     0  2624 6537. 27597.
2     0     1 tv         0  3832 9469. 42168.
3     0     1 web        0  1183 2817. 11237.
4     0     2 mobile     0  1865 4097. 15612.
5     0     2 tv         0  2633 5707. 21180.
6     0     2 web        0   823 1784.  6332.

Counts, sums and sums of squares per cell are enough for a regression on categorical covariates without clusters. With clusters, the account is the unit, and the compressed rows need the second moments of the account summaries, as in device_summaries.

All covariates

The full model has five covariates and five outcome columns, so $u_a$ has $d o + 2 d^2 =$ 102 entries. Writing one dplyr column per entry would work, but it is tedious, so we write the loop in C++ using Rcpp. compress_profiles visits the profile rows once, in account order. It adds each row’s products into a buffer that holds its account’s summary $u_a$. When the account’s rows end, it adds $u_a$ into its stratum’s first moments and $u_a u_a^\top$ into its stratum’s second moments, and clears the buffer, so the account summaries are never stored. The buffer keeps each distinct entry of $u_a$ once, since $S_a$ and $S^T_a$ are symmetric, and the outer product skips the entries of $u_a$ that are zero, such as all of $S^T_a$ for an account that was not prompted, or the indicators of devices the account does not use. This is the only step that visits every row. account_summaries builds the same $u_a$ but returns them, one row per account; delta vectors early need them for their second pass.

Code

#include <Rcpp.h>
#include <vector>
using namespace Rcpp;

// Profile rows to account summaries u_a = (vec S^Y, vec S, vec S^T) in one pass, with R = (1, x) and Y the outcome columns.
// [[Rcpp::export]]
NumericMatrix account_summaries(IntegerVector account, NumericMatrix R, NumericVector T, NumericMatrix Y, int n_accounts) {
  int n = R.nrow(), d = R.ncol(), o = Y.ncol();
  NumericMatrix u(n_accounts, d * o + 2 * d * d);
  for (int i = 0; i < n; i++) {
    int a = account[i] - 1;
    for (int l = 0; l < o; l++) for (int j = 0; j < d; j++) u(a, l * d + j) += R(i, j) * Y(i, l);     // S^Y
    for (int m = 0; m < d; m++) for (int j = 0; j < d; j++) {
      double rr = R(i, j) * R(i, m);
      u(a, d * o + m * d + j) += rr;                                                                // S
      u(a, d * o + d * d + m * d + j) += T[i] * rr;                                                 // S^T
    }
  }
  return u;
}

// Profile rows to the compressed dataset in one pass; the rows of each account must be together.
// The buffer holds the distinct entries of u_a: S^Y, then the upper triangles of S and S^T.
// [[Rcpp::export]]
List compress_profiles(IntegerVector account, IntegerVector stratum, NumericMatrix R, NumericVector T,
                       NumericMatrix Y, int n_strata) {
  int n = R.nrow(), d = R.ncol(), o = Y.ncol();
  int nt = d * (d + 1) / 2, q = d * o + 2 * nt;
  std::vector<double> u(q, 0.0);
  std::vector<int> nz;
  IntegerVector G(n_strata);
  NumericMatrix first(n_strata, q);
  std::vector<NumericMatrix> second;
  for (int h = 0; h < n_strata; h++) second.push_back(NumericMatrix(q, q));
  auto add_account = [&](int h) {                                         // add u_a and u_a u_a^T into stratum h
    nz.clear();
    for (int k = 0; k < q; k++) if (u[k] != 0) nz.push_back(k);
    G[h] += 1;
    NumericMatrix& M = second[h];
    for (size_t a = 0; a < nz.size(); a++) {
      first(h, nz[a]) += u[nz[a]];
      for (size_t b = a; b < nz.size(); b++) M(nz[b], nz[a]) += u[nz[a]] * u[nz[b]];   // lower triangle
    }
    for (int k : nz) u[k] = 0;
  };
  for (int i = 0; i < n; i++) {
    if (i > 0 && account[i] != account[i - 1]) add_account(stratum[i - 1] - 1);
    for (int l = 0; l < o; l++) for (int j = 0; j < d; j++) u[l * d + j] += R(i, j) * Y(i, l);   // S^Y
    for (int m = 0; m < d; m++) for (int j = 0; j <= m; j++) {
      double rr = R(i, j) * R(i, m);
      int t = m * (m + 1) / 2 + j;
      u[d * o + t] += rr;                                                                  // S
      u[d * o + nt + t] += T[i] * rr;                                                      // S^T
    }
  }
  add_account(stratum[n - 1] - 1);
  List second_moments(n_strata);
  for (int h = 0; h < n_strata; h++) {
    NumericMatrix& M = second[h];
    for (int a = 0; a < q; a++) for (int b = a + 1; b < q; b++) M(a, b) = M(b, a);         // fill the upper triangle
    second_moments[h] = M;
  }
  return List::create(_["G"] = G, _["first"] = first, _["second"] = second_moments);
}

Code

distinct_index <- function(d, o) {                     # where each entry of u_a sits among its distinct entries
  tri <- outer(0:(d - 1), 0:(d - 1), \(j, m) pmax(j, m) * (pmax(j, m) + 1) / 2 + pmin(j, m))
  nt <- d * (d + 1) / 2
  c(seq_len(d * o), d * o + 1 + as.vector(tri), d * o + nt + 1 + as.vector(tri))
}

compress_rows <- function(df, R, Y = as.matrix(df[y_cols])) {        # profile rows to the compressed dataset
  stopifnot(!is.unsorted(df$account))                                 # each account's rows are together
  accounts <- df |>
    distinct(account, Z, day, w) |>
    arrange(account)
  strata <- group_by(accounts, Z, day, w)
  h <- group_indices(strata)
  k <- compress_profiles(df$account, h[match(df$account, accounts$account)], R, df$T, Y, max(h))
  i <- distinct_index(ncol(R), ncol(Y))
  group_keys(strata) |>
    mutate(G = k$G, first_moments = map(seq_len(n()), \(s) k$first[s, i]), second_moments = map(k$second, \(m) m[i, i]))
}

compressed <- compress_rows(df, centered(df))
compressed |>
  mutate(length_u = map_int(first_moments, length), dim_second_moments = map_chr(second_moments, \(m) paste(dim(m), collapse = " x "))) |>
  select(Z, day, w, G, length_u, dim_second_moments)

# A tibble: 8 × 6
      Z   day     w     G length_u dim_second_moments
  <int> <int> <dbl> <int>    <int> <chr>             
1     0     1  1.05  7639      102 102 x 102         
2     0     2  1.12  5321      102 102 x 102         
3     0     3  1.33  2988      102 102 x 102         
4     0     4  2.00  1014      102 102 x 102         
5     1     1 19.5    412      102 102 x 102         
6     1     2  9.39   634      102 102 x 102         
7     1     3  4.05   981      102 102 x 102         
8     1     4  2.00  1011      102 102 x 102

The compressed dataset has 8 compressed rows, one per stratum, whatever the number of profile rows. Fitting from it reproduces the fit on profile rows to machine precision. Table 4 compares its effects with iv_robust on all 43,989 profile rows, centering the covariates at the compliers’ means so that the coefficient on $T$ is the effect, and computing the ratio from the four outcome columns.

Code

fit_c <- fit_compressed(compressed, d, y_cols)
x_C_compressed <- complier_mean_compressed(compressed, d, o, x_ref)
internal <- c(coef = max(abs(fit_c$B - fit$B)), cov = max(abs(fit_c$V - fit$V)) / max(abs(fit$V)))
effects_c <- effects(fit_c, x_C_compressed, x_ref)
er <- map(set_names(c("y", "N", "D", "TN", "TD")), \(col) textbook(df, x_C_compressed, TRUE, TRUE, col))
compare_tab <- tibble(
  quantity = c("LATE on hours", "standard error of the LATE", "ratio effect on rebuffers per play"),
  compressed = c(effects_c$estimate[1], effects_c$se[1], effects_c$estimate[2]),
  estimatr = c(er$y[["estimate"]], er$y[["se"]],
               er$TN[["estimate"]] / er$TD[["estimate"]] - (er$TN[["estimate"]] - er$N[["estimate"]]) / (er$TD[["estimate"]] - er$D[["estimate"]]))) |>
  mutate(difference = compressed - estimatr)

Code

compare_tab |>
  mutate(across(c(compressed, estimatr), \(v) sprintf("%.12f", v)), difference = sprintf("%.1e", difference)) |>
  kable(col.names = c("Quantity", "Compressed dataset", "estimatr on profile rows", "Difference"))

Table 4: Effects from the compressed dataset and from estimatr’s iv_robust on every profile row, with the same weights and clusters.

Quantity	Compressed dataset	estimatr on profile rows	Difference
LATE on hours	0.663910193939	0.663910193939	4.0e-15
standard error of the LATE	0.057160548380	0.057160548380	6.2e-16
ratio effect on rebuffers per play	-0.025478235508	-0.025478235508	-2.1e-17

The differences are at the level of floating point rounding. estimatr fits one outcome at a time, so it cannot give the ratio’s standard error, which needs the covariance between outcomes. Against our own fit on the profile rows, the full compressed covariance, including the covariance between outcomes, agrees to a relative difference of 2e-14.

Delta Vectors Optimization

Delta vectors are a convenient construct that compute treatment effects by combining the right coefficients and covariances for you. This works for an average effect, as well as conditional average effects. We saw its flexibility in the earlier section Section 2.4.3. Earlier, we used it after estimating the coefficients and their covariances.

Compression changes our way of thinking. We need to be conscious of data storage. We can compress the data to estimate coefficients and covariances, then combine those outputs with delta vectors to get effects. A step stands out though. To compute the covariances from compressed data, the compressed rows had to hold the second moments $\sum_a u_a u_a^\top$. These can be expensive. If we know in advance which effects we want, say the LATE, we can use the delta vectors earlier and never build $\sum_a u_a u_a^\top$, with exactly the same standard errors. The price is flexibility and a second pass over the data. We describe the two strategies, one using it flexibly after getting $\hat B$ and covariances, and one earlier, and compare their costs in Benchmark.

Delta vectors at the end (most flexible)

The coefficient matrix $\hat B$ has $2d = $ 12 rows and $o = $ 5 columns, so the covariance of its 60 entries is a 60 by 60 matrix. To get it, the compressed rows must hold $\sum_a u_a u_a^\top$, which has 10,404 entries per stratum. An effect only needs one number from it. For the LATE, the gradient with respect to $\text{vec}(\hat B)$ is the delta vector $\delta$ in the block of the hours column and zeros elsewhere:

Code

effect_list <- effect_functions(x_C_compressed, x_ref, y_cols)
theta <- as.vector(fit_c$B)
grad_late <- gradient(\(th) effect_list$tau_LATE(matrix(th, nrow(fit_c$B), dimnames = list(NULL, y_cols))), theta)
matrix(round(grad_late, 3), nrow = 2 * d, dimnames = list(c(colnames(centered(example)), paste0("T*", colnames(centered(example)))), y_cols))

                y N D TN TD
1           0.000 0 0  0  0
x1_c        0.000 0 0  0  0
x2_c        0.000 0 0  0  0
kids_c      0.000 0 0  0  0
mobile_c    0.000 0 0  0  0
web_c       0.000 0 0  0  0
T*1         1.000 0 0  0  0
T*x1_c      0.286 0 0  0  0
T*x2_c     -0.239 0 0  0  0
T*kids_c    0.108 0 0  0  0
T*mobile_c  0.347 0 0  0  0
T*web_c     0.136 0 0  0  0

Code

c(dim_V = dim(fit_c$V), se_LATE = sqrt(drop(t(grad_late) %*% fit_c$V %*% grad_late)))

     dim_V1      dim_V2     se_LATE 
60.00000000 60.00000000  0.05716055

Delta vectors earlier (most efficient)

The variance of an effect with gradient $g$ is $g^\top \widehat{\text{Cov}}(\text{vec}(\hat B)) \, g$. Writing out the sandwich, it is a sum of squares over accounts,

\[g^\top \widehat{\text{Cov}}(\text{vec}(\hat B)) \, g = \sum_a \psi_a^2, \qquad \psi_a = g^\top \left(I_o \otimes (M_1^\top W M_2)^{-1}\right) \, \text{vec}\left(\textstyle\sum_p w_a \, m_{1,ap} \, \hat\varepsilon_{ap}^\top\right),\]

where $\psi_a$ is the influence of account $a$ on the effect. The inner sums are $w_a P_{Z_a} L(\hat B) \, u_a$, where the matrix $P_{Z_a} = I_o \otimes \begin{pmatrix} I_d \\ Z_a I_d \end{pmatrix}$ maps $\text{vec}(E_a)$ to the stacked inner sums of $E_a$ and $Z_a E_a$. Everything in front of $u_a$ is the same for all accounts in a stratum, so

\[\psi_a = c_h^\top u_a, \qquad c_h = w_h \, L(\hat B)^\top P_{Z_h}^\top \left(I_o \otimes (M_1^\top W M_2)^{-\top}\right) g.\]

This changes the order of the work. We fit $\hat B$ from the first moments $\sum u_a$ alone. Once $\hat B$ and the effects we want are known, we compute $c_h$ for each stratum, and a second pass over the account summaries gives each account’s influence with one dot product. We never form $\sum u_a u_a^\top$.

Code

compressed_first <- compressed |>
  select(Z, day, w, first_moments)                     # pass 1 only keeps sum u_a
fit_1 <- fit_first_moments(compressed_first, d, y_cols)
x_C_1 <- complier_mean_compressed(compressed_first, d, o, x_ref)
effect_list_1 <- effect_functions(x_C_1, x_ref, y_cols)
grads <- map(effect_list_1, \(effect) gradient(\(th) effect(matrix(th, nrow(fit_1$B), dimnames = list(NULL, y_cols))), as.vector(fit_1$B)))

projection <- function(fit, compressed_first, grad, d, o) {               # c_h for every stratum, one row each
  L <- L_matrix(fit$B, d)
  v <- kronecker(diag(o), t(solve(fit$M1_W_M2))) %*% grad
  do.call(rbind, map(seq_len(nrow(compressed_first)), \(k)
    as.vector(compressed_first$w[k] * t(L) %*% t(stack_instruments(compressed_first$Z[k], d, o)) %*% v)))
}
C_late <- projection(fit_1, compressed_first, grads$tau_LATE, d, o)
dim(C_late)

[1]   8 102

Code

accounts <- df |>
  distinct(account, Z, day, w) |>
  arrange(account)
u <- account_summaries(df$account, centered(df), df$T, as.matrix(df[y_cols]), nrow(accounts))   # pass 2 rebuilds u_a
h <- group_indices(group_by(accounts, Z, day, w))
psi_late <- rowSums(u * C_late[h, ])                                 # pass 2: one dot product per account
tibble(account = accounts$account, psi = psi_late) |>
  filter(account %in% example_accounts)

# A tibble: 2 × 2
  account       psi
    <int>     <dbl>
1       1  0.000174
2      15 -0.00284

The prompted account moves the LATE more than the unprompted one, because the prompted arm of its cohort is small and its weight is large. The standard error is the square root of the sum of squared influences, and it agrees with the full covariance:

Code

se_early <- map_dbl(grads, \(g) sqrt(sum(rowSums(u * projection(fit_1, compressed_first, g, d, o)[h, ])^2)))
tibble(effect = names(se_early), se_end = effects_c$se, se_early = se_early, difference = se_early - effects_c$se)

# A tibble: 2 × 4
  effect    se_end se_early difference
  <chr>      <dbl>    <dbl>      <dbl>
1 tau_LATE 0.0572   0.0572    6.25e-17
2 tau_R    0.00410  0.00410  -7.81e-18

Each method has its place. Delta vectors at the end need the second moments, but once we have them, we can ask for any effect later: a new segment, a new point on the treatment line, a slope. Delta vectors early must know the effects before the second pass, but they are cheaper. Forming $\sum u_a u_a^\top$ costs $(\dim u)^2$ operations per account, while the projections cost $\dim u$ per account and effect.

Benchmark

We now have three ways from the profile rows to the effects: the textbook fit on the profile rows, and the compressed dataset with delta vectors at the end or early. This section compares how long each takes and how many numbers each keeps.

Compute

To see what each method costs, we run the simulation of the running example with a million accounts instead of 20,000, and time the textbook fit, one iv_robust per outcome column, against the two compressed methods. Delta vectors at the end compress the profile rows in one pass. Delta vectors early build the account summaries, add them into first moments, and take the second pass over them.

Code

set.seed(1)
big <- simulate_accounts(1e6) |>
  prepare()
big_c <- big |>
  mutate(across(all_of(x_cols), \(v) v - x_ref[cur_column()]))
t_textbook <- system.time(walk(y_cols, \(col) iv_robust(
  as.formula(paste(col, "~ (x1 + x2 + kids + mobile + web) * T | (x1 + x2 + kids + mobile + web) * Z")),
  data = big_c, weights = w, clusters = account, se_type = "CR0")))[["elapsed"]]
t_end <- system.time({
  big_compressed <- compress_rows(big, centered(big))
  big_fit <- fit_compressed(big_compressed, d, y_cols)
  se_end_big <- effects(big_fit, complier_mean_compressed(big_compressed, d, o, x_ref), x_ref)$se
})[["elapsed"]]
t_early <- system.time({
  big_accounts <- big |>
    distinct(account, Z, day, w) |>
    arrange(account)
  big_u <- account_summaries(big$account, centered(big), big$T, as.matrix(big[y_cols]), nrow(big_accounts))
  h_big <- group_indices(group_by(big_accounts, Z, day, w))
  big_first <- group_keys(group_by(big_accounts, Z, day, w)) |>
    mutate(first_moments = map(seq_len(n()), \(k) colSums(big_u[h_big == k, , drop = FALSE])))
  big_fit_1 <- fit_first_moments(big_first, d, y_cols)
  effect_list_big <- effect_functions(complier_mean_compressed(big_first, d, o, x_ref), x_ref, y_cols)
  se_early_big <- map_dbl(effect_list_big, \(effect) {
    grad <- gradient(\(th) effect(matrix(th, nrow(big_fit_1$B), dimnames = list(NULL, y_cols))), as.vector(big_fit_1$B))
    sqrt(sum(rowSums(big_u * projection(big_fit_1, big_first, grad, d, o)[h_big, ])^2))
  })
})[["elapsed"]]
n_big <- nrow(big)
n_big_accounts <- nrow(big_accounts)
rm(big, big_c, big_u, big_accounts, big_first, h_big)
invisible(gc())

Code

tibble(step = c("textbook: iv_robust on profile rows, once per outcome column",
                "delta vectors at the end: one pass to the compressed dataset, covariance, effects",
                "delta vectors early: account summaries, first moments, projections"),
       seconds = c(t_textbook, t_end, t_early)) |>
  kable(digits = 3, col.names = c("Step", "Seconds"))

Table 5: Seconds to estimate the LATE and the ratio, with their standard errors, on a million accounts, starting from the profile rows.

Step	Seconds
textbook: iv_robust on profile rows, once per outcome column	11.527
delta vectors at the end: one pass to the compressed dataset, covariance, effects	0.742
delta vectors early: account summaries, first moments, projections	1.160

Table 5 shows the times on the machine that rendered this post. The 2,189,477 profile rows reduce to 1,000,000 account summaries in one pass over the profile rows, and to 8 compressed rows. The two compressed methods agree on the standard errors to within 6e-16.

Storage

The two compressed methods also keep different numbers. Delta vectors at the end keep, for every stratum, the count of accounts, the first moments $\sum_a u_a$ and the second moments $\sum_a u_a u_a^\top$, and any effect can be computed from them later. Delta vectors early keep only the first moments after the first pass. The second pass visits the accounts again, one at a time: it rebuilds $u_a$, takes its dot product with $c_h$, and adds $\psi_a^2$ to a running total, one total per effect. Nothing is kept per account, so after the second pass the stored numbers are the first moments and one sum of squares per effect. Our code for delta vectors early keeps the account summaries in memory for convenience; at scale, the second pass reads the profile rows again instead.

Both methods grow with the number of covariates. $u_a$ has $d o + d(d+1)$ distinct entries, with $d = k + 1$, because $S_a$ and $S^T_a$ hold the product of every pair of covariates. That is quadratic in $k$, and so are the first moments that delta vectors early keep. The second moments keep the products of every pair of entries of $u_a$, so delta vectors at the end grow like $k^4$.

Table 6 compares the numbers in the profile rows of our experiment with the numbers each method keeps, for the three models of the running example, and for the full model with more continuous covariates on the same 43,989 profile rows. A profile row holds the account, $Z$, the day, $T$, the $k$ covariates and the raw outcomes; the columns $TN$ and $TD$ are computed from them. The methods’ counts take each distinct number once: a symmetric matrix by its upper triangle, and the repeated or always zero entries of $u_a$ once or not at all. The two models with one covariate estimate one effect, the LATE on hours, and the others estimate two, the LATE and the ratio. The reduction is the number of numbers in the profile rows divided by the number the method keeps.

Code

n_strata <- nrow(compressed)
storage <- function(case, k, raw_outcomes, distinct_u, n_effects) {
  tibble(case = case, k = k,
         rows = nrow(df) * (4 + k + raw_outcomes),                                        # account, Z, day, T, covariates, raw outcomes
         end = n_strata * (1 + distinct_u + distinct_u * (distinct_u + 1) / 2),           # count, first and second moments
         early = n_strata * distinct_u + n_effects,                                       # first moments, one sum of squares per effect
         distinct_u = distinct_u)
}
full_u <- \(k) (k + 1) * o + (k + 1) * (k + 2)
storage_tab <- bind_rows(
  storage("one continuous covariate, hours", 1, 1, ncol(cuped_summaries) - 4, 1),
  storage("one categorical covariate, hours", 2, 1, ncol(device_summaries) - 4, 1),
  storage("all covariates, hours and rebuffers per play", length(x_cols), 3, full_u(length(x_cols)), length(grads)),
  map_dfr(c(10, 20, 50), \(k) storage(paste(k, "continuous covariates, hours and rebuffers per play"), k, 3, full_u(k), 2))) |>
  mutate(reduction_end = rows / end, reduction_early = rows / early)
rows_big <- n_big * (4 + length(x_cols) + 3)                                  # the full model's profile rows in the compute benchmark

Code

times <- \(v) paste0(ifelse(v >= 10, format(round(v), big.mark = ",", trim = TRUE), formatC(signif(v, 2), format = "fg", digits = 2)), "x")
storage_tab |>
  transmute(case, k, rows = format(rows, big.mark = ","),
            end = format(end, big.mark = ","), reduction_end = times(reduction_end),
            early = format(early, big.mark = ","), reduction_early = times(reduction_early)) |>
  kable(col.names = c("Case", "Covariates $k$", "Numbers in the profile rows", "Kept, delta vectors at the end", "Reduction",
                      "Kept, delta vectors early", "Reduction"))

Table 6: Numbers in the profile rows of our experiment and numbers kept by each method, counting each distinct number once. The reduction is the first divided by the second. The first three rows are the models of the running example.

Case	Covariates $k$	Numbers in the profile rows	Kept, delta vectors at the end	Reduction	Kept, delta vectors early	Reduction
one continuous covariate, hours	1	263,934	360	733x	65	4,061x
one categorical covariate, hours	2	307,923	440	700x	73	4,218x
all covariates, hours and rebuffers per play	5	527,868	21,608	24x	578	913x
10 continuous covariates, hours and rebuffers per play	10	747,813	142,128	5.3x	1,498	499x
20 continuous covariates, hours and rebuffers per play	20	1,187,703	1,292,768	0.92x	4,538	262x
50 continuous covariates, hours and rebuffers per play	50	2,507,373	33,837,488	0.074x	23,258	108x

For the full model, the 43,989 profile rows hold 527,868 numbers. Delta vectors at the end keep 21,608, 24 times fewer, and delta vectors early keep 578, 913 times fewer. The account summaries, 72 numbers for each of the 20,000 accounts, would hold more numbers than the profile rows; this is why neither method stores them, and each one is added into its compressed row, or into its running totals, as soon as it is built.

The methods keep the same numbers however many profile rows there are, so the reduction grows in proportion to the size of the experiment. The 2,189,477 profile rows of the compute benchmark hold 26,273,724 numbers for the full model, a reduction of 1,216 times at the end and 45,456 times early. More covariates shrink the reduction. With 20 covariates, delta vectors at the end already keep more numbers than the profile rows of our experiment, and with 50 they keep 33,837,488, while delta vectors early keep 23,258. The price of the smaller storage is flexibility and a second read: the effects must be chosen before the second pass, and a new effect, such as a new segment or a slope, needs another pass over the profile rows. A categorical covariate with $L$ levels counts as $L - 1$ indicators here, which overstates its size: products of indicators of the same covariate are zero or repeated, as the device case showed.

When There Are Many Covariates

Table 6 showed that delta vectors early keep storage quadratic in $k$, but they fix the effects in advance. Delta vectors at the end can answer any effect later, but their storage grows like $k^4$, and experimentation platforms often adjust for dozens of pre-experiment signals. To keep the flexibility of delta vectors at the end, we reduce the covariates. In this section we add covariates to the running example and summarize them with prognostic scores.

More covariates for the running example

We add $m = 20$ covariates to every profile row. Each is a noisy measurement of the profile’s engagement, and half of them also carry the account’s network quality, like the activity signals a service records before an experiment. Each signal on its own is weak, but together they predict hours well. The signals are drawn after everything else in the simulation, so the earlier profile rows, and the true effects, do not change.

Code

add_signals <- function(profiles, m = 20) {
  signals <- map(seq_len(m), \(j) 0.4 * profiles$engagement_p + 0.2 * profiles$quality * (j %% 2) + rnorm(nrow(profiles)))
  bind_cols(profiles, as_tibble(set_names(signals, paste0("g", seq_len(m)))))
}
set.seed(31)
df_many <- profiles |>
  add_signals() |>
  prepare()
signal_cols <- paste0("g", 1:20)
sizes <- function(k, categorical = FALSE) {
  first <- if (categorical) (k + 1) * o + 2 * (k + 1) else full_u(k)
  totals <- o + 2 * (k + 1)
  c(exact = n_strata * (1 + first + first * (first + 1) / 2), second_order = n_strata * (1 + first + totals * (totals + 1) / 2))
}
X_many <- cbind(centered(df_many)[, -1], as.matrix(df_many[signal_cols]))    # 5 base covariates and 20 signals, centered
dim(X_many)

[1] 43989    25

With $k = 25$ covariates, the compressed dataset for delta vectors at the end would hold 2,778,888 numbers, counting each distinct number once, against 1,407,648 in the profile rows: the account, $Z$, the day, $T$, 25 covariates and 3 raw outcomes per row. The compressed dataset would be larger than the data.

Summarizing the covariates with prognostic scores

A covariate only reduces variance through what it predicts about the outcome, so we can replace the 25 covariates with a prediction of each outcome, a prognostic score, from a model trained on data from before the experiment. A fixed function of pre-experiment covariates is itself a pre-experiment covariate, so the scores are as valid as any other covariate, and a better model only makes them more precise. DoorDash calls this CUPAC, control using predictions as covariates.

We simulate a pre-experiment sample of the same population with nobody prompted, and fit one linear model per raw outcome, hours, rebuffers and plays. Any model would do, including a gradient boosted tree.

Code

set.seed(99)
history <- simulate_accounts(50000, force_Z = 0) |>
  add_signals() |>
  prepare()
X_history <- cbind(centered(history)[, -1], as.matrix(history[signal_cols]))
score_models <- lm(as.matrix(history[c("y", "N", "D")]) ~ X_history) |>    # one fit for the three outcomes
  coef() |>
  as.data.frame() |>
  as.list()
score <- function(X) {
  map_dfc(score_models, \(b) drop(cbind(1, X) %*% b) - mean(drop(cbind(1, X_history) %*% b))) |>
    as.matrix()
}
X_scores <- score(X_many)
head(X_scores)

               y            N          D
[1,] -0.07084308 -0.039162732 -0.5381469
[2,] -0.73417387 -0.098503839 -0.8858300
[3,] -1.69499355 -0.132632913 -2.6397010
[4,]  2.60128168  0.154814566  2.5229958
[5,] -0.55769760  0.007430742 -0.4327662
[6,] -0.56787040 -0.045352782 -1.2735420

The model now has $k = 3$ covariates, one score per raw outcome, whatever the number of signals.

Table 7 compares the three scores with all 25 covariates on this experiment, both with the compressed dataset for delta vectors at the end.

Code

fit_exact <- function(df, X) {
  compressed <- compress_rows(df, cbind(1, X))
  zero <- set_names(rep(0, ncol(X)), paste0("x", seq_len(ncol(X))))
  fit <- fit_compressed(compressed, ncol(X) + 1, y_cols)
  effects(fit, complier_mean_compressed(compressed, ncol(X) + 1, o, zero), zero)
}

scores_tab <- tibble(method = c("all 25 covariates", "3 prognostic scores"), k = c(25, 3),
                     numbers = c(sizes(25)[["exact"]], sizes(3)[["exact"]]),
                     fit = list(fit_exact(df_many, X_many), fit_exact(df_many, X_scores)))
exact_fits <- set_names(scores_tab$fit, scores_tab$method)
scores_tab <- scores_tab |>
  mutate(LATE = map_dbl(fit, \(f) f$estimate[1]), se = map_dbl(fit, \(f) f$se[1])) |>
  select(-fit)

Code

scores_tab |>
  mutate(numbers = format(numbers, big.mark = ",")) |>
  kable(digits = 4, col.names = c("Method", "Covariates $k$", "Numbers stored", "LATE", "SE"))

Table 7: All 25 covariates against 3 prognostic scores, each fit from its compressed dataset for delta vectors at the end. Numbers stored count each distinct number once.

Method	Covariates $k$	Numbers stored	LATE	SE
all 25 covariates	25	2,778,888	0.6652	0.0526
3 prognostic scores	3	6,888	0.6656	0.0528

The scores are as precise as all 25 covariates, with a standard error of 0.0528 against 0.0526, and their compressed dataset holds 6,888 numbers instead of 2,778,888. Its size does not depend on the number of covariates the scores summarize, only on the number of outcomes. Two practical notes:

The scores are predictions, but they do not add a term to the standard error. Once trained on data from before the experiment, the score is a fixed function of observed covariates, and the effect is consistent whatever that function is. The appendix explains why and checks it by retraining the scores in every simulated experiment.
Effects along the scores are hard to interpret. When we want the effect along a particular covariate, such as the effect for kids profiles, we keep that covariate in the model next to the scores.

The appendix compares the scores with three other strategies, a second order approximation of the compressed dataset, binning and clustering, and confirms over 400 simulated experiments that the scores are as precise as all 25 covariates.

End to End

Here is the whole analysis of our experiment in a few lines, starting from the profile rows. Both versions use the same weights and clusters.

With estimatr, we center the covariates at the compliers’ means, so that the coefficient on $T$ is the effect, and fit one two stage least squares model per outcome column. The ratio’s estimate combines four fits, but estimatr cannot give its standard error.

Code

x_C_rows <- complier_mean(df)
df_c <- df |>
  mutate(across(all_of(x_cols), \(v) v - x_C_rows[cur_column()]))
fits <- map(set_names(y_cols), \(col) iv_robust(
  as.formula(paste(col, "~ (x1 + x2 + kids + mobile + web) * T | (x1 + x2 + kids + mobile + web) * Z")),
  data = df_c, weights = w, clusters = account, se_type = "CR0"))
l <- map(fits, \(f) f$coefficients[["T"]])
tibble(effect = c("tau_LATE", "tau_R"),
       estimate = c(l$y, l$TN / l$TD - (l$TN - l$N) / (l$TD - l$D)),
       se = c(fits$y$std.error[["T"]], NA))

# A tibble: 2 × 3
  effect   estimate      se
  <chr>       <dbl>   <dbl>
1 tau_LATE   0.664   0.0572
2 tau_R     -0.0255 NA

With compression, we build the account summaries in one pass, compress them into one row per stratum, fit the model from the compressed dataset, and read off both effects with their standard errors.

Code

compressed_e2e <- compress_rows(df, centered(df))
fit_e2e <- fit_compressed(compressed_e2e, d, y_cols)
effects(fit_e2e, complier_mean_compressed(compressed_e2e, d, o, x_ref), x_ref)

# A tibble: 2 × 3
  effect   estimate      se
  <chr>       <dbl>   <dbl>
1 tau_LATE   0.664  0.0572 
2 tau_R     -0.0255 0.00410

Summary

One model covers every design in this post: weighted two stage least squares, with the treatment interacted with each covariate, cluster robust standard errors, and one outcome column per metric, plus the columns $TN$ and $TD$ for a ratio metric. Turning off the encouragement, the clusters, the ramp or the ratio gives back the textbook analysis of each simpler design, from Lin’s regression to the Wald estimator. Two details matter for getting the right effect: the delta vector reads the treatment line at the compliers’ covariates, and a ratio metric needs the compliers’ own levels.

The model needs only one compressed row per assignment and cohort: the count of accounts, the sum of the account summaries $u_a$, and the sum of their outer products. With one continuous covariate the compressed rows hold the CUPED aggregates and their products, and with categorical covariates they hold counts and sums by cell. When we know the effects we want, we can skip the outer products: the first moments give the fit, and one dot product per account gives each effect’s exact standard error.

The two ways to get standard errors trade storage for flexibility. Delta vectors at the end keep the outer products and can answer any effect later, but their storage grows like $k^4$ with the number of covariates; delta vectors early grow like $k^2$, but fix the effects in advance. With many covariates, prognostic scores trained before the experiment keep the compressed dataset independent of $k$, and are as precise as the covariates they summarize.

Appendix

Why the treatment line is fit to the compliers

As in the main text, $x$ denotes the centered covariates. Consider one cohort, where by (A2) every account is prompted with the same probability $e$, whatever its $x$. The two stage least squares estimate solves the sample version of $E[m_1 \varepsilon] = 0$, where $\varepsilon = y - m_2^\top \beta$. The instruments come in two halves, $(1, x)$ and $Z (1, x)$, so this is two sets of conditions, $E\left[\begin{pmatrix} 1 \\ x \end{pmatrix} \varepsilon\right] = 0$ and $E\left[Z \begin{pmatrix} 1 \\ x \end{pmatrix} \varepsilon\right] = 0$. Subtracting $e$ times the first set from the second gives

\[E\left[(Z - e) \begin{pmatrix} 1 \\ x \end{pmatrix} \varepsilon\right] = 0.\]

By (A2) the prompt is independent of $x$, so $E[Z - e \mid x] = 0$, and the baseline line, which does not involve $T$, drops out:

\[E\left[(Z - e) \begin{pmatrix} 1 \\ x \end{pmatrix} \left(y - T \left(\beta_{10} + x^\top \beta_{1x}\right)\right)\right] = 0.\]

Conditional on $x$, $E[(Z - e) \, y \mid x] = e(1-e) \left(E[y \mid Z = 1, x] - E[y \mid Z = 0, x]\right)$, and by (A3) the difference of the two conditional means is $\pi_C(x) \, \tau_C(x)$, the share of compliers at $x$ times their average effect, as in the post on averages. In the same way, $E[(Z - e) \, T \mid x] = e(1-e) \, \pi_C(x)$. Substituting both and dividing by $e(1-e)$,

\[E_x\left[\pi_C(x) \left(\tau_C(x) - \beta_{10} - x^\top \beta_{1x}\right) \begin{pmatrix} 1 \\ x \end{pmatrix}\right] = 0.\]

These are the normal equations of a weighted least squares fit of the treatment line $\beta_{10} + x^\top \beta_{1x}$ to the complier effects $\tau_C(x)$, with weight $\pi_C(x) f(x)$, where $f(x)$ is the probability density of $x$. The first entry of $(1, x)$ is 1, so the first equation says that the weighted average of the line equals the weighted average of $\tau_C(x)$. Dividing the weight by its total, the overall share of compliers $\pi_C = E[\pi_C(x)]$, gives the probability density of $x$ among compliers, $f(x \mid C)$, which turns the weighted average of $\tau_C(x)$ into the LATE and the weighted average of the line into its value at the compliers’ mean, $\bar x_C - x_{ref}$ in centered units.

For the outcome $TN$, the same steps apply with $E[(Z - e) \, TN \mid x] = e(1-e) \, \pi_C(x) \, E[N(1) \mid C, x]$, because the prompt changes $TN$ only for compliers, from 0 to their rebuffers with the setting on. So the treatment line of $TN$ is fit to $E[N(1) \mid C, x]$, and at $\bar x_C$ it equals $E[N(1) \mid C]$.

Proofs of the reductions

Write $\hat\beta = (M_1^\top W M_2)^{-1} M_1^\top W y$, and write $s_a = \sum_p w_a m_{1,ap} \hat\varepsilon_{ap}$ for the inner sums of account $a$, so that $\widehat{\text{Cov}}(\hat\beta) = (M_1^\top W M_2)^{-1} \left(\sum_a s_a s_a^\top\right) (M_1^\top W M_2)^{-\top}$.

No encouragement. If $T = Z$, then $m_{2,ap} = (1, x_{ap}, Z_a, Z_a x_{ap}) = m_{1,ap}$, so $M_1 = M_2$ and $\hat\beta = (M_2^\top W M_2)^{-1} M_2^\top W y$, which is weighted least squares. The sandwich becomes $(M_2^\top W M_2)^{-1} \left(\sum_a s_a s_a^\top\right)(M_2^\top W M_2)^{-1}$ with $s_a = \sum_p w_a m_{2,ap} \hat\varepsilon_{ap}$, the robust covariance of weighted least squares. Every profile is a complier, so $\bar x_C$ is the mean of $x$, and $\delta^\top \beta$ is the difference between the two arms’ lines at the mean, which is the coefficient on $T$ in Lin’s regression with covariates centered at their mean.

No clusters. If each account has one profile, then $s_a = w_a m_{1,a} \hat\varepsilon_a$ and $\sum_a s_a s_a^\top = \sum_a w_a^2 \hat\varepsilon_a^2 m_{1,a} m_{1,a}^\top$, the meat of the heteroskedasticity robust covariance (HC0). The usual small sample corrections also agree: with $G$ clusters, $n$ rows and $\dim(\beta)$ coefficients, the cluster correction $\frac{G}{G-1} \frac{n-1}{n-\dim(\beta)}$ equals the HC1 correction $\frac{n}{n-\dim(\beta)}$ when $G = n$.

No ramp. If every weight equals a constant $w$, then $\hat\beta = (w M_1^\top M_2)^{-1} w M_1^\top y = (M_1^\top M_2)^{-1} M_1^\top y$. The inner sums are $w$ times their unweighted versions $s'_a$, so the covariance is $(w M_1^\top M_2)^{-1} \left(w^2 \sum_a s'_a s'^\top_a\right) (w M_1^\top M_2)^{-\top}$, and the $w$ cancels.

No ratio. If $D = 1$ for every row, the outcome $D$ is fit exactly by the intercept: $\hat\beta_D$ has a 1 in the intercept and zeros elsewhere, and its residuals are zero. The outcome $TD = T$ is fit exactly by the $T$ column, so $\hat\beta_{TD}$ has a 1 there and zeros elsewhere. Then $\delta^\top \hat\beta_{TD} = 1$ and $\delta^\top \hat\beta_D = 0$, so

\[\hat\tau_R = \frac{\delta^\top \hat\beta_{TN}}{1} - \frac{\delta^\top (\hat\beta_{TN} - \hat\beta_N)}{1 - 0} = \delta^\top \hat\beta_N,\]

the LATE on the metric $N$. The columns $D$ and $TD$ have zero residuals, so they add nothing to the covariance.

No covariates. With no covariates, $(1, x)$ is just $1$, so $M_1 = (1, Z)$ and $M_2 = (1, T)$, and $\hat\beta_{10}$ is the weighted Wald ratio $\frac{\bar y_{Z=1} - \bar y_{Z=0}}{\bar T_{Z=1} - \bar T_{Z=0}}$, or the difference in means when $T = Z$.

Each reduction changes a different part of the model: the instruments, the grouping of the inner sums, the weights, or the outcome columns. So any combination of them is also a special case.

Monte Carlo

We repeat the experiment 500 times and compute both effects from the compressed dataset, with delta vectors at the end, each time. We also compute the ratio with the baseline line, read at $\bar x_C$, in place of the compliers’ levels, to show why the columns $TN$ and $TD$ are needed, and the kids slope, the coefficient on $T$ times the kids indicator, whose true value is $-0.2$.

Code

analyze <- function(seed) {
  set.seed(seed)
  df <- simulate_accounts() |>
    prepare()
  comp <- compress_rows(df, centered(df))
  fit <- fit_compressed(comp, d, y_cols)
  x_C <- complier_mean_compressed(comp, d, o, x_ref)
  base <- c(1, x_C - x_ref, rep(0, d)); delta <- c(rep(0, d), 1, x_C - x_ref)
  baseline_ratio <- \(B) { m0N <- sum(base * B[, "N"]); m0D <- sum(base * B[, "D"])
                           (m0N + sum(delta * B[, "N"])) / (m0D + sum(delta * B[, "D"])) - m0N / m0D }
  f <- \(th) baseline_ratio(matrix(th, nrow(fit$B), dimnames = list(NULL, y_cols)))
  theta <- as.vector(fit$B); grad <- gradient(f, theta)
  kids <- c(rep(0, d), 0, 0, 0, 1, 0, 0); i <- (match("y", y_cols) - 1) * 2 * d + seq_len(2 * d)
  bind_rows(effects(fit, x_C, x_ref),
            tibble(effect = "tau_R with the baseline line", estimate = Re(f(theta)),
                   se = sqrt(drop(t(grad) %*% fit$V %*% grad))),
            tibble(effect = "kids slope", estimate = sum(kids * fit$B[, "y"]),
                   se = sqrt(drop(t(kids) %*% fit$V[i, i] %*% kids)))) |>
    mutate(rep = seed)
}
invisible(gc())
mc <- mclapply(1:500, analyze, mc.cores = cores) |>
  bind_rows()
targets <- c(tau_LATE = truth[["tau_LATE"]], tau_R = truth[["tau_R"]],
             "tau_R with the baseline line" = truth[["tau_R"]], "kids slope" = -0.2)
mc_tab <- mc |>
  mutate(target = targets[effect]) |>
  group_by(effect) |>
  summarise(bias = mean(estimate - target), mc_se = sd(estimate) / sqrt(n()), sd = sd(estimate),
            se_over_sd = mean(se) / sd(estimate), coverage = mean(abs(estimate - target) <= qnorm(0.975) * se))

Code

mc_tab |>
  kable(digits = 4, col.names = c("Quantity", "Bias", "MC SE of bias", "SD", "SE/SD", "Coverage"))

Table 8: 500 simulated experiments. Bias is the average error, with its Monte Carlo standard error. SE/SD compares the average standard error with the spread of the estimates, and coverage is the share of 95% intervals that contain the target.

Quantity	Bias	MC SE of bias	SD	SE/SD	Coverage
kids slope	-0.0060	0.0091	0.2042	0.9074	0.920
tau_LATE	-0.0024	0.0027	0.0608	0.9803	0.950
tau_R	-0.0002	0.0002	0.0040	1.0089	0.934
tau_R with the baseline line	-0.0019	0.0002	0.0039	1.0058	0.924

Reading Table 8:

Both effects are unbiased, up to the Monte Carlo error of the bias, and their standard errors match the spread of the estimates, with coverage of 95.0% for the LATE and 93.4% for the ratio.
Using the baseline line instead of the compliers’ levels biases the ratio by -0.0019, about 11 Monte Carlo standard errors and 47% of a standard deviation. It estimates a different quantity, the ratio with the baseline of all profiles, so its bias does not shrink as the experiment grows while its standard error does, and its coverage falls toward zero in larger experiments.
The kids slope is also unbiased, so the same compressed dataset serves questions about heterogeneity as well as average effects. Its coverage, 92.0%, is a little below 95%, and its standard errors are 9% smaller than the spread of its estimates.

Other strategies for many covariates

The main text summarized 25 covariates with 3 prognostic scores. Here we compare that with three other ways to keep the compressed dataset small: approximating its second moments, binning the covariates, and clustering them.

A second order approximation: keep only the residual totals

The covariance needs, for every account, the sums $E_a$ of its residuals multiplied by its intercept and covariates. The second moments of the account summaries are what make $E_a$ computable for any coefficients, and they are the part that grows like $k^4$. For the effects we care about, most of $E_a$ barely matters.

An effect reads the treatment line at some covariates $x_0$. The first section of this appendix showed that the treatment line is a fit to the compliers, so in large samples its coefficients are driven by the moments of $(1, x)$ among compliers. For a line fit with an intercept, the influence of a data point with covariates $x$ on the line’s value at $x_0$ is proportional to

\[1 + (x_0 - \bar x_C)^\top \Sigma_C^{-1} (x - \bar x_C),\]

where $\bar x_C$ and $\Sigma_C$ are the mean and covariance of the covariates among compliers. This is the same leverage that appears in any regression with an intercept. The LATE and the ratio read the line at $x_0 = \bar x_C$, where the second term vanishes and every profile’s residual gets the same weight. So for these effects, an account’s influence depends approximately only on its residuals added up, $\sum_p \hat\varepsilon_{ap}$, which is the first row of $E_a$. Its second moments only need the account totals of $y$, of $(1, x)$ and of $T (1, x)$: $o + 2(k + 1)$ numbers per account instead of $o(k + 1) + (k+1)(k + 2)$.

The estimate still needs the first moments, but those never needed the account summaries: they are sums over the profile rows of each stratum, so a crossprod per stratum computes them. The second moments only need each account’s totals, which rowsum computes in one pass. The full account summaries are never built.

Code

compress_second_order <- function(df, X) {
  R <- cbind(1, X)
  Y <- as.matrix(df[y_cols])
  accounts <- df |>
    distinct(account, Z, day, w) |>
    arrange(account)
  strata <- group_by(accounts, Z, day, w)
  h_account <- group_indices(strata)
  h_row <- h_account[df$account]
  totals <- rowsum(cbind(Y, R, df$T * R), df$account)                       # account totals of y, (1, x), T (1, x)
  rows <- split(seq_len(nrow(df)), factor(h_row, seq_len(max(h_account))))
  group_keys(strata) |>
    mutate(G = tabulate(h_account),
           S = map(rows, \(i) crossprod(R[i, , drop = FALSE])),
           S_T = map(rows, \(i) crossprod(R[i, , drop = FALSE] * df$T[i], R[i, , drop = FALSE])),
           S_Y = map(rows, \(i) crossprod(R[i, , drop = FALSE], Y[i, , drop = FALSE])),
           second_moments = map(seq_len(n()), \(k) crossprod(totals[h_account == k, , drop = FALSE])))
}

fit_second_order <- function(compressed) {
  d <- nrow(compressed$S[[1]])
  o <- ncol(compressed$S_Y[[1]])
  M1_W_M2 <- reduce(map(seq_len(nrow(compressed)), \(k) with(compressed[k, ],
    w * rbind(cbind(S[[1]], S_T[[1]]), Z * cbind(S[[1]], S_T[[1]])))), `+`)
  M1_W_Y <- reduce(map(seq_len(nrow(compressed)), \(k) with(compressed[k, ], w * rbind(S_Y[[1]], Z * S_Y[[1]]))), `+`)
  B <- solve(M1_W_M2, M1_W_Y)
  colnames(B) <- y_cols
  L_totals <- cbind(diag(o), -t(B[1:d, , drop = FALSE]), -t(B[d + 1:d, , drop = FALSE]))   # residual totals from account totals
  first_row <- (0:(o - 1)) * d + 1
  meat <- reduce(map(seq_len(nrow(compressed)), \(k) {
    cov_E <- matrix(0, d * o, d * o)
    cov_E[first_row, first_row] <- L_totals %*% compressed$second_moments[[k]] %*% t(L_totals)
    compressed$w[k]^2 * stack_instruments(compressed$Z[k], d, o) %*% cov_E %*% t(stack_instruments(compressed$Z[k], d, o))
  }), `+`)
  M1_W_M2_inv <- kronecker(diag(o), solve(M1_W_M2))
  arm <- \(z_value) reduce(map(which(compressed$Z == z_value), \(k)
    compressed$w[k] * c(compressed$S_T[[k]][, 1], compressed$S[[k]][1, 1])), `+`)
  share <- \(v) v[1:d] / v[d + 1]
  diff <- share(arm(1)) - share(arm(0))
  list(B = B, V = M1_W_M2_inv %*% meat %*% t(M1_W_M2_inv), x_C = diff[-1] / diff[1])
}

compressed_many <- compress_second_order(df_many, X_many)
compressed_many |>
  transmute(Z, day, G, dim_S = map_chr(S, \(m) paste(dim(m), collapse = " x ")),
            dim_second_moments = map_chr(second_moments, \(m) paste(dim(m), collapse = " x ")))

# A tibble: 8 × 5
      Z   day     G dim_S   dim_second_moments
  <int> <int> <int> <chr>   <chr>             
1     0     1  7639 26 x 26 57 x 57           
2     0     2  5321 26 x 26 57 x 57           
3     0     3  2988 26 x 26 57 x 57           
4     0     4  1014 26 x 26 57 x 57           
5     1     1   412 26 x 26 57 x 57           
6     1     2   634 26 x 26 57 x 57           
7     1     3   981 26 x 26 57 x 57           
8     1     4  1011 26 x 26 57 x 57

The first moments hold two $26 \times 26$ matrices, $S$ and $S^T$, and one $26 \times 5$ matrix, $S^Y$, per stratum, and the second moments are a $57 \times 57$ matrix: the totals of 5 outcome columns, of the intercept and 25 covariates, and of the same 26 times $T$.

Binning and clustering

Binning or clustering the covariates turns them into one categorical covariate. That helps a lot when each account has one profile row: a categorical covariate can then be a grouping key, and the compressed dataset becomes counts and sums per cell. With several profiles per account, it helps much less. A categorical covariate with $L$ levels still enters the model as $L - 1$ indicators, so it costs about as much as $L - 1$ continuous covariates, and binning each of 25 covariates into deciles would create 225 indicators. Binning only saves space when it bins a few directions that matter.

We try two versions with 10 and 20 levels. The first bins the hours score into deciles, which keeps the direction that predicts the outcome. The second runs k-means on the standardized covariates, which ignores the outcome. Both are fit on the pre-experiment sample.

Code

bin_edges <- quantile(drop(cbind(1, X_history) %*% score_models$y), seq(0, 1, 0.1))
indicators <- function(level, n_levels) sapply(2:n_levels, \(j) as.numeric(level == j))
score_bins <- function(X) {
  level <- findInterval(drop(cbind(1, X) %*% score_models$y), bin_edges[2:10]) + 1
  indicators(level, 10)
}
scale_center <- colMeans(X_history)
scale_sd <- apply(X_history, 2, sd)
set.seed(1)
clusters <- kmeans(scale(X_history, scale_center, scale_sd), centers = 20, nstart = 3, iter.max = 50)
cluster_of <- function(X) {
  X_std <- scale(X, scale_center, scale_sd)
  distances <- -2 * X_std %*% t(clusters$centers) + rep(rowSums(clusters$centers^2), each = nrow(X_std))   # squared distance, less |x|^2
  max.col(-distances, ties.method = "first")
}
X_bins <- score_bins(X_many)
X_kmeans <- indicators(cluster_of(X_many), 20)
c(bins = ncol(X_bins), kmeans = ncol(X_kmeans))

  bins kmeans 
     9     19

Comparison

Table 9 compares the strategies on this experiment, with one row per strategy, next to two references: no covariates and the 5 base covariates. Every row uses the exact compressed dataset, except the second order approximation of all 25 covariates. For each one we report how many numbers its compressed dataset stores, counting each distinct number once, and the LATE with its standard error. The function sizes, from the section on many covariates, counts these numbers. A compressed dataset with $k$ continuous covariates stores, per stratum, the count, $o(k+1) + (k+1)(k+2)$ first moments, and either all their pairwise products (exact) or the products of the $o + 2(k + 1)$ account totals (second order). Indicators of one categorical covariate have fewer distinct products, because two indicators of the same covariate are never both 1.

Code

fit_second <- function(df, X) {
  fit <- fit_second_order(compress_second_order(df, X))
  effects(fit, fit$x_C, rep(0, ncol(X)))
}
strategies <- list(
  "no covariates" = list(X = matrix(numeric(0), nrow(df_many), 0), categorical = FALSE, second_order = FALSE),
  "5 base covariates" = list(X = centered(df_many)[, -1], categorical = FALSE, second_order = FALSE),
  "all 25 covariates" = list(X = X_many, categorical = FALSE, second_order = FALSE),
  "all 25 covariates, second order approximation" = list(X = X_many, categorical = FALSE, second_order = TRUE),
  "3 prognostic scores" = list(X = X_scores, categorical = FALSE, second_order = FALSE),
  "hours score in 10 bins" = list(X = X_bins, categorical = TRUE, second_order = FALSE),
  "k-means, 20 clusters" = list(X = X_kmeans, categorical = TRUE, second_order = FALSE))
many_tab <- imap_dfr(strategies, \(m, name) {
  fit <- if (m$second_order) fit_second(df_many, m$X) else exact_fits[[name]] %||% fit_exact(df_many, m$X)   # reuse the fits from the main text
  size <- sizes(ncol(m$X), m$categorical)
  tibble(strategy = name, k = ncol(m$X), numbers = size[[if (m$second_order) "second_order" else "exact"]],
         LATE = fit$estimate[1], se = fit$se[1])
})

Code

many_tab |>
  mutate(numbers = format(numbers, big.mark = ",")) |>
  kable(digits = 4, col.names = c("Strategy", "Covariates $k$", "Numbers stored", "LATE", "SE"))

Table 9: Strategies for 25 covariates on the running example, one row each. Numbers stored count each distinct number once. Every row uses the exact compressed dataset, except the second order approximation.

Strategy	Covariates $k$	Numbers stored	LATE	SE
no covariates	0	288	0.6720	0.0920
5 base covariates	5	21,608	0.6639	0.0572
all 25 covariates	25	2,778,888	0.6652	0.0526
all 25 covariates, second order approximation	25	19,888	0.6652	0.0526
3 prognostic scores	3	6,888	0.6656	0.0528
hours score in 10 bins	9	20,448	0.6544	0.0665
k-means, 20 clusters	19	80,088	0.6575	0.0717

The second order approximation stores 19,888 numbers instead of 2,778,888, and its standard error, 0.0526, is within 0.1% of the exact one, 0.0526. The three prognostic scores give a standard error of 0.0528, as small as all 25 covariates, with fewer numbers still, 6,888, and their storage does not grow with the number of covariates. Binning the hours score into deciles keeps some of the gain over no covariates, but here it does worse than the 5 base covariates, and the k-means clusters do worse still: clusters that ignore the outcome mostly separate profiles along directions that do not predict it.

One experiment can be lucky, so we repeat it 400 times, refitting every strategy each time. The score models, the bins and the clusters stay fixed, as they would in practice, since they come from data before the experiment. Each replication also retrains the scores on its own pre-experiment sample, which we use in the next section.

An exact fit with all 25 covariates takes about 20 seconds per experiment, so every strategy here takes its standard errors from the second order approximation, which matched the exact standard error to within 0.1% in Table 9. The point estimates do not depend on the approximation, so the spread of the estimates holds for the exact compressed dataset too, while the SE/SD ratio and the coverage test the approximation itself.

Code

analyze_many <- function(seed) {
  set.seed(seed)
  df <- simulate_accounts() |>
    add_signals() |>
    prepare()
  X <- cbind(centered(df)[, -1], as.matrix(df[signal_cols]))
  history <- simulate_accounts(10000, force_Z = 0) |>                      # a fresh pre-experiment sample, for the next section
    add_signals() |>
    prepare()
  X_history <- cbind(centered(history)[, -1], as.matrix(history[signal_cols]))
  models <- lm(as.matrix(history[c("y", "N", "D")]) ~ X_history) |>
    coef() |>
    as.data.frame() |>
    as.list()
  X_retrained <- map_dfc(models, \(b) drop(cbind(1, X) %*% b) - mean(drop(cbind(1, X_history) %*% b))) |>
    as.matrix()
  list("no covariates" = matrix(numeric(0), nrow(df), 0), "5 base covariates" = centered(df)[, -1],
       "all 25 covariates" = X, "3 prognostic scores" = score(X), "hours score in 10 bins" = score_bins(X),
       "k-means, 20 clusters" = indicators(cluster_of(X), 20), "3 prognostic scores, retrained" = X_retrained) |>
    imap_dfr(\(Xm, name) fit_second(df, Xm) |>
               mutate(strategy = name, rep = seed))
}
invisible(gc())
many_mc <- mclapply(1:400, analyze_many, mc.cores = cores) |>
  bind_rows()
many_mc_tab <- many_mc |>
  filter(effect == "tau_LATE", strategy != "3 prognostic scores, retrained") |>
  group_by(strategy) |>
  summarise(bias = mean(estimate - truth[["tau_LATE"]]), sd = sd(estimate), se_over_sd = mean(se) / sd(estimate),
            coverage = mean(abs(estimate - truth[["tau_LATE"]]) <= qnorm(0.975) * se)) |>
  arrange(sd)

Code

many_mc_tab |>
  kable(digits = 4, col.names = c("Strategy", "Bias", "SD", "SE/SD", "Coverage"))

Table 10: The LATE over 400 simulated experiments, one row per strategy, with standard errors from the second order approximation.

Strategy	Bias	SD	SE/SD	Coverage
all 25 covariates	-0.0023	0.0566	0.9590	0.9400
3 prognostic scores	-0.0023	0.0567	0.9608	0.9425
hours score in 10 bins	-0.0003	0.0622	0.9617	0.9425
5 base covariates	-0.0021	0.0626	0.9511	0.9400
k-means, 20 clusters	-0.0033	0.0699	0.9829	0.9600
no covariates	-0.0018	0.0900	0.9817	0.9475

Every strategy is unbiased, up to the Monte Carlo error of about 0.004 in the bias. The second order standard errors are between 95% and 98% of the spread of the estimates, and the coverage is between 94.0% and 96.0%, so the approximation is close to calibrated for the LATE. Across experiments, the prognostic scores are as precise as all 25 covariates, with a standard deviation of 0.0567 against 0.0566, and both beat the 5 base covariates (0.0626).

Table 11 shows how the compressed dataset grows with the number of covariates under each approach, with our 5 outcome columns and 8 strata.

Code

tibble(k = c(5, 10, 25, 50, 100)) |>
  mutate(exact = map_dbl(k, \(k) sizes(k)[["exact"]]),
         second_order = map_dbl(k, \(k) sizes(k)[["second_order"]]),
         scores_exact = sizes(3)[["exact"]],
         scores_second_order = sizes(3)[["second_order"]]) |>
  mutate(across(-k, \(v) format(v, big.mark = ","))) |>
  kable(col.names = c("Covariates $k$", "Exact", "Second order", "Scores, exact", "Scores, second order"))

Table 11: Numbers stored by the compressed dataset as the number of continuous covariates grows. Prognostic scores keep one covariate per raw outcome, whatever the number of covariates they summarize.

Covariates $k$	Exact	Second order	Scores, exact	Scores, second order
5	21,608	1,808	6,888	1,056
10	142,128	4,528	6,888	1,056
25	2,778,888	19,888	6,888	1,056
50	33,837,488	69,488	6,888	1,056
100	467,294,688	258,688	6,888	1,056

The exact compressed dataset grows like $k^4$, the second order one like $k^2$, and with prognostic scores it does not grow at all. The recommendations follow:

Reduce many covariates to a prognostic score per outcome, trained on data from before the experiment. This keeps the variance reduction and makes the compressed dataset independent of the number of covariates.
Keep individual covariates in the model only when we want to read effects along them, such as the effect for kids profiles. Those questions read the treatment line away from the compliers’ covariates, so they need the exact second moments for the covariates involved.
For the LATE and the ratio, use the second order compressed dataset, or delta vectors early when the profile rows can be read twice.
Avoid clustering or binning covariates without regard to the outcome. With several profiles per account it saves little space, and it gives up precision.

Why the score’s modeling error does not enter the standard error

A prognostic score $\hat y(x)$ is a prediction, and a different pre-experiment sample would give a different $\hat y$. Should the standard error of the effect include that uncertainty? It does not need to, for two reasons.

First, conditional on the pre-experiment sample, $\hat y$ is a fixed function. It is trained on data that is independent of the experiment, so for the experiment, $\hat y(x_{ap})$ is computed from observed covariates, exactly like $\log x$ or any other transformation we might choose. We condition on it the same way we condition on $x$. The gap between $\hat y(x)$ and the true $E[y \mid x]$ is real, but the model never assumes that the score is the true conditional mean, and its coefficient is estimated in the experiment, so the uncertainty of that coefficient is already in the cluster robust covariance.

Second, the effect does not depend on which $\hat y$ we use. With a randomized prompt, two stage least squares with any pre-experiment covariate is consistent for the same LATE. A worse score leaves more noise and a wider interval, but not a bias. Averaging over pre-experiment samples as well as experiments, the law of total variance gives

\[\text{Var}(\hat\tau) = E\left[\text{Var}(\hat\tau \mid \hat y)\right] + \text{Var}\left(E[\hat\tau \mid \hat y]\right).\]

The cluster robust standard error estimates the first term. The second term is close to zero, because $E[\hat\tau \mid \hat y]$ is close to $\tau$ for every $\hat y$: different pre-experiment samples give different but equally valid adjustments of the same effect.

This argument needs the score to be trained without the experiment’s outcomes. A flexible model trained on the experiment itself could absorb part of the effect or fit its noise, which would bias the estimate and shrink the standard error. Training on pre-experiment data, as here, or cross-fitting, where the score of each profile comes from a model trained on other profiles, avoids that.

The Monte Carlo of the previous section kept the scores fixed. To check the second point, we use the scores it retrained on a fresh pre-experiment sample of 10,000 accounts in every simulated experiment, so the spread of the estimates includes the variation of $\hat y$. As in the previous section, the standard errors come from the second order approximation.

Code

fresh_tab <- many_mc |>
  filter(strategy == "3 prognostic scores, retrained") |>
  group_by(effect) |>
  summarise(bias = mean(estimate - truth[effect]), sd = sd(estimate), se_over_sd = mean(se) / sd(estimate),
            coverage = mean(abs(estimate - truth[effect]) <= qnorm(0.975) * se))

Code

fresh_tab |>
  kable(digits = 4, format.args = list(scientific = FALSE), col.names = c("Effect", "Bias", "SD", "SE/SD", "Coverage"))

Table 12: The 400 simulated experiments of the previous section, each with prognostic scores retrained on its own pre-experiment sample.

Effect	Bias	SD	SE/SD	Coverage
tau_LATE	-0.0023	0.0568	0.959	0.9400
tau_R	-0.0001	0.0039	1.019	0.9475

Even with the scores retrained every time, the standard errors, which ignore the variation of $\hat y$, are within Monte Carlo error of the spread of the estimates, with coverage of 94.0% for the LATE and 94.8% for the ratio. With 400 replications, SE/SD is itself uncertain by about $1/\sqrt{2 \cdot 399} \approx 3.5\%$, and the coverage by about $\sqrt{0.95 \cdot 0.05 / 400} \approx 1.1$ percentage points. The standard deviation of the LATE, 0.0568, is close to the 0.0567 we saw with fixed scores, so retraining the scores adds little variance. The bias of the LATE, -0.0023, is about -0.8 Monte Carlo standard errors from zero, and the same experiments with fixed scores gave -0.0023.