Table of Contents
Fetching ...

Fast Rerandomization for Balancing Covariates in Randomized Experiments: A Metropolis-Hastings Framework

Jiuyao Lu, Tianruo Zhang, Ke Zhu

TL;DR

This work tackles covariate imbalance in randomized experiments, where rerandomization using the Mahalanobis distance $M(\mathbf{W})$ can achieve strong balance but suffers from prohibitively slow rejection sampling at high dimensions. It introduces a Metropolis-Hastings framework with pair-switching to define a stationary distribution $\pi(\mathbf{W}) \propto M(\mathbf{W})^{-1/T}$ and then applies a sampling-importance resampling step to recover uniform draws over the acceptable set $\mathcal{W}_a$, preserving classical inference guarantees. The resulting PSRSRR algorithm dramatically speeds up sampling by orders of magnitude while maintaining exact and asymptotic validity, demonstrated through extensive simulations and real-data applications (STAR and reserpine). The approach provides a practical, theory-grounded path to strict rerandomization thresholds, with potential extensions to complex experimental designs and covariate-prioritized balance criteria. Overall, PSRSRR makes fast, uniform rerandomization feasible at stringent balance thresholds, enabling robust Fisher randomization tests and efficient estimation in modern causal inference pipelines.

Abstract

Balancing covariates is critical for credible and efficient randomized experiments. Rerandomization addresses this by repeatedly generating treatment assignments until covariate balance meets a prespecified threshold. By shrinking this threshold, it can achieve arbitrarily strong balance, with established results guaranteeing optimal estimation and valid inference in both finite-sample and asymptotic settings across diverse complex experimental settings. Despite its rigorous theoretical foundations, practical use is limited by the extreme inefficiency of rejection sampling, which becomes prohibitively slow under small thresholds and often forces practitioners to adopt suboptimal settings, leading to degraded performance. Existing work focusing on acceleration typically fail to maintain the uniformity over the acceptable assignment space, thus losing the theoretical grounds of classical rerandomization. Building upon a Metropolis-Hastings framework, we address this challenge by introducing an additional sampling-importance resampling step, which restores uniformity and preserves statistical guarantees. Our proposed algorithm, PSRSRR, achieves speedups ranging from 10 to 10,000 times while maintaining exact and asymptotic validity, as demonstrated by simulations and two real-data applications.

Fast Rerandomization for Balancing Covariates in Randomized Experiments: A Metropolis-Hastings Framework

TL;DR

This work tackles covariate imbalance in randomized experiments, where rerandomization using the Mahalanobis distance can achieve strong balance but suffers from prohibitively slow rejection sampling at high dimensions. It introduces a Metropolis-Hastings framework with pair-switching to define a stationary distribution and then applies a sampling-importance resampling step to recover uniform draws over the acceptable set , preserving classical inference guarantees. The resulting PSRSRR algorithm dramatically speeds up sampling by orders of magnitude while maintaining exact and asymptotic validity, demonstrated through extensive simulations and real-data applications (STAR and reserpine). The approach provides a practical, theory-grounded path to strict rerandomization thresholds, with potential extensions to complex experimental designs and covariate-prioritized balance criteria. Overall, PSRSRR makes fast, uniform rerandomization feasible at stringent balance thresholds, enabling robust Fisher randomization tests and efficient estimation in modern causal inference pipelines.

Abstract

Balancing covariates is critical for credible and efficient randomized experiments. Rerandomization addresses this by repeatedly generating treatment assignments until covariate balance meets a prespecified threshold. By shrinking this threshold, it can achieve arbitrarily strong balance, with established results guaranteeing optimal estimation and valid inference in both finite-sample and asymptotic settings across diverse complex experimental settings. Despite its rigorous theoretical foundations, practical use is limited by the extreme inefficiency of rejection sampling, which becomes prohibitively slow under small thresholds and often forces practitioners to adopt suboptimal settings, leading to degraded performance. Existing work focusing on acceleration typically fail to maintain the uniformity over the acceptable assignment space, thus losing the theoretical grounds of classical rerandomization. Building upon a Metropolis-Hastings framework, we address this challenge by introducing an additional sampling-importance resampling step, which restores uniformity and preserves statistical guarantees. Our proposed algorithm, PSRSRR, achieves speedups ranging from 10 to 10,000 times while maintaining exact and asymptotic validity, as demonstrated by simulations and two real-data applications.
Paper Structure (47 sections, 10 theorems, 57 equations, 13 figures, 8 tables, 3 algorithms)

This paper contains 47 sections, 10 theorems, 57 equations, 13 figures, 8 tables, 3 algorithms.

Key Result

Theorem 3.1

The limiting distribution of the Markov chain $\{\mathbf{W}^{(t)}\}_{t \ge 0}$ with temperature $T$ is $\pi(\mathbf{W})=\frac{M(\mathbf{W})^{-1/T}}{\sum_{\mathbf{W}^* \in \mathcal{W}} M(\mathbf{W}^*)^{-1/T}}$ for any $\mathbf{W} \in \mathcal{W}$. $\pi$ is also the stationary distribution.

Figures (13)

  • Figure 1: Comparison plots of, from top to bottom, MSE (relative to CR), CI length (relative to CR), coverage rate, power and type I error, by sample size $n$ and number of covariates $p$. The red dashed lines represent the nominal level of 95% CI coverage rate and 0.05 type I error. Better performance is indicated by lower MSE and CI Length, a coverage rate at or above 95%, higher power, and a type I error at or below 0.05.
  • Figure 2: Comparison of sampling time (on a $\log_{10}$ scale) by sample size $n$ and number of covariates $p$. CR is plotted to provide a baseline for the computational cost of a single random draw, rather than as a benchmark to outperform.
  • Figure 3: Box-plot of standardized differences in covariate means for STAR data. The diamonds indicate the standardized difference for the actual assignment in the experiment. The values on the right side are PRIVs, the empirical percent reductions in variance compared with CR of each method.
  • Figure 4: Temperature Analysis: Comparison of MSE
  • Figure 5: Temperature Analysis: Comparison of Sampling Time
  • ...and 8 more figures

Theorems & Definitions (11)

  • Theorem 3.1
  • Theorem 3.2
  • Corollary 3.3
  • Lemma A.1
  • Definition A.2: Reweighting Operators
  • Theorem A.3
  • Lemma A.4
  • Lemma A.5
  • Lemma A.6
  • Lemma A.7
  • ...and 1 more