DeepBlip: Estimating Conditional Average Treatment Effects Over Time

TL;DR

DeepBlip introduces a neural network framework for structural nested mean models, enabling end-to-end training to estimate time-varying treatment effects with improved accuracy and interpretability in clinical data.

Contribution

It is the first neural framework for SNMMs that allows simultaneous learning of blip functions using a novel double optimization trick.

Findings

Achieves state-of-the-art performance on clinical datasets

Correctly adjusts for time-varying confounding

Provides unbiased estimates with robustness to model misspecification

Abstract

Structural nested mean models (SNMMs) are a principled approach to estimate the treatment effects over time. A particular strength of SNMMs is to break the joint effect of treatment sequences over time into localized, time-specific ``blip effects''. This decomposition promotes interpretability through the incremental effects and enables the efficient offline evaluation of optimal treatment policies without re-computation. However, neural frameworks for SNMMs are lacking, as their inherently sequential g-estimation scheme prevents end-to-end, gradient-based training. Here, we propose DeepBlip, the first neural framework for SNMMs, which overcomes this limitation with a novel double optimization trick to enable simultaneous learning of all blip functions. Our DeepBlip seamlessly integrates sequential neural networks like LSTMs or transformers to capture complex temporal dependencies. By design, our method correctly adjusts for time-varying confounding to produce unbiased estimates, and its Neyman-orthogonal loss function ensures robustness to nuisance model misspecification. Finally, we evaluate our DeepBlip across various clinical datasets, where it achieves state-of-the-art performance.