Learning to Ask Medical Questions using Reinforcement Learning

Uri Shaham, Tom Zahavy, Cesar Caraballo, Shiwani Mahajan, Daisy Massey, Harlan Krumholz
Proceedings of the 5th Machine Learning for Healthcare Conference, PMLR 126:2-26, 2020.

Abstract

We propose a novel reinforcement learning-based approach for adaptive and iterative feature selection. Given a masked vector of input features, a reinforcement learning agent iteratively selects certain features to be unmasked, and uses them to predict an outcome when it is sufficiently confident. The algorithm makes use of a novel environment setting, corresponding to a non-stationary Markov Decision Process. A key component of our approach is a guesser network, trained to predict the outcome from the selected features and parametrizing the reward function. Applying our method to a national survey dataset, we show that it not only outperforms strong baselines when requiring the prediction to be made based on a small number of input features, but is also highly more interpretable. Our code is publicly available at https://github.com/ushaham/adaptiveFS.

Cite this Paper

BibTeX
@InProceedings{pmlr-v126-shaham20a, title = {Learning to Ask Medical Questions using Reinforcement Learning}, author = {Shaham, Uri and Zahavy, Tom and Caraballo, Cesar and Mahajan, Shiwani and Massey, Daisy and Krumholz, Harlan}, booktitle = {Proceedings of the 5th Machine Learning for Healthcare Conference}, pages = {2--26}, year = {2020}, editor = {Doshi-Velez, Finale and Fackler, Jim and Jung, Ken and Kale, David and Ranganath, Rajesh and Wallace, Byron and Wiens, Jenna}, volume = {126}, series = {Proceedings of Machine Learning Research}, month = {07--08 Aug}, publisher = {PMLR}, pdf = {http://proceedings.mlr.press/v126/shaham20a/shaham20a.pdf}, url = {https://proceedings.mlr.press/v126/shaham20a.html}, abstract = {We propose a novel reinforcement learning-based approach for adaptive and iterative feature selection. Given a masked vector of input features, a reinforcement learning agent iteratively selects certain features to be unmasked, and uses them to predict an outcome when it is sufficiently confident. The algorithm makes use of a novel environment setting, corresponding to a non-stationary Markov Decision Process. A key component of our approach is a guesser network, trained to predict the outcome from the selected features and parametrizing the reward function. Applying our method to a national survey dataset, we show that it not only outperforms strong baselines when requiring the prediction to be made based on a small number of input features, but is also highly more interpretable. Our code is publicly available at https://github.com/ushaham/adaptiveFS.} }
Endnote
%0 Conference Paper %T Learning to Ask Medical Questions using Reinforcement Learning %A Uri Shaham %A Tom Zahavy %A Cesar Caraballo %A Shiwani Mahajan %A Daisy Massey %A Harlan Krumholz %B Proceedings of the 5th Machine Learning for Healthcare Conference %C Proceedings of Machine Learning Research %D 2020 %E Finale Doshi-Velez %E Jim Fackler %E Ken Jung %E David Kale %E Rajesh Ranganath %E Byron Wallace %E Jenna Wiens %F pmlr-v126-shaham20a %I PMLR %P 2--26 %U https://proceedings.mlr.press/v126/shaham20a.html %V 126 %X We propose a novel reinforcement learning-based approach for adaptive and iterative feature selection. Given a masked vector of input features, a reinforcement learning agent iteratively selects certain features to be unmasked, and uses them to predict an outcome when it is sufficiently confident. The algorithm makes use of a novel environment setting, corresponding to a non-stationary Markov Decision Process. A key component of our approach is a guesser network, trained to predict the outcome from the selected features and parametrizing the reward function. Applying our method to a national survey dataset, we show that it not only outperforms strong baselines when requiring the prediction to be made based on a small number of input features, but is also highly more interpretable. Our code is publicly available at https://github.com/ushaham/adaptiveFS.
APA
Shaham, U., Zahavy, T., Caraballo, C., Mahajan, S., Massey, D. & Krumholz, H.. (2020). Learning to Ask Medical Questions using Reinforcement Learning. Proceedings of the 5th Machine Learning for Healthcare Conference, in Proceedings of Machine Learning Research 126:2-26 Available from https://proceedings.mlr.press/v126/shaham20a.html.

Related Material