Active Feature Acquisition for Personalised Treatment Assignment

Julianna Piskorz, Nicolás Astorga, Jeroen Berrevoets, Mihaela van der Schaar
Proceedings of The 28th International Conference on Artificial Intelligence and Statistics, PMLR 258:4330-4338, 2025.

Abstract

Making treatment effect estimation actionable for personalized decision-making requires overcoming the costs and delays of acquiring necessary features. While many machine learning models estimate Conditional Average Treatment Effects (CATE), they mostly assume that \emph{all} relevant features are readily available at prediction time – a scenario that is rarely realistic. In practice, acquiring features, such as medical tests, can be both expensive and time-consuming, highlighting the need for strategies that select the most informative features for each individual, enhancing decision accuracy while controlling costs. Existing active feature acquisition (AFA) methods, developed for supervised learning, fail to address the unique challenges of CATE, such as confounding, overlap, and the structural similarities of potential outcomes under different treatments. To tackle these challenges, we propose specialised feature acquisition metrics and estimation strategies tailored to the CATE setting. We demonstrate the effectiveness of our methods through experiments on synthetic datasets designed to reflect common biases and data issues. In doing so, this work aims to bridge the gap between cutting-edge CATE estimation techniques and their practical, cost-efficient application in personalised treatment assignment.

Cite this Paper

BibTeX
@InProceedings{pmlr-v258-piskorz25a, title = {Active Feature Acquisition for Personalised Treatment Assignment}, author = {Piskorz, Julianna and Astorga, Nicol{\'a}s and Berrevoets, Jeroen and van der Schaar, Mihaela}, booktitle = {Proceedings of The 28th International Conference on Artificial Intelligence and Statistics}, pages = {4330--4338}, year = {2025}, editor = {Li, Yingzhen and Mandt, Stephan and Agrawal, Shipra and Khan, Emtiyaz}, volume = {258}, series = {Proceedings of Machine Learning Research}, month = {03--05 May}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v258/main/assets/piskorz25a/piskorz25a.pdf}, url = {https://proceedings.mlr.press/v258/piskorz25a.html}, abstract = {Making treatment effect estimation actionable for personalized decision-making requires overcoming the costs and delays of acquiring necessary features. While many machine learning models estimate Conditional Average Treatment Effects (CATE), they mostly assume that \emph{all} relevant features are readily available at prediction time – a scenario that is rarely realistic. In practice, acquiring features, such as medical tests, can be both expensive and time-consuming, highlighting the need for strategies that select the most informative features for each individual, enhancing decision accuracy while controlling costs. Existing active feature acquisition (AFA) methods, developed for supervised learning, fail to address the unique challenges of CATE, such as confounding, overlap, and the structural similarities of potential outcomes under different treatments. To tackle these challenges, we propose specialised feature acquisition metrics and estimation strategies tailored to the CATE setting. We demonstrate the effectiveness of our methods through experiments on synthetic datasets designed to reflect common biases and data issues. In doing so, this work aims to bridge the gap between cutting-edge CATE estimation techniques and their practical, cost-efficient application in personalised treatment assignment.} }
Endnote
%0 Conference Paper %T Active Feature Acquisition for Personalised Treatment Assignment %A Julianna Piskorz %A Nicolás Astorga %A Jeroen Berrevoets %A Mihaela van der Schaar %B Proceedings of The 28th International Conference on Artificial Intelligence and Statistics %C Proceedings of Machine Learning Research %D 2025 %E Yingzhen Li %E Stephan Mandt %E Shipra Agrawal %E Emtiyaz Khan %F pmlr-v258-piskorz25a %I PMLR %P 4330--4338 %U https://proceedings.mlr.press/v258/piskorz25a.html %V 258 %X Making treatment effect estimation actionable for personalized decision-making requires overcoming the costs and delays of acquiring necessary features. While many machine learning models estimate Conditional Average Treatment Effects (CATE), they mostly assume that \emph{all} relevant features are readily available at prediction time – a scenario that is rarely realistic. In practice, acquiring features, such as medical tests, can be both expensive and time-consuming, highlighting the need for strategies that select the most informative features for each individual, enhancing decision accuracy while controlling costs. Existing active feature acquisition (AFA) methods, developed for supervised learning, fail to address the unique challenges of CATE, such as confounding, overlap, and the structural similarities of potential outcomes under different treatments. To tackle these challenges, we propose specialised feature acquisition metrics and estimation strategies tailored to the CATE setting. We demonstrate the effectiveness of our methods through experiments on synthetic datasets designed to reflect common biases and data issues. In doing so, this work aims to bridge the gap between cutting-edge CATE estimation techniques and their practical, cost-efficient application in personalised treatment assignment.
APA
Piskorz, J., Astorga, N., Berrevoets, J. & van der Schaar, M.. (2025). Active Feature Acquisition for Personalised Treatment Assignment. Proceedings of The 28th International Conference on Artificial Intelligence and Statistics, in Proceedings of Machine Learning Research 258:4330-4338 Available from https://proceedings.mlr.press/v258/piskorz25a.html.

Related Material