{"type":"ScholarlyArticle","title":"Pre-training of Graph Augmented Transformers for Medication Recommendation","authors":[{"name":"Junyuan Shang","slug":"junyuan-shang","given":"Junyuan","family":"Shang","affiliation":"Analytics Center of Excellence , IQVIA , Cambridge , MA , USA; Georgia Institute of Technology , Atlanta , GA , USA"},{"name":"Tengfei Ma","slug":"tengfei-ma","given":"Tengfei","family":"Ma","affiliation":"IBM Research AI , Yorktown Heights , NY , USA"},{"name":"Cao Xiao","slug":"cao-xiao","given":"Cao","family":"Xiao","affiliation":"Analytics Center of Excellence , IQVIA , Cambridge , MA , USA"},{"name":"Jimeng Sun","slug":"jimeng-sun","given":"Jimeng","family":"Sun","affiliation":"Georgia Institute of Technology , Atlanta , GA , USA"}],"abstract":"Medication recommendation is an important healthcare application.  It is commonly formulated as a temporal prediction task. Hence, most existing works only utilize longitudinal electronic health records (EHRs) from a small number of patients with multiple visits ignoring a large number of patients with a single visit (selection bias). Moreover, important hierarchical knowledge such as diagnosis hierarchy is not leveraged in the representation learning process. Despite the success of deep learning techniques in computational phenotyping, most previous approaches have two limitations: task-oriented representation and ignoring hierarchies of medical codes. \r\nTo address these challenges, we propose G-BERT,  a new model to combine the power of Graph Neural Networks (GNNs) and BERT (Bidirectional Encoder Representations from Transformers) for medical code representation and medication recommendation. We use  GNNs to represent the internal hierarchical structures of medical codes. Then we integrate the GNN representation into a transformer-based visit encoder and pre-train it on EHR data from patients only with a single visit. The pre-trained visit encoder and representation are then fine-tuned for downstream predictive tasks on longitudinal EHRs from patients with multiple visits. G-BERT is the first to bring the language model pre-training schema into the healthcare domain and it achieved state-of-the-art performance on the medication recommendation task.","doi":"10.24963/ijcai.2019/825","url":"https://ijcai.org/proceedings/2019/825","pdf":"https://ijcai.org/proceedings/2019/0825.pdf","markdown":"https://ijcai.org/proceedings/2019/825.md","bibtex":"https://ijcai.org/proceedings/2019/bibtex/825","volume":{"slug":"2019","title":"Proceedings of the Twenty-Eighth International Joint Conference on Artificial Intelligence","short_title":"IJCAI 2019","url":"https://ijcai.org/proceedings/2019/","year":2019,"location":"Macao","dates":"10-16 August 2019","editors":"Sarit Kraus","isbn":"978-0-9992411-4-1","publisher":"International Joint Conferences on Artificial Intelligence Organization","doi_prefix":"10.24963/ijcai.2019"},"section":"AI for Improving Human Well-being","track":"AI for Improving Human Well-being","number":825,"pages":{"first":5953,"last":5959},"keywords":["Special Track on AI for Improving Human-Well Being: Health applications (Special Track on AI and Human Wellbeing)"],"videos":[],"references":[{"position":1,"raw":"[Baytas et al., 2017] Inci M. Baytas, Cao Xiao, Xi Zhang, Fei Wang, Anil K. Jain, and Jiayu Zhou. Patient subtyp- ing via time-aware lstm networks. In Proceedings of the 23rd ACM SIGKDD International Conference on Knowl- edge Discovery and Data Mining, 2017.","title":"Patient subtyping via time-aware lstm networks","authors":"Inci M Baytas, Cao Xiao, Xi Zhang, Fei Wang, Anil K Jain, Jiayu Zhou","year":"2017","venue":"Proceedings of the 23rd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining","doi":null,"url":null,"internal_url":null},{"position":2,"raw":"[Cho et al., 2014] Kyunghyun Cho, Bart Van Merriënboer, Dzmitry Bahdanau, and Yoshua Bengio. On the proper- ties of neural machine translation: Encoder-decoder ap- proaches. arXiv preprint arXiv:1409.1259, 2014.","title":"On the properties of neural machine translation: Encoder-decoder approaches","authors":"Cho","year":"2014","venue":null,"doi":null,"url":null,"internal_url":null},{"position":3,"raw":"[Choi et al., 2016] Edward Choi, Mohammad Taha Ba- hadori, Jimeng Sun, Joshua Kulas, Andy Schuetz, and Walter Stewart. Retain: An interpretable predictive model for healthcare using reverse time attention mechanism. In Advances in Neural Information Processing Systems, pages 3504-3512, 2016.","title":"Retain: An interpretable predictive model for healthcare using reverse time attention mechanism","authors":"Edward Choi, Mohammad Taha Bahadori, Jimeng Sun, Joshua Kulas, Andy Schuetz, Walter Stewart","year":"2016","venue":"Advances in Neural Information Processing Systems","doi":null,"url":null,"internal_url":null},{"position":4,"raw":"[Choi et al., 2017] Edward Choi, Mohammad Taha Ba- hadori, Le Song, Walter F Stewart, and Jimeng Sun. Gram: Graph-based attention model for healthcare representation learning. In SIGKDD, 2017.","title":"Gram: Graph-based attention model for healthcare representation learning","authors":"Edward Choi, Mohammad Taha Bahadori, Le Song, Walter F Stewart, Jimeng Sun","year":"2017","venue":"SIGKDD","doi":null,"url":null,"internal_url":null},{"position":5,"raw":"[Choi et al., 2018] Edward Choi, Cao Xiao, Walter Stew- art, and Jimeng Sun. Mime: Multilevel medical embed- ding of electronic health records for predictive healthcare.","title":"Mime: Multilevel medical embedding of electronic health records for predictive healthcare","authors":"Edward Choi, Cao Xiao, Walter Stewart, Jimeng Sun","year":null,"venue":null,"doi":null,"url":null,"internal_url":null},{"position":6,"raw":"In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems 31, pages 4547- 4557.","title":"Advances in Neural Information Processing Systems","authors":"S Bengio, H Wallach, H Larochelle, K Grauman, N","year":null,"venue":null,"doi":null,"url":null,"internal_url":null},{"position":7,"raw":"Curran Associates, Inc., 2018.","title":null,"authors":"Curran Associates, Inc","year":"2018","venue":null,"doi":null,"url":null,"internal_url":null},{"position":8,"raw":"[Devlin et al., 2018] Jacob Devlin, Ming-Wei Chang, Ken- ton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understand- ing. arXiv preprint arXiv:1810.04805, 2018.","title":"Bert: Pre-training of deep bidirectional transformers for language understanding","authors":"Jacob Devlin, Ming-Wei Chang, Kenton Lee, Kristina Toutanova","year":"2018","venue":null,"doi":null,"url":null,"internal_url":null},{"position":9,"raw":"[Erhan et al., 2010] Dumitru Erhan, Yoshua Bengio, Aaron Courville, Pierre-Antoine Manzagol, Pascal Vincent, and Samy Bengio. Why does unsupervised pre-training help deep learning? Journal of Machine Learning Research, 11(Feb):625-660, 2010.","title":"Why does unsupervised pre-training help deep learning","authors":"Yoshua Dumitru Erhan, Aaron Bengio, Pierre-Antoine Courville, Pascal Manzagol, Samy Vincent, Bengio","year":"2010","venue":"Journal of Machine Learning Research","doi":null,"url":null,"internal_url":null},{"position":10,"raw":"[Gilmer et al., 2017] J. Gilmer, S.S. Schoenholz, P.F. Riley, O. Vinyals, and G.E. Dahl. Neural message passing for quantum chemistry. In ICML, 2017.","title":"Neural message passing for quantum chemistry","authors":"J Gilmer, S S Schoenholz, P F Riley, O Vinyals, G E Dahl","year":"2017","venue":"ICML","doi":null,"url":null,"internal_url":null},{"position":11,"raw":"[Hinton et al., 2006] Geoffrey E Hinton, Simon Osindero, and Yee-Whye Teh. A fast learning algorithm for deep belief nets. Neural computation, 18(7):1527-1554, 2006.","title":"A fast learning algorithm for deep belief nets","authors":"Geoffrey E Hinton, Simon Osindero, Yee-Whye Teh","year":"2006","venue":"Neural computation","doi":null,"url":null,"internal_url":null},{"position":12,"raw":"[Johnson et al., 2016] Alistair EW Johnson, Tom J Pollard, Lu Shen, H Lehman Li-wei, Mengling Feng, Mohammad Ghassemi, Benjamin Moody, Peter Szolovits, Leo An- thony Celi, and Roger G Mark. Mimic-iii, a freely ac- cessible critical care database. Scientific data, 3:160035, 2016.","title":"Mimic-iii, a freely accessible critical care database","authors":"E W Alistair, Tom J Johnson, Lu Pollard, H Shen, Mengling Lehman Li-Wei, Mohammad Feng, Benjamin Ghassemi, Peter Moody, Leo Szolovits, Roger G Anthony Celi, Mark","year":"2016","venue":"Scientific data","doi":null,"url":null,"internal_url":null},{"position":13,"raw":"[Kingma and Ba, 2014] Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. CoRR, abs/1412.6980, 2014.","title":"Adam: A method for stochastic optimization","authors":"P Diederik, Jimmy Kingma, Ba","year":"2014","venue":null,"doi":null,"url":null,"internal_url":null},{"position":14,"raw":"[Kipf and Welling, 2017] Thomas N Kipf and Max Welling. Semi-supervised classification with graph convolutional networks. In ICLR, 2017.","title":"Semi-supervised classification with graph convolutional networks","authors":"N Thomas, Max Kipf, Welling","year":"2017","venue":"ICLR","doi":null,"url":null,"internal_url":null},{"position":15,"raw":"[Lipton et al., 2015] Zachary C Lipton, David C Kale, Charles Elkan, and Randall Wetzel. Learning to diag- nose with lstm recurrent neural networks. arXiv preprint arXiv:1511.03677, 2015.","title":"Learning to diagnose with lstm recurrent neural networks","authors":"David C Zachary C Lipton, Charles Kale, Randall Elkan, Wetzel","year":"2015","venue":null,"doi":null,"url":null,"internal_url":null},{"position":16,"raw":"[Luaces et al., 2012] Oscar Luaces, Jorge Díez, José Barran- quero, Juan José del Coz, and Antonio Bahamonde. Binary relevance efficacy for multilabel classification. Progress in Artificial Intelligence, 1(4):303-313, 2012.","title":"Binary relevance efficacy for multilabel classification","authors":"Luaces","year":"2012","venue":"Progress in Artificial Intelligence","doi":null,"url":null,"internal_url":null},{"position":17,"raw":"[Ma et al., 2018] Tengfei Ma, Cao Xiao, and Fei Wang. Health-atm: A deep architecture for multifaceted patient health record representation and risk prediction. In Pro- ceedings of the 2018 SIAM International Conference on Data Mining, pages 261-269. SIAM, 2018.","title":"Health-atm: A deep architecture for multifaceted patient health record representation and risk prediction","authors":"Ma","year":"2018","venue":"Proceedings of the 2018 SIAM International Conference on Data Mining","doi":null,"url":null,"internal_url":null},{"position":18,"raw":"[Paszke et al., 2017] Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary De- Vito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch. 2017.","title":"Automatic differentiation in pytorch","authors":"Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary De-Vito, Zeming Lin, Alban Desmaison, Luca Antiga, Adam Lerer","year":"2017","venue":null,"doi":null,"url":null,"internal_url":null},{"position":19,"raw":"[Peters et al., 2018] Matthew Peters, Mark Neumann, Mohit Iyyer, Matt Gardner, Christopher Clark, Kenton Lee, and Luke Zettlemoyer. Deep contextualized word representa- tions. In NAACL, volume 1, pages 2227-2237, 2018.","title":"Deep contextualized word representations","authors":"Matthew Peters, Mark Neumann, Mohit Iyyer, Matt Gardner, Christopher Clark, Kenton Lee, Luke Zettlemoyer","year":"2018","venue":"NAACL","doi":null,"url":null,"internal_url":null},{"position":20,"raw":"[Radford et al., 2018] Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever. Improving language understanding by generative pre-training. 2018.","title":"Improving language understanding by generative pre-training","authors":"Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever","year":"2018","venue":null,"doi":null,"url":null,"internal_url":null},{"position":21,"raw":"[Ramachandran et al., 2016] Prajit Ramachandran, Peter J Liu, and Quoc V Le. Unsupervised pretraining for sequence to sequence learning. arXiv preprint arXiv:1611.02683, 2016.","title":"Unsupervised pretraining for sequence to sequence learning","authors":"Ramachandran","year":"2016","venue":null,"doi":null,"url":null,"internal_url":null},{"position":22,"raw":"[Shang et al., 2019] Junyuan Shang, Cao Xiao, Tengfei Ma, Hongyan Li, and Jimeng Sun. Gamenet: Graph augmented memory networks for recommending medication combi- nation. AAAI, 2019.","title":"Gamenet: Graph augmented memory networks for recommending medication combination","authors":"Junyuan Shang, Cao Xiao, Tengfei Ma, Hongyan Li, Jimeng Sun","year":"2019","venue":null,"doi":null,"url":null,"internal_url":null},{"position":23,"raw":"[Vaswani et al., 2017] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Sys- tems, pages 5998-6008, 2017.","title":"Attention is all you need","authors":"Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, Illia Polosukhin","year":"2017","venue":"Advances in Neural Information Processing Systems","doi":null,"url":null,"internal_url":null},{"position":24,"raw":"[Velickovic et al., 2017] Petar Velickovic, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Lio, and Yoshua Bengio. Graph attention networks. arXiv preprint arXiv:1710.10903, 1(2), 2017.","title":"Graph attention networks","authors":"Petar Velickovic, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Lio, Yoshua Bengio","year":"2017","venue":null,"doi":null,"url":null,"internal_url":null},{"position":25,"raw":"[Xiao et al., 2018a] Cao Xiao, Edward Choi, and Jimeng Sun. Opportunities and challenges in developing deep learning models using electronic health records data: a systematic review. Journal of the American Medical In- formatics Association, 2018.","title":"Opportunities and challenges in developing deep learning models using electronic health records data: a systematic review","authors":"Xiao","year":"2018","venue":"Journal of the American Medical Informatics Association","doi":null,"url":null,"internal_url":null},{"position":26,"raw":"[Xiao et al., 2018b] Cao Xiao, Tengfei Ma, Adji B. Dieng, David M. Blei, and Fei Wang. Readmission prediction via deep contextual embedding of clinical concepts. PLOS ONE, 13(4):1-15, 04 2018.","title":"Readmission prediction via deep contextual embedding of clinical concepts","authors":"Xiao","year":"2018","venue":"PLOS ONE","doi":null,"url":null,"internal_url":null},{"position":27,"raw":"[Zhang et al., 2017] Yutao Zhang, Robert Chen, Jie Tang, Walter F Stewart, and Jimeng Sun. Leap: Learning to pre- scribe effective and safe treatment combinations for mul- timorbidity. In SIGKDD, pages 1315-1324, 2017.","title":"Leap: Learning to prescribe effective and safe treatment combinations for multimorbidity","authors":"Zhang","year":"2017","venue":"SIGKDD","doi":null,"url":null,"internal_url":null}],"license_note":"Copyright International Joint Conferences on Artificial Intelligence."}