PaperlayerView exact revision on arXiv

Source-linked AI summary

Clinical Decision Transformer: Intended Treatment Recommendation through Goal Prompting

Seunghyun Lee, Da Young Lee, Sujeong Im, Nan Hee Kim, Sung-Min Park

arXiv:2302.00612v1cs.AI

TL;DR

EHR-based clinical recommendations must handle informative missing values and heterogeneous patient histories, while prior recommender systems are described as imitating recorded behavior. CDT conditions medication sequences on a desired clinical goal and patient history, and the paper reports that it can recommend medications in the intended direction of that goal.

  • Problem

    Clinical recommender systems based on foundation models are described as imitating clinician behavior, while EHR data also contain informative missing values and heterogeneous patient histories that require contextualization.

    Source: §1 · pp. 1–2

  • Method

    CDT uses goal-conditioned sequencing and a GPT-based model to generate medications from a desired outcome goal and patient history, with embeddings and attention designed for missing tabular data.

    Source: §1 · p. 3, §3 · p. 4, §4.1 · p. 4

  • Results

    CDT is reported to recommend medications in intended directions according to preassigned goals, unlike behavior cloning of recorded prescriptions.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The paper presents goal prompting as an approach to clinical medication recommendation that targets intended outcomes rather than reproducing recorded prescriptions.

    Source: §1 · pp. 2–3

    The treatment-effect comparisons cover only the very next admission, and the evaluation model was not designed for autoregressive inference.

    Source: §5.3 · p. 6

Abstract

from arXiv · show

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.

1. Introduction

EHR recommendations must account for informative missing laboratory values and heterogeneous patient histories. CDT uses goal prompts to recommend medications toward intended clinical outcomes, rather than merely imitating recorded prescriptions.

1. Introduction

  • Missing laboratory values can be informative, and heterogeneous treatment histories make contextualization important for individual-oriented use of EHR data.The paper notes that missingness can correlate with mortality and diagnosis categories, and that there is no one-size-fits-all treatment.Source: §1 · p. 1
  • CDT recommends medications toward a preassigned clinical goal, unlike behavior-cloning systems that imitate recorded prescriptions.The model conditions recommendations on the goal and an individual's EHR history.Source: §1 · p. 2
  • CDT applies column embeddings, [MISS] tokens, and admission-wise attention masks to represent tabular EHR data with missing values.Source: §1 · p. 3
  • The authors evaluate recommended-medication treatment effects using an auxiliary model for counterfactual estimation.Source: §1 · p. 3

2. Related Work

Related work frames clinical recommendation as sequence modeling while highlighting the challenge of missing clinical time-series data and evaluating treatments different from observed prescriptions.

Decision-making via sequence modeling

  • Decision Transformer and Trajectory Transformer cast decision-making as sequence modeling, with Trajectory Transformer also generating goal-conditioned trajectories.Source: §2 · p. 3

Multivariate-time series with missing values

  • Omitting or separately imputing missing data can introduce bias or leave missing context unmodeled; this study instead uses column embeddings and missing-value tokens for end-to-end training.Source: §2 · p. 3

Counterfactual estimation using deep learning

  • The paper uses a Counterfactual Recurrent Network encoder to estimate outcomes for CDT-recommended medications that differ from factual prescriptions.The related work describes counterfactual estimation as essential for evaluating treatment effects using observed data.Source: §2 · p. 3

3. Problem Formulation

The problem formulation represents patient records as sequences of clinical states, medication sets, and treatment outcomes across admissions. It constructs goal-conditioned subsequences using low outcomes as targets for medication recommendations.

3. Problem Formulation

  • Each patient record comprises clinical states and medication histories across hospital admissions, with prescribable medications represented as sets.The formulation includes baseline covariates and laboratory-test columns in the tabular clinical states.Source: §3 · p. 3
  • Treatment outcome Y_t is a clinical variable such as A1c, assumed to be better at relatively lower values except for extremely low values.Source: §3 · p. 3
  • For lower-valued goals, patient sequences are divided at outcome minima, and the minima are prepended to the corresponding partial histories as goals.Missing outcomes are not considered target outcomes.Source: §3 · pp. 3–4, Figure 2 · p. 4
  • The resulting training sequences support medication estimation conditioned on the assigned outcome goal, prior sequence, and current states.Source: §3 · p. 4

