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

1. Problem Statement

Randomized panel experiments generate sequences of outcomes suitable for studying dynamic causal effects. While randomization tests offer finite-sample validity grounded in the known assignment mechanism, they often lack power when treatment effects are heterogeneous. Conversely, flexible treatment-effect models (such as Conditional Average Treatment Effect, or CATE, models) can capture complex heterogeneity but typically require strong modeling assumptions or sample splitting to maintain validity in randomization inference.

The core challenge addressed is the trade-off in constructing model-assisted randomization tests:

1. Validity: The nuisance model (e.g., a CATE estimator) must be independent of the realized treatment assignments ($A[N]$) to ensure the randomization distribution remains valid. Fitting a model on $A[N]$ and using it in the test statistic breaks the symmetry required for valid p-values.
2. Power: To improve power, the test statistic should utilize accurate estimates of treatment effect heterogeneity.
3. Current Limitations: Existing solutions either refit the model for every simulated assignment (computationally prohibitive) or use sample splitting (sacrificing statistical power by discarding data).

This paper proposes a method to estimate CATEs without using the realized assignments, thereby preserving validity without the efficiency loss of sample splitting.

2. Methodology

The proposed approach, Fit CATE Once, leverages the covariance structure of residualized outcomes in panel data to estimate the magnitude of treatment effects, leaving only the sign to be determined.

2.1. Setup and Assumptions

• Data: $N$ units observed over $T$ time periods with covariates $X_i$, outcomes $Y_{i,t}$, and a treatment start time $A_i$.
• Additive Outcome Model: The observed outcome follows $Y_{i,t} = \mu_{0,t}(X_i) + \sum_{l=0}^{t-1} \mathbb{1}\{A_i = t-l\}\tau_l(X_i) + \epsilon_{i,t}$.
• No Anticipation: Treatment starting after time $t$ does not affect $Y_{i,t}$.
• Residualization: Outcomes are residualized by subtracting the conditional mean $\mu_t(X_i)$, yielding residuals $R_{i,t}$.

2.2. Unsigned CATE Estimation

The key insight is that under the additive model, the second-order moments (covariances) of the residualized outcomes encode the magnitude of treatment effects, independent of the realized assignment signs.

• Diagonal Moments (Lag-Invariant Effects): If the treatment effect is constant across lags ($\tau_l(x) \equiv \tau^*(x)$), the diagonal entries of the residual covariance matrix $\Sigma_R(x)$ depend on the squared effect magnitude $\tau^{*2}(x)$ and a known design factor $v_t(x)$. By comparing diagonal moments across time, the nuisance variance can be differenced out, identifying $|\tau^*(x)|$.
• Off-Diagonal Moments (Lag-Varying Effects): If effects vary by lag and residual errors are serially uncorrelated (Assumption 4), the off-diagonal residual covariances depend on the CATE vector $\tau(x)$ through quadratic forms: $c(x) = F_x(\tau) = (\tau^\top H_{t,s}(x) \tau)_{(t,s)}$.
  • This creates a moment-matching problem where the off-diagonal covariances identify the vector $\tau(x)$ up to a global sign.
  • The authors propose a convex relaxation (estimating the rank-one matrix $B = \tau\tau^\top$ via semidefinite programming) and a nonlinear least squares (NLS) approach to solve for the unsigned estimator $\hat{\tau}_\pm(x)$.

2.3. Sign Orientation

Since the moment-based estimators only recover the magnitude (or direction up to a sign), a final step is required to determine the sign.

• Assignment-Free Orientation: The sign is chosen to best fit the observed residualized outcome trajectory for each unit, minimizing the squared error between the observed residuals and the model predicted by the unsigned estimator.
• Partial Assignment Orientation: Alternatively, a small subset of observed assignments can be used solely to learn the sign, leaving the majority of units for the randomization test. This is computationally cheaper than re-fitting a full CATE model.

2.4. CATE-Assisted Randomization Test

Once the signed CATE estimator $\hat{\tau}(X_i)$ is obtained, it is used to construct a test statistic (e.g., a likelihood-ratio score or AIPW score) that is more sensitive to heterogeneous effects.

• Validity: Because the CATE estimator is derived from residual covariances (which depend on covariates and outcomes but not the realized $A[N]$), the conditional independence condition $\hat{\eta} \perp A[N] \mid X, Y$ holds.
• Procedure: The test computes a p-value by comparing the observed statistic against the distribution of statistics generated by permuting the assignment vector $A[N]$ while keeping the nuisance model fixed.

3. Key Contributions

1. Assignment-Free CATE Estimation: The paper establishes that in panel experiments, the covariance structure of residualized outcomes allows for the identification of CATE magnitudes (and vectors up to sign) without fitting on realized assignments.
2. Theoretical Guarantees:
   • Identification: Proves local and global identification of the unsigned CATE vector up to a global sign under rank conditions on the Jacobian of the moment map.
   • Consistency: Establishes the consistency of the convex relaxation and NLS estimators for the unsigned CATE as $N \to \infty$.
3. Valid Model-Assisted Tests: Demonstrates that using these assignment-free estimators in randomization tests preserves finite-sample validity (Type I error control) while avoiding the power loss associated with sample splitting.
4. Subgroup Discovery: Shows that assignment-free CATE estimates can be used post-hoc to discover heterogeneous subgroups and conduct valid randomization tests within those subgroups.

4. Experimental Results

The authors evaluate the method through synthetic, semi-synthetic, and consistency experiments:

• Type I Error Control: Across various sample sizes and lag structures (lag-invariant and lag-varying), the proposed CATE-assisted randomization test (RT CATE) controls Type I error at the nominal level (0.1), comparable to unadjusted (RT DM) and covariate-adjusted (RT cDM) tests.
• Power: RT CATE consistently achieves higher power than:
  • Unadjusted tests (RT DM).
  • Covariate-adjusted tests using only outcome means (RT cDM).
  • Sample-splitting alternatives (RT SS), which sacrifice data for validity.
  • This gain is particularly pronounced when treatment effects are heterogeneous across units or lags.
• Consistency: The unsigned estimator converges to the true CATE magnitude (up to sign) as sample size increases.
• Warm-Starting: Using the unsigned estimator as a "warm start" allows for accurate signed CATE estimation with significantly fewer revealed assignments than learning the full CATE vector from scratch.
• Semi-Synthetic Validation: Using a county-level teen employment dataset (based on Callaway and Sant'Anna, 2021), the method outperforms standard event-study Two-Way Fixed Effects (TWFE) regressions in both size control (avoiding over-rejection) and power.
• Subgroup Analysis: The method successfully recovers threshold-based subgroups and supports valid inference within those discovered groups.

5. Significance and Claims

The paper claims that Fit CATE Once resolves the fundamental tension between validity and power in model-assisted randomization tests for panel data. By exploiting the second-order moment structure of residuals, the method extracts treatment effect information that is "assignment-free," allowing researchers to use flexible, high-power test statistics without compromising the randomization inference framework.

The authors emphasize that this approach is not limited to the specific panel designs studied but suggests a broader class of model-assisted tests where residuals depend on assignments only through treatment effects. The method is presented as a practical solution that avoids the computational burden of refitting models for every permutation and the statistical inefficiency of sample splitting.

Limitations noted by the authors:

• The off-diagonal estimator assumes serially uncorrelated residual errors; correlated errors would confound treatment effects with error structure.
• The sign recovery step relies on the signal-to-noise ratio and the timing of treatment; early treatment starts provide more information for sign recovery than late starts.
• Global identification relies on rank conditions that, while generally satisfied in standard designs, may fail in specific degenerate cases.