Multi-arm experiments

Most examples on randomization inference assume that you have an experiment with two arms: a control and a treatment. What if you have more than two? You can still do randomization inference, but you need to make some choices.1

Theory

Suppose you ran a randomized experiment with three arms:

  • \(T \in \{0,1,2\}\)
    • \(T=0\): Control
    • \(T=1\): Treatment A
    • \(T=2\): Treatment B

A multi-arm experiment does not have “one” treatment effect. You must decide which contrast you want to test and report. The three most common are:

  1. A vs Control: \(H_0:\ \tau_{A0}=0\)
  2. B vs Control: \(H_0:\ \tau_{B0}=0\)
  3. A vs B: \(H_0:\ \tau_{AB}=0\)

In this example, I focus on the first one: A vs Control.

Randomization inference tests a sharp null about outcomes under different assignments:

\[ H_0:\ Y_i(1)=Y_i(0)\ \text{for all } i \text{ in the units eligible to switch between } T=0 \text{ and } T=1. \]

The key phrase is “eligible to switch”.

Option 1

If you only care about A vs Control, analyze only units with \(T \in \{0,1\}\). This is the simplest option.

  • Estimation sample: \(T \in \{0,1\}\)
  • Re-randomization: shuffle A/Control labels within that sample
  • Interpretation: RI is conditional on the experiment that could have assigned only A or Control to those units.

Option 2

If you want more precision, you can use the full dataset to estimate a covariate-adjusted statistic, but only allow reassignments among the A/Control units. Treatment-B units stay fixed as \(T=2\) in every permutation.

  • Estimation sample: \(T \in \{0,1,2\}\)
  • Re-randomization: shuffle A/Control labels only among units with \(T \in \{0,1\}\), keeping all \(T=2\) units fixed
  • Interpretation: RI is conditional on the realized set of B assignments. You test A vs Control “holding B fixed”.

If you keep B in the regression, you should include a B indicator (or use a full set of arm indicators) so B is not implicitly pooled into Control.

Discussion

If you have no covariates and saturated arm indicators: \[ y_i = \alpha + \beta_A \mathbf{1}[T_i=1] + \beta_B \mathbf{1}[T_i=2] + \varepsilon_i, \] \(\beta_A\) is the difference in mean outcomes between A and Control. Including B observations does not change the meaning of \(\beta_A=\bar y_{A} - \bar y_{0}\). That is, option 1 and option 2 will typically produce the same \(\hat\beta_A\) (up to edge-case numerical details).

If you do have covariates, \[ y_i = \alpha + \beta_A \mathbf{1}[T_i=1] + \beta_B \mathbf{1}[T_i=2] + X_i'\gamma + \varepsilon_i, \] \(\hat\beta_A\) is computed after “partialling out” \(X\). If you estimate \(\gamma\) using all arms (option 2) versus only A/Control (Option 1), then \(\hat\gamma\) can change, which can change \(\hat\beta_A\).

Which one is “correct”? That is up to you. Are you looking for conditional or an unconditional estimate?

Implementation

ritest (Python) requires the permutation variable to be binary, so for A vs Control you typically construct:

\[ A_i = \mathbf{1}[T_i=1]. \]

Option 1 (drop B)

  1. Filter to \(T \in \{0,1\}\)
  2. Create \(A\)
  3. Call ritest normally (permuting A)

Minimal sketch:

df_sub = df[df["treatment"].isin([0, 1])].copy()
df_sub["A"] = (df_sub["treatment"] == 1).astype(int)

res = ritest(
    df=df_sub,
    permute_var="A",
    formula="y ~ 1 + A + x1 + x2",   # include covariates if you want
    stat="A",
    reps=5000,
    seed=123,
)

Option 2 (keep B fixed)

This option requires much more work, it emulates Stata’s fixlevels(#) option.

  1. Work on the full dataset.
  2. Create A and a separate B indicator.
  3. Build a permutation matrix that only permutes A among units with \(T \in \{0,1\}\).2
  4. Pass that matrix via permutations= on the ritest call..

Sketch:

df = df.copy()
df["A"] = (df["treatment"] == 1).astype(int)
df["B"] = (df["treatment"] == 2).astype(int)

A = df["A"].to_numpy()
mask = (df["B"].to_numpy() == 0)     # only Control/A units are shuffled

perms = np.empty((reps, len(df)), dtype=A.dtype)
for r in range(reps):
    pr = A.copy()
    pr[mask] = rng.permutation(pr[mask])
    perms[r] = pr

res = ritest(
    df=df,
    permute_var="A",
    permutations=perms,              # B stays fixed because A=0 for B always
    formula="y ~ 1 + A + B + x1 + x2",
    stat="A",
    seed=123,
)
Back to top

Footnotes

  1. This section borrows heavily from Simon Heß’s discussion here.↩︎

  2. If your experimental design involves stratification or clustering, you must build permutations that respect that same structure when you supply permutations=.↩︎