4. Clinical Decision Transformer

CDT uses a GPT-based sequence model to recommend medications conditioned on desired treatment outcomes and patient histories. Its design embeds tabular values and missingness, then contextualizes state tokens within and across admissions.

4.1. Model architecture

  • CDT uses causal GPT sequence modeling but conditions recommendations on a desired future treatment outcome, requiring one goal prompt rather than a return-to-go objective.The paper motivates this choice by the high missing rate of EHR covariates, which makes return-to-go difficult to calculate.Source: §4.1 · p. 4

4.2. Column embedding

  • Column embeddings combine each clinical variable's identifier and value; separate trainable [MISS]k tokens represent missing laboratory-test values.Baseline categorical variables and continuous laboratory values use different embedding procedures.Source: §4.2 · p. 4
  • Medication sets and discretized outcome goals are embedded, and identical admission embeddings are added to state and medication tokens from the same admission.Source: §4.2 · pp. 4–5

4.3. Admission-wise attention mask

  • The admission-wise attention mask extends state-token attention through the final state token of the same admission, contextualizing states within and across admissions.Source: §4.3 · p. 5, Figure 3 · p. 5

5. Experiments

The experiments evaluate CDT on a diabetes EHR dataset, testing goal-conditioned treatment effects and recommendation behavior against behavior cloning. They also examine contextual attention and ablate two architectural components.

5.1. Diabetes EHR dataset

  • The diabetes EHR cohort included 4,788 patients and 9,406 goal-conditioned subsequences; laboratory columns had a 44.39% average missing rate.The median laboratory missing rate was 22.37%, and A1c was used as the goal-prompt outcome.Source: §5.1 · p. 5

5.2. Training

  • CDT used GPT-2 to predict clinician-prescribed medication sets from a prepended goal, prior admissions, and current states, with an 8:1:1 train-validation-test split.The comparison behavior-cloning model used an LSTM trained to imitate factual prescriptions.Source: §5.2 · p. 5

5.3. Recommendation through a goal prompting

  • Goal prompts shifted estimated A1c toward the normal range or down/up relative to factual prescriptions, whereas behavior cloning remained similar to factual outcomes.The authors report that prompt-based recommendations had a relatively significant treatment effect according to the conditioned goal, unlike behavior cloning’s marginal improvement.Source: Figure 4 · p. 6, §5.3 · p. 6

5.4. Recommendation analysis

  • Approximately 33% cold-start recall rose to 94% with longer input histories, while FPR fell from around 4% to 1%; CDT had higher recall and lower FPR than BC on longer inputs.At cold start, BC had higher recall but also higher FPR than CDT.Source: §5.4 · pp. 6–7, Figure 5 · p. 7
  • With longer histories, CDT showed state-focused attention to the latest admission and medication-focused attention to the last admission’s prescriptions; at cold start, age, hypertension, and A1c informed state context.Missing values were contextualized using other diagnoses and laboratory information from the same admission.Source: Figure 6 · p. 7, §5.4 · p. 7
  • After GPT contextual embedding, admission representations clustered by factual prescriptions, whereas input admissions were entangled.The admissions were projected into two dimensions using t-SNE.Source: §5.4 · pp. 7–8, Figure 7 · p. 8

5.5. Ablation study

  • Removing column embedding reduced recommendation performance at every input length, while the admission-wise attention mask helped only in the cold-start case.The authors attribute the column-embedding result to difficulty representing tabular data without it.Source: Table 3 · p. 8, §5.5 · p. 8

6. Conclusion

The conclusion presents CDT as a goal-prompted medication recommender trained on diabetes EHR histories. It reports intended treatment effects across A1c goal ranges and identifies scalable EHR use as a future research direction.

6. Conclusion

  • CDT recommends medication sequences toward prompted A1c goals using treatment histories and current clinical states.The study extracted records from 4,788 patients with diabetes and used goal-conditioned sequencing to create goal-sequence pairs.Source: §6 · p. 8
  • The GPT-based model combines admission-wise attention masking and column embeddings to contextualize inter- and intra-admission EHR information.The paper reports that CDT achieved intended treatment effects for NormalA1c, LowerA1c, and HigherA1c prompts.Source: §6 · p. 8
  • The authors identify scalable use of large-scale, disease-wide EHR systems through goal prompting as a future research direction.Source: §6 · p. 8

