NEAR: Neighborhood Effect-Aware Regularization for Two-Stage CATE Estimation

Estimating the conditional average treatment effect (CATE) from observational data is a central task in causal inference. A common approach is two-stage CATE learning, which first estimates nuisance functions and then uses them to fit a target CATE model. However, the resulting second-stage CATE targets or losses depend on estimated nuisance functions and can be noisy. Existing methods use regularization and information sharing to stabilize second-stage estimation, yet it remains unclear how noisy first-stage CATE estimates should guide the selection of observations to regularize together. We propose Neighborhood Effect-Aware Regularization (NEAR), a modular second-stage regularizer that uses first-stage plug-in CATE estimates to construct a fixed graph and regularizes CATE predictions over pairs with similar plug-in CATEs. We show how plug-in CATE estimation errors control true CATE variation along the selected pairs and derive finite-sample CATE risk bounds for squared-loss second-stage regression using a stabilized version of the doubly robust pseudo-outcome. Across six two-stage learners and three benchmarks, NEAR lowers mean CATE estimation error. Graph support comparisons and perturbation diagnostics further support using plug-in CATE similarity to select pairs for regularization.