[Paper Review] Learning Optimal Prescriptive Trees from Observational Data

Nathanael Jo, Sina Aghaei+2 more|arXiv (Cornell University)|Aug 31, 2021

Advanced Causal Inference Techniques39 references4 citations

TL;DR

This paper proposes a mixed-integer optimization (MIO)-based method to learn optimal, interpretable prescriptive trees from observational data without requiring randomization or strong assumptions. It achieves asymptotic exactness in out-of-sample performance and uniquely supports domain constraints like budget and fairness, outperforming state-of-the-art methods by up to threefold in finite samples.

ABSTRACT

We consider the problem of learning an optimal prescriptive tree (i.e., an interpretable treatment assignment policy in the form of a binary tree) of moderate depth, from observational data. This problem arises in numerous socially important domains such as public health and personalized medicine, where interpretable and data-driven interventions are sought based on data gathered in deployment -- through passive collection of data -- rather than from randomized trials. We propose a method for learning optimal prescriptive trees using mixed-integer optimization (MIO) technology. We show that under mild conditions our method is asymptotically exact in the sense that it converges to an optimal out-of-sample treatment assignment policy as the number of historical data samples tends to infinity. Contrary to existing literature, our approach: 1) does not require data to be randomized, 2) does not impose stringent assumptions on the learned trees, and 3) has the ability to model domain specific constraints. Through extensive computational experiments, we demonstrate that our asymptotic guarantees translate to significant performance improvements in finite samples, as well as showcase our uniquely flexible modeling power by incorporating budget and fairness constraints.

Motivation & Objective

• To develop a method for learning interpretable, optimal treatment assignment policies from observational data, where randomized trials are unethical or impractical.
• To address the limitations of existing prescriptive tree methods that require randomized data, impose deep tree structures, or lack support for domain-specific constraints.
• To ensure the learned policy is asymptotically optimal in out-of-sample performance under mild regularity conditions.
• To enable practical deployment in high-stakes domains such as healthcare and social services by incorporating fairness and budget constraints.

Proposed method

• Formulates the prescriptive tree learning problem as a mixed-integer optimization (MIO) problem to jointly select tree structure and treatment assignments.
• Uses doubly robust estimators to estimate potential outcomes from observational data, improving robustness when either the outcome or propensity score model is misspecified.
• Employs a data-driven objective function based on the doubly robust estimator to maximize expected outcomes while ensuring asymptotic consistency.
• Incorporates constraints into the MIO framework to model domain-specific requirements such as budget limits (e.g., only 20% of population can receive treatment) and fairness (e.g., outcome disparity ≤ δ between demographic groups).
• Uses a decomposition approach to handle large-scale datasets by solving subproblems iteratively, improving computational tractability.
• Employs a sample-based approximation of the expected outcome to enable optimization over finite historical data while maintaining theoretical guarantees.

Experimental results

Research questions

• RQ1 Can a prescriptive tree be learned from observational data without requiring randomized treatment assignment, while still achieving asymptotic optimality in out-of-sample performance?
• RQ2 How does the proposed MIO-based method compare to existing state-of-the-art methods in terms of out-of-sample performance across different data-generating mechanisms?
• RQ3 To what extent can the method incorporate practical constraints such as budget limits and fairness requirements in real-world policy settings?
• RQ4 Does the use of doubly robust estimators improve performance stability when either the outcome or propensity score model is misspecified?
• RQ5 How well do the theoretical asymptotic guarantees translate into finite-sample performance improvements?

Key findings

• The proposed method achieves up to threefold performance improvement over state-of-the-art methods in out-of-sample treatment assignment accuracy, particularly in non-randomized settings.
• The method maintains high performance across diverse experimental designs, including conditionally randomized and non-randomized data, with the doubly robust estimator showing robustness when one of the models (outcome or propensity) is misspecified.
• When fairness constraints are imposed, reducing the disparity threshold δ to 0.01 leads to a median outcome disparity of -0.02 in the randomized setting and full parity in the conditionally randomized setting with r=0.11.
• The method successfully enforces budget constraints, allowing control over the proportion of the population receiving treatment, with no degradation in performance.
• The DM estimator was found to produce outcome estimates closest to the true expected outcomes, though the DR estimator showed better robustness in practice due to reduced bias from extreme propensity weights.
• Empirical results confirm that the asymptotic theoretical guarantees translate into strong finite-sample performance, especially under realistic observational data conditions.