Propensity Patchwork Kriging for Scalable Inference on Heterogeneous Treatment Effects

TL;DR

This paper introduces Propensity Patchwork Kriging, a scalable Gaussian process method for estimating heterogeneous treatment effects by partitioning data based on propensity scores, reducing computational costs while maintaining estimate continuity.

Contribution

It extends Patchwork Kriging to causal inference, enabling efficient HTE estimation by focusing on propensity score partitions and continuity constraints.

Findings

Method reduces computational cost significantly.

Maintains continuity of HTE estimates across regions.

Demonstrated effectiveness through simulations and real data.

Abstract

Gaussian process-based models are attractive for estimating heterogeneous treatment effects (HTE), but their computational cost limits scalability in causal inference settings. In this work, we address this challenge by extending Patchwork Kriging into the causal inference framework. Our proposed method partitions the data according to the estimated propensity score and applies Patchwork Kriging to enforce continuity of HTE estimates across adjacent regions. By imposing continuity constraints only along the propensity score dimension, rather than the full covariate space, the proposed approach substantially reduces computational cost while avoiding discontinuities inherent in simple local approximations. The resulting method can be interpreted as a smoothing extension of stratification and provides an efficient approach to HTE estimation. The proposed method is demonstrated through simulation studies and a real data application.