Deep Learning of Continuous and Structured Policies for Aggregated Heterogeneous Treatment Effects

TL;DR

This paper introduces a deep learning framework for estimating and ranking heterogeneous treatment effects involving multiple structured treatment policies, including continuous and discrete variables, improving decision-making in complex treatment scenarios.

Contribution

It develops a novel methodology integrating multiple treatment factors into deep models, including a Neural-Augmented Naive Bayes layer for flexible treatment effect estimation and ranking.

Findings

Enhanced performance on public datasets

Effective modeling of continuous and structured treatment policies

Improved subject ranking based on aggregated HTE functions

Abstract

As estimation of Heterogeneous Treatment Effect (HTE) is increasingly adopted across a wide range of scientific and industrial applications, the treatment action space can naturally expand, from a binary treatment variable to a structured treatment policy. This policy may include several policy factors such as a continuous treatment intensity variable, or discrete treatment assignments. From first principles, we derive the formulation for incorporating multiple treatment policy variables into the functional forms of individual and average treatment effects. Building on this, we develop a methodology to directly rank subjects using aggregated HTE functions. In particular, we construct a Neural-Augmented Naive Bayes layer within a deep learning framework to incorporate an arbitrary number of factors that satisfies the Naive Bayes assumption. The factored layer is then applied with continuous treatment variables, treatment assignment, and direct ranking of aggregated treatment effect functions. Together, these algorithms build towards a generic framework for deep learning of heterogeneous treatment policies, and we show their power to improve performance with public datasets.