TERRA: A Transformer-Enabled Recursive R-learner for Longitudinal Heterogeneous Treatment Effect Estimation

TL;DR

TERRA introduces a novel transformer-based recursive approach for estimating heterogeneous treatment effects over time, effectively capturing complex temporal dependencies and addressing biases in longitudinal data.

Contribution

It combines transformer encoding of treatment histories with recursive residual learning to improve longitudinal HTE estimation beyond existing methods.

Findings

Outperforms baseline models in accuracy and stability

Effectively models long-range temporal dependencies

Addresses post-treatment bias with recursive residual learning

Abstract

Accurately estimating heterogeneous treatment effects (HTE) in longitudinal settings is essential for personalized decision-making across healthcare, public policy, education, and digital marketing. However, time-varying interventions introduce many unique challenges, such as carryover effects, time-varying heterogeneity, and post-treatment bias, which are not addressed by standard HTE methods. To address these challenges, we introduce TERRA (Transformer-Enabled Recursive R-learner), which facilitates longitudinal HTE estimation with flexible temporal modeling and learning. TERRA has two components. First, we use a Transformer architecture to encode full treatment-feature histories, enabling the representation of long-range temporal dependencies and carryover effects, hence capturing individual- and time-specific treatment effect variation more comprehensively. Second, we develop a recursive residual-learning formulation that generalizes the classical structural nested mean models (SNMMs) beyond parametric specifications, addressing post-treatment bias while reducing reliance on functional assumptions. In simulations and data applications, TERRA consistently outperforms strong baselines in HTE estimation in both accuracy and stability, highlighting the value of combining principled causal structure with high-capacity sequence models for longitudinal HTE.