[Paper Review] Recursive Partitioning for Personalization using Observational Data

Nathan Kallus|arXiv (Cornell University)|Aug 31, 2016

Statistical Methods and Inference20 citations

TL;DR

This paper proposes a novel recursive partitioning approach for personalized treatment selection using observational data, framing personalization as a single learning task rather than m separate regressions. It introduces a new impurity measure for treatment regimes, develops a personalized tree (PT), forest (PF), and optimal partitioning (OPT) via mixed-integer programming, and demonstrates superior performance and interpretability in personalized medicine and job training applications, achieving up to 47% of the potential benefit from perfect personalization.

ABSTRACT

We study the problem of learning to choose from m discrete treatment options (e.g., news item or medical drug) the one with best causal effect for a particular instance (e.g., user or patient) where the training data consists of passive observations of covariates, treatment, and the outcome of the treatment. The standard approach to this problem is regress and compare: split the training data by treatment, fit a regression model in each split, and, for a new instance, predict all m outcomes and pick the best. By reformulating the problem as a single learning task rather than m separate ones, we propose a new approach based on recursively partitioning the data into regimes where different treatments are optimal. We extend this approach to an optimal partitioning approach that finds a globally optimal partition, achieving a compact, interpretable, and impactful personalization model. We develop new tools for validating and evaluating personalization models on observational data and use these to demonstrate the power of our novel approaches in a personalized medicine and a job training application.

Motivation & Objective

• To address the challenge of learning optimal treatment assignments from observational data where randomized trials are impractical or unethical.
• To overcome limitations of the standard 'regress and compare' approach, which treats each treatment as a separate regression task and fails to account for inter-treatment comparisons in observational settings.
• To develop a unified learning framework that directly optimizes for personalization by recursively partitioning the covariate space into regimes where specific treatments are optimal.
• To create interpretable, compact, and globally optimal personalization models using mixed-integer programming and ensemble methods.
• To introduce new validation techniques, including submatching and coefficient of personalization, to fairly evaluate and tune personalization models on observational data.

Proposed method

• Reformulates personalization as a single learning task by defining a new impurity measure that quantifies the expected outcome difference within a partition when assigning the best treatment.
• Develops the Personalization Tree (PT), a greedy recursive partitioning algorithm that minimizes the sum of within-partition personalization impurities to identify optimal treatment regimes.
• Introduces the Personalization Forest (PF), an ensemble of PTs trained on bootstrap samples to improve robustness and generalization.
• Proposes the Optimal Personalization Tree (OPT), which uses a mixed-integer program (MIP) to globally optimize the partitioning of the covariate space, ensuring the best possible model under the constraints.
• Employs a novel submatching technique to create matched test sets from observational data, enabling valid out-of-sample evaluation of personalization models.
• Defines and uses the coefficient of personalization (P1, P2) to quantify model performance as a percentage of the maximum possible benefit from perfect personalization.

Experimental results

Research questions

• RQ1 Can a recursive partitioning framework outperform standard 'regress and compare' methods in personalization from observational data?
• RQ2 How can a unified learning approach that jointly optimizes treatment assignment across all options improve model performance and interpretability compared to separate regressions per treatment?
• RQ3 To what extent can a globally optimal partitioning approach via mixed-integer programming yield better personalization models than greedy or ensemble alternatives?
• RQ4 How can valid evaluation and validation of personalization models be achieved on observational data without experimental controls?
• RQ5 What is the real-world impact of personalization, and how can it be quantified meaningfully using a new metric like the coefficient of personalization?

Key findings

• The Optimal Personalization Tree (OPT) achieved the best out-of-sample risk of 0.356 in the warfarin dosing application, corresponding to 47% of the maximum possible benefit from perfect personalization.
• In the job training application, the Personalization Forest (PF) achieved an average net income of $4,200.80, representing 28% of the benefit from perfect personalization relative to standard of care.
• OPT outperformed other methods at small sample sizes (n = 100) due to its ability to find simple, interpretable trees (depth 2–3), while performance declined at large n due to MIP solver limitations.
• The proposed submatching technique enabled valid evaluation by creating matched test sets with minimal covariate imbalance, ensuring reliable out-of-sample assessment.
• The coefficient of personalization (P1, P2) successfully quantified model performance as a percentage of the theoretical maximum benefit, providing a clear and interpretable metric for real-world impact.
• The OPT and PF methods demonstrated superior interpretability and compactness compared to baseline methods like CART and causal forests, especially in high-dimensional settings.