Fit CATE Once: Model-Assisted Randomization Tests Without Sample Splitting

2026-05-09

AI Summary

This study addresses the challenge of flexibly modeling heterogeneous treatment effects while preserving the validity of randomization-based inference. The authors propose a model-assisted randomization test that avoids sample splitting by estimating the unsigned conditional average treatment effect (CATE) through the residual covariance structure, retaining the original treatment assignment for inference, and optimizing sign assignments to best fit the observed outcomes. This approach represents the first seamless integration of flexible CATE modeling with randomization testing, achieving strict Type I error control while substantially improving statistical power. Moreover, it enables the identification of heterogeneous subgroups exhibiting distinct treatment responses. Empirical evaluations demonstrate that the method outperforms conventional covariate adjustment and sample-splitting strategies in both power and precision.

Abstract

Randomization tests and flexible treatment-effect models offer complementary strengths for analyzing data from randomized panel experiments: the former provide valid inference under the known assignment mechanism, while the latter can capture complex patterns of effect heterogeneity. We develop model-assisted randomization tests that combine these strengths without sample splitting. The key idea is to estimate an unsigned version of the conditional average treatment effect (CATE) from the covariance structure of residualized outcomes, while leaving the realized assignments for randomization inference. The remaining sign can be chosen to best fit the observed outcomes. We establish identification and consistency for the proposed unsigned CATE estimators, as well as validity for the CATE-assisted randomization tests. Across synthetic and semi-synthetic experiments, the CATE-assisted randomization tests control Type I error and achieve higher power than covariate-adjusted and sample-split alternatives. Finally, we show that the assignment-free CATE estimates can be used to discover heterogeneous subgroups and test subgroup-specific treatment effects.

Problem

• conditional average treatment effect
• randomization tests
• effect heterogeneity
• model-assisted inference
• panel experiments

Innovation

• model-assisted randomization test
• conditional average treatment effect
• no sample splitting
• effect heterogeneity
• residualized outcomes