Fundings

The work acknowledges support from Korean government-funded research programs and a National IT Industry Promotion Agency grant.

Fundings

  • Support included NRF grants funded by the Korean government and Ministry of Education, plus a National IT Industry Promotion Agency grant.The listed support also includes the Bio & Medical Technology Development Program.Source: Fundings · pp. 8–9

A. Evaluation Model Selection

The evaluation-model comparison selected CRN-E using mean square error on factual test data, since counterfactual ground truth was unavailable. The candidates use recurrent or Transformer backbones with approaches to balance representations against prescription information.

A. Evaluation Model Selection

  • The candidate set comprised CRN-E, a Transformer with gradient reversal (T-GRL), and a Transformer with domain confusion loss (T-DC).CRN-E uses an LSTM backbone; T-GRL and T-DC use GPT-based Transformer backbones.Source: A. Evaluation Model Selection · p. 11
  • The evaluation models estimate A1c at the very next admission rather than autoregressively, to avoid unreliable bias with the dataset's high missing rate.T-DC combines prediction and domain-confusion losses, balanced by α.Source: A. Evaluation Model Selection · pp. 11–12
  • CRN-E trains its representation to predict A1c while reducing its predictiveness of the factual prescription through gradient reversal.Its overall loss subtracts the prescription loss, with α setting the trade-off between outcome and prescription predictions.Source: A. Evaluation Model Selection · p. 11
  • CRN-E was selected as the evaluation model after candidates were compared by mean square error on factual test data.The authors lacked ground-truth values for counterfactual estimation.Source: A. Evaluation Model Selection · p. 12, Table 6 · p. 12

B. Hyperparameter Tuning for Clinical Decision Transformer

CDT hyperparameters were tuned using training and validation data, with the validation-set Youden index guiding selection.

B. Hyperparameter Tuning for Clinical Decision Transformer

C. Behavior Cloning Model

The behavior-cloning model uses an LSTM to imitate clinicians’ medication prescriptions, without manipulating recommendations toward an intended outcome.

C. Behavior Cloning Model

  • The behavior-cloning model uses LSTM networks to imitate clinicians’ medication-prescribing behavior.Unlike CDT, it is not allowed to manipulate recommendations with any intention.Source: C. Behavior Cloning Model · p. 13
  • It is trained with cross-entropy loss over the factual medication sets prescribed by clinicians.The training approach matches CDT training, although the model’s aim differs.Source: C. Behavior Cloning Model · p. 13
  • Its hyperparameters are presented in Table 8.The supplied passage identifies the table but does not provide its values.Source: Table 8 · p. 13

D. Diabetes Dataset

The dataset was extracted from an EHR system and includes 4,788 patients diagnosed with diabetes who were admitted between January 2015 and July 2021.

D. Diabetes Dataset

  • 4,788 patients diagnosed with diabetes and admitted between January 2015 and July 2021 were screened from an EHR system.The extraction was approved by the National Health Insurance Corporation of Korea review committee and the Korea University Ansan Hospital Institutional Review Board (2021AS0264).Source: D. Diabetes Dataset · p. 14
  • Table 9 presents baseline characteristics of the extracted dataset.The supplied caption does not provide characteristic values.Source: Table 9 · p. 14
  • Table 10 describes the dataset for CDT after goal-conditioned sequencing.The supplied caption does not provide dataset values.Source: Table 10 · p. 14
  • Table 11 describes the evaluation-model dataset after records without A1c data were removed.The supplied caption does not provide dataset values.Source: Table 11 · p. 14
  • Table 12 lists the prescribable medications.The supplied caption does not provide the medication names.Source: Table 12 · p. 15
  • Tables 13 and 14 describe baseline columns and laboratory-test columns, respectively.The supplied captions do not provide column names or values.Source: Table 13 · p. 15, Table 14 · p. 16
Loading 2302.00612v1…