FastRerandomize: Fast Rerandomization Using Accelerated Computing

TL;DR

FastRerandomize is an R package that significantly accelerates rerandomization in experimental design through GPU/TPU support, memory efficiency, and optimized kernels, enabling high-dimensional, scalable, and precise covariate balance checks.

Contribution

The paper introduces a novel, scalable rerandomization method with GPU acceleration, memory-efficient storage, and auto-vectorized kernels, making high-dimensional experimental design feasible.

Findings

Order-of-magnitude speedups over baseline workflows

Enables exact or Monte Carlo rerandomization at large scales

Improves precision and power in high-dimensional experiments

Abstract

We present fastrerandomize, an R package for fast, scalable rerandomization in experimental design. Rerandomization improves precision by discarding treatment assignments that fail a prespecified covariate-balance criterion, but existing implementations can become computationally prohibitive as the number of units or covariates grows. fastrerandomize introduces three complementary advances: (i) optional GPU/TPU acceleration to parallelize balance checks, (ii) memory-efficient key-only storage that avoids retaining full assignment matrices, and (iii) auto-vectorized, just-in-time compiled kernels for batched candidate generation and inference. This approach enables exact or Monte Carlo rerandomization at previously intractable scales, making it practical to adopt the tighter balance thresholds required in modern high-dimensional experiments while simultaneously quantifying the resulting gains in precision and power for a given covariate set. Our approach also supports randomization-based testing conditioned on acceptance. In controlled benchmarks, we observe order-of-magnitude speedups over baseline workflows, with larger gains as the sample size or dimensionality grows, translating into improved precision of causal estimates.