Nonlinear Covariate Balance in Experimental Design

Qing Chen, Peng Zhang
Proceedings of the 43rd International Conference on Machine Learning, PMLR 306:18358-18379, 2026.

Abstract

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.

Cite this Paper

BibTeX
@InProceedings{pmlr-v306-chen26gn, title = {Nonlinear Covariate Balance in Experimental Design}, author = {Chen, Qing and Zhang, Peng}, booktitle = {Proceedings of the 43rd International Conference on Machine Learning}, pages = {18358--18379}, year = {2026}, editor = {Zhang, Tong and Dudik, Miroslav and Jaggi, Martin and Agarwal, Alekh and Li, Sharon and Schuurmans, Dale and Zhu, Jerry and Berkenkamp, Felix and Dong, Hanze and Bietti, Alberto}, volume = {306}, series = {Proceedings of Machine Learning Research}, month = {06--11 Jul}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v306/main/assets/chen26gn/chen26gn.pdf}, url = {https://proceedings.mlr.press/v306/chen26gn.html}, abstract = {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.} }
Endnote
%0 Conference Paper %T Nonlinear Covariate Balance in Experimental Design %A Qing Chen %A Peng Zhang %B Proceedings of the 43rd International Conference on Machine Learning %C Proceedings of Machine Learning Research %D 2026 %E Tong Zhang %E Miroslav Dudik %E Martin Jaggi %E Alekh Agarwal %E Sharon Li %E Dale Schuurmans %E Jerry Zhu %E Felix Berkenkamp %E Hanze Dong %E Alberto Bietti %F pmlr-v306-chen26gn %I PMLR %P 18358--18379 %U https://proceedings.mlr.press/v306/chen26gn.html %V 306 %X 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.
APA
Chen, Q. & Zhang, P.. (2026). Nonlinear Covariate Balance in Experimental Design. Proceedings of the 43rd International Conference on Machine Learning, in Proceedings of Machine Learning Research 306:18358-18379 Available from https://proceedings.mlr.press/v306/chen26gn.html.

Related Material