Optimal Correlation for Bernoulli Trials with Covariates

TL;DR

This paper investigates optimal correlation structures for treatment assignment in Bernoulli trials with covariates, proposing designs that minimize variance of the IPW estimator and demonstrating their effectiveness through theoretical proofs and simulations.

Contribution

It introduces a novel proof for the optimal two-cluster design and proposes a hybrid approach with multiple clusters to improve robustness.

Findings

The two-cluster design can be computed via a 0-1 knapsack problem.

A shift-invariant version of the design improves on stratified designs.

Simulations show strong performance of the proposed methods.

Abstract

Given covariates for $n$ units, each of which is to receive a treatment with probability $1/2$, we study the question of how best to correlate their treatment assignments to minimize the variance of the IPW estimator of the average treatment effect. Past work by \cite{bai2022} found that the optimal stratified experiment is a matched-pair design, where the matching depends on oracle knowledge of the distributions of potential outcomes given covariates. We study the strictly broader class of all admissible correlation structures, for which \cite{cytrynbaum2023} recently showed that the optimal design is to divide the units into two clusters and uniformly assign treatment to exactly one of them. This design can be computed by solving a 0-1 knapsack problem that uses the same oracle information. We derive a novel proof of this fact using a result about admissible Bernoulli correlations. We also show how to construct a shift-invariant version of the design that still improves on the optimal stratified design by ensuring that exactly half of the units are treated. A method with just two clusters is not robust to a bad proxy for the oracle, and we mitigate this with a hybrid that uses $O(n^\alpha)$ clusters for $0<\alpha<1$. Under certain assumptions, we also derive a CLT for the IPW estimator under these designs and a consistent estimator of the variance. We compare the proposed designs to the optimal stratified design in simulated examples and find strong performance.