Nonlinear Covariate Balance in Experimental Design

We study experimental designs that balance nonlinear functions of covariates, extending classical methods that primarily target linear balance. Building on the Gram-Schmidt Walk (GSW) framework of Harshaw et al (2024) for linear covariate balancing, we introduce a design that directly controls imbalance in nonlinear structure, including polynomial and more general smooth function classes. Like GSW, the proposed design retains sufficient robustness against model misspecification. Our implementation operates directly on a Gram matrix, avoiding the expensive step of explicitly constructing the nonlinear covariate expansions. We further accelerate the nonlinear design via a low-rank approximation of the Gram matrix, achieving runtimes comparable to the GSW of Harshaw et al (2024) while preserving nonlinear covariate balance and robustness.

Randomized controlled trials are used to test whether a new treatment works. A key design question is how to assign subjects to treatment and control groups. Many randomization methods aim to balance simple averages of background features, but outcomes may depend on nonlinear patterns, such as feature squares or feature interactions. We propose a new randomization method that balances these richer patterns. The method is efficient and, in simulations, improves accuracy when outcomes depend on nonlinear feature structure. This gives practitioners a new tool for designing experiments when nonlinear structure matters.