Personalized Assignment to One of Many Treatment Arms via Regularized and Clustered Joint Assignment Forests

TL;DR

This paper introduces regularized and clustered forest-based methods for personalized treatment assignment in randomized trials, improving accuracy by pooling information across multiple treatment arms and reducing variance.

Contribution

It proposes novel forest algorithms that incorporate regularization and clustering to enhance personalized treatment assignment across many arms.

Findings

Regularized forests outperform separate arm-wise predictions.

Clustering treatment arms improves assignment accuracy.

Methods achieve significant utility gains in simulations.

Abstract

We consider learning personalized assignments to one of many treatment arms from a randomized controlled trial. Standard methods that estimate heterogeneous treatment effects separately for each arm may perform poorly in this case due to excess variance. We instead propose methods that pool information across treatment arms: First, we consider a regularized forest-based assignment algorithm based on greedy recursive partitioning that shrinks effect estimates across arms. Second, we augment our algorithm by a clustering scheme that combines treatment arms with consistently similar outcomes. In a simulation study, we compare the performance of these approaches to predicting arm-wise outcomes separately, and document gains of directly optimizing the treatment assignment with regularization and clustering. In a theoretical model, we illustrate how a high number of treatment arms makes finding the best arm hard, while we can achieve sizable utility gains from personalization by regularized optimization.