Optimizing precision in stepped-wedge designs via machine learning and quadratic inference functions
Liangbo Lyu, Bingkai Wang
TL;DR
This work tackles precision in stepped-wedge designs by marrying cross-fitted, flexible covariate adjustment with adaptive correlation learning via quadratic inference functions (QIF). The proposed estimators are consistent and asymptotically normal under mild $L_2$-convergence of nuisance components, and, when mean models are correctly specified, can attain minimal asymptotic variance; crucially, QIF never degrades efficiency relative to independence and often improves precision by combining multiple correlation structures. Through simulations and two real-world applications, the authors demonstrate substantial finite-sample gains from combining ML-based covariate adjustment with QIF, especially in larger samples, while maintaining valid inference and accommodating treatment-effect heterogeneity across exposure duration and calendar time. The framework provides a principled, implementable alternative to standard mixed-effects approaches, with robust performance across diverse design features and settings, and it offers clear paths for extensions to covariate-adaptive randomization and missing data scenarios.
Abstract
Stepped-wedge designs are increasingly used in randomized experiments to accommodate logistical and ethical constraints by staggering treatment roll-out over time. Despite their popularity, existing analytical methods largely rely on parametric models with linear covariate adjustment and prespecified correlation structures, which may limit achievable precision in practice. We propose a new class of estimators for the causal average treatment effect in stepped-wedge designs that optimizes precision through flexible, machine-learning-based covariate adjustment to capture complex outcome-covariate relationships, together with quadratic inference functions to adaptively learn the correlation structure. We establish consistency and asymptotic normality under mild conditions requiring only $L_2$ convergence of nuisance estimators, even under model misspecification, and characterize when the estimator attains the minimal asymptotic variance. Moreover, we prove that the proposed estimator never reduces efficiency relative to an independence working correlation. The proposed method further accommodates treatment-effect heterogeneity across both exposure duration and calendar time. Finally, we demonstrate our methods through simulation studies and reanalyses of two empirical studies that differ substantially in research area and key design parameters.