NOFLITE: Learning to Predict Individual Treatment Effect Distributions
Abstract: Estimating the effect of a treatment on an individual's outcome of interest is an important challenge in various fields, such as healthcare, economics, marketing, and education. Previous work in machine learning has focused on estimating the expected value of the treatment effect. However, effective personalized decision-making requires more than just the treatment expected effect; it requires knowing the entire treatment effect distribution. Knowing this distribution allows analyzing the treatment's expected utility or quantifying the uncertainty regarding a treatment's effect. This information is essential for prescribing optimal treatments. The ability of a model to predict accurate individual treatment effect distributions is captured by its likelihood. In light of this, we propose a novel neural architecture, NOFLITE, that uses normalizing flows to directly optimize this likelihood, while simultaneously learning flexible estimates of the individual treatment effect distribution. Experiments on various semi-synthetic data sets show that NOFLITE outperforms existing methods in terms of loglikelihood. Moreover, we illustrate how the predicted distributions can enable an in-depth analysis of the treatment effect and more accurate decision-making.
Submission Length: Regular submission (no more than 12 pages of main content)
Code: https://github.com/toonvds/NOFLITE
Assigned Action Editor: ~Fredrik_Daniel_Johansson1
License: Creative Commons Attribution 4.0 International (CC BY 4.0)
Submission Number: 1388
Loading