Clinical Decision Transformer: Intended Treatment Recommendation through Goal Prompting

Abstract

With recent achievements in tasks requiring context awareness, foundation models have been adopted to treat large-scale data from electronic health record (EHR) systems. However, previous clinical recommender systems based on foundation models have a limited purpose of imitating clinicians' behavior and do not directly consider a problem of missing values. In this paper, we propose Clinical Decision Transformer (CDT), a recommender system that generates a sequence of medications to reach a desired range of clinical states given as goal prompts. For this, we conducted goal-conditioned sequencing, which generated a subsequence of treatment history with prepended future goal state, and trained the CDT to model sequential medications required to reach that goal state. For contextual embedding over intra-admission and inter-admissions, we adopted a GPT-based architecture with an admission-wise attention mask and column embedding. In an experiment, we extracted a diabetes dataset from an EHR system, which contained treatment histories of 4788 patients. We observed that the CDT achieved the intended treatment effect according to goal prompt ranges (e.g., NormalA1c, LowerA1c, and HigherA1c), contrary to the case with behavior cloning. To the best of our knowledge, this is the first study to explore clinical recommendations from the perspective of goal prompting. See https://clinical-decision-transformer.github.io for code and additional information.