Fast Rerandomization via the BRAIN
Jiuyao Lu, Daogao Liu, Zhanran Lin, Xiaomeng Wang
TL;DR
Rerandomization improves causal inference by balancing covariates, but traditional acceptance-rejection sampling is prohibitively slow when thousands of balanced allocations are needed. The authors propose BRAIN, a fast, metaheuristic rerandomization method that minimizes the Mahalanobis-distance balance criterion $M(W)=(W-\frac{n_t}{n}1_n)^T H (W-\frac{n_t}{n}1_n)$ over binary allocations $W$ with $\sum_i W_i=n_t$, using a descent-based local search over $L$ disjoint swaps and a perturbation of $S$ swaps, terminating when $M(W)\le a$. BRAIN preserves the standard randomization-based properties: the estimator $\widehat{\tau}$ is unbiased and achieves variance reduction relative to complete randomization, with a quantifiable lower bound $\frac{V_{CR}-\mathrm{Var}[\widehat{\tau}]}{V_{CR}} \ge (1-\frac{a}{p})R^2$, where $R^2=\frac{n\beta^T S_{XX} \beta}{n_t n_c \mathrm{V}_{CR}}$. The method extends to sequential, stratified, and cluster randomized designs (SeqBRAIN, StratBRAIN, ClustBRAIN) with maintained statistical guarantees. Empirical results from simulations and a real clinical trial show BRAIN delivers comparable statistical performance to existing rerandomization methods while achieving orders-of-magnitude faster sampling, enabling fast, scalable, randomization-based inference in diverse experimental settings.
Abstract
Randomized experiments are a crucial tool for causal inference in many different fields. Rerandomization addresses any covariate imbalance in such experiments by resampling treatment assignments until certain balance criteria are satisfied. However, rerandomization based on naïve acceptance-rejection sampling is computationally inefficient, especially when numerous independent assignments are required to perform randomization-based statistical inference. Existing acceleration methods are suboptimal and not applicable in structured experiments, including stratified and clustered experiments. Based on metaheuristics in integer programming, we propose BRAIN -- a novel computationally-lightweight methodology that ensures covariate balance in randomized experiments while significantly accelerating the computation. Our BRAIN method provides unbiased treatment effect estimators with reduced variance compared to complete randomization, preserving the desirable statistical properties of traditional rerandomization. Simulation studies and a real data example demonstrate the benefits of our method in fast sampling while retaining the appealing statistical guarantees.