DT-BEHRT: Disease Trajectory-aware Transformer for Interpretable Patient Representation Learning

This paper introduces DT-BEHRT, a graph-enhanced transformer model that improves predictive performance and interpretability in electronic health record analysis by explicitly modeling disease trajectories within organ systems and employing a novel pre-training strategy to capture heterogeneous medical code interactions.

The Gist

Imagine your medical history not as a giant, messy pile of paperwork, but as a complex, unfolding story. For years, computers trying to read these stories (Electronic Health Records, or EHRs) have struggled because they treated every piece of information the same way. They didn't know the difference between a diagnosis (the main plot twist), a medication (the character's tool), or a lab test (a clue).

The paper introduces DT-BEHRT, a new AI system designed to read these medical stories the way a skilled doctor does: by understanding the plot, the characters, and how the story changes over time.

Here is a breakdown of how it works, using simple analogies:

1. The Problem: The "One-Size-Fits-All" Mistake

Previous AI models were like a librarian who shoves every book into the same bin. Whether it's a mystery novel, a cookbook, or a biography, they are all just "books."

• In medicine: Old models treated a diagnosis of "Heart Failure" the same way they treated a prescription for "Aspirin."
• The issue: In real life, diagnoses are the driving force of a patient's story. They connect to other diseases (like how a heart issue connects to lung issues). Treatments are more like actions that happen over time. Mixing them up confuses the AI.

2. The Solution: DT-BEHRT (The "Smart Detective")

DT-BEHRT is a "Disease Trajectory-aware Transformer." That's a fancy way of saying it's an AI that pays attention to how diseases move and change over a patient's life. It uses three main "tools" to solve the mystery:

Tool A: The "Chapter Summarizer" (Disease Aggregation)

• The Analogy: Imagine a book with 19 different chapters (like "Respiratory System," "Heart," "Digestive System"). Instead of reading every single word, this tool creates a summary card for each chapter.
• How it works: If a patient has five different codes for lung problems, DT-BEHRT groups them under a "Respiratory System" card. It asks: "How many different lung issues has this person had?" This helps the AI see the big picture of a specific body system without getting lost in the details.

Tool B: The "Time-Traveling Map" (Disease Progression)

• The Analogy: Think of a patient's visits as stops on a train journey. Old models just looked at the list of stops. DT-BEHRT draws a map connecting the stops.
• How it works: It builds a graph (a web of connections) that shows how a patient moved from Visit 1 to Visit 2. Did the heart condition get worse? Did a new lung issue appear after the heart issue? This tool tracks the direction of the disease, understanding that "Heart Failure" today might lead to "Kidney Trouble" tomorrow.

Tool C: The "Final Report" (Patient Representation)

• The Analogy: After gathering the chapter summaries and the time-travel map, the AI writes a final executive summary of the patient's entire health story.
• How it works: It combines everything—the specific codes, the organ system summaries, and the progression map—into one single "fingerprint" for the patient. This fingerprint is what the AI uses to make predictions (like "Will this patient be readmitted to the hospital?").

3. The Training: "Fill-in-the-Blanks" and "Guess the Category"

To teach this AI to be smart, the researchers didn't just show it answers; they made it play games during training:

• The Masking Game: They hid random medical codes in a patient's history and asked the AI to guess what was missing. This taught it how different diseases and treatments usually appear together (e.g., if you see "Diabetes," you often see "Insulin").
• The Ancestor Game: They hid a specific disease code and asked the AI to guess the category it belongs to (e.g., if the code is "Pneumonia," the AI must guess it belongs to "Respiratory System"). This forced the AI to learn the medical "family tree" and understand the big picture, not just the small details.

4. Why It Matters: The "Doctor's Intuition"

The best part about DT-BEHRT is that it's interpretable.

• Old AI: "I predict a 90% chance of readmission because of these 50 random codes." (The doctor has no idea why).
• DT-BEHRT: "I predict a 90% chance of readmission because the patient's Respiratory System is struggling, and their Heart Condition has been getting worse over the last three visits."

It mimics how a human doctor thinks. It doesn't just look at the data; it looks at the story of the disease.

The Bottom Line

DT-BEHRT is like upgrading from a calculator to a narrative analyst. It understands that a patient's health isn't just a list of symptoms, but a dynamic story where diseases interact, evolve, and tell a specific tale about the body's systems. By respecting this story, it makes better predictions and gives doctors a clear explanation of why those predictions were made.

Technical Summary

1. Problem Statement

The paper addresses critical limitations in existing predictive modeling approaches for Electronic Health Records (EHRs). While structured EHRs contain longitudinal patient data (visits with medical codes), current methods struggle with two main issues:

• Homogeneous Treatment of Codes: Existing sequence-based (e.g., BEHRT, Med-BERT) and graph-based models often treat all medical codes (diagnoses, medications, procedures, labs) uniformly. However, these codes play fundamentally different roles. For instance, diagnosis codes drive disease trajectories and interact densely within organ systems, whereas treatment codes (medications/procedures) are often temporally sequential but less interactive within a single visit.
• Inadequate Temporal and Structural Modeling: Sequence models often fail to capture complex intra-visit code dependencies because code ordering within a visit is unreliable. Conversely, pure graph models often struggle to capture sequential dependencies across multiple hospital visits.
• Lack of Interpretability: Many high-performing models act as "black boxes," failing to align with the disease-centered reasoning used by clinicians.

2. Methodology: DT-BEHRT Architecture

The authors propose DT-BEHRT, a graph-enhanced sequential architecture designed to disentangle disease trajectories by explicitly modeling diagnosis-centric interactions. The model consists of four core modules:

A. Sequence Representation (SR) Module

• Input: A sequence of hospital visits, where each visit contains a set of medical codes.
• Embeddings: Codes are enriched with type embeddings (diagnosis, medication, lab, procedure) and visit-index embeddings.
• Mechanism: Uses a standard Transformer encoder (BERT-style) with a special [SEQ] token to summarize the entire sequence. It treats intra-visit code ordering as unordered (using a mask) to avoid relying on unreliable timestamps.

B. Disease Aggregation (DA) Module

• Goal: To capture high-level semantic patterns and interactions of diseases within specific organ systems.
• Mechanism: Leverages the ICD-9 ontology (19 top-level chapters). The model introduces special DA tokens, one for each organ system chapter.
• Logic: If a patient has a sufficient number of distinct diagnosis codes belonging to a specific ICD-9 chapter, the corresponding DA token is appended to the sequence.
• Regularization: A covariance regularization loss is applied to DA tokens to ensure they capture decorrelated, diverse organ-system abstractions rather than redundant information.
• Attention Mask: Restricts attention so that DA tokens only attend to diagnosis codes within their specific organ system chapter.

C. Disease Progression (DP) Module

• Goal: To model the temporal evolution of diseases across visits.
• Mechanism: Constructs a heterogeneous graph for each patient layer.
  • Nodes: Virtual "visit nodes" (DP nodes) and diagnosis nodes.
  • Edges: Directed edges connect visit nodes to their diagnosis nodes; sequential edges connect visit nodes chronologically ($t \to t+1$); self-loops preserve information.
• Learning: Uses a Graph Attention Network (GAT) to update visit node representations by aggregating information from current diagnoses and previous visits, effectively capturing asynchronous progression patterns.

D. Patient Representation (PR) Module

• Integration: Combines three sources of information:
  1. The [SEQ] token (global sequence summary).
  2. DA tokens (organ-system level abstractions).
  3. DP tokens (temporal progression trends).
• Mechanism: An attention-based pooling mechanism uses the [SEQ] token as a query to weigh the importance of DA and DP tokens, producing the final patient representation vector ($\mathbf{h}_{[CLS]}$).

3. Pre-Training Framework

To enhance robustness and align the different modules, the authors design a novel pre-training strategy:

1. Global Code Masking Prediction (GCMP): Instead of masking individual tokens, the model masks all occurrences of a unique code across the entire trajectory. The model must predict the masked code type (Diagnosis, Medication, Lab, Procedure) simultaneously. This encourages learning trajectory-level co-occurrence semantics (e.g., comorbidities).
2. Ancestor Code Prediction (ACP): An auxiliary task where the model predicts the ICD-9 chapter (ancestor) of a masked diagnosis code. This is performed using both the SR module's [SEQ] representation and the DP module's final visit representation. This forces the graph and sequence modules to align with the ontology structure.

4. Key Contributions

1. Novel Architecture: DT-BEHRT is the first to explicitly disentangle diagnosis-centric interactions (via DA) and temporal progression (via DP) while treating other code types differently, moving beyond homogeneous modeling.
2. Tailored Pre-training: The combination of trajectory-level masking and ontology-informed ancestor prediction improves semantic alignment across modules and representation robustness.
3. Interpretability: The model's design mirrors clinical reasoning by highlighting specific organ systems (DA) and temporal disease trajectories (DP), making its predictions explainable to clinicians.

5. Experimental Results

The model was evaluated on three public EHR datasets: MIMIC-III, MIMIC-IV, and eICU.

• General Outcome Prediction: DT-BEHRT achieved state-of-the-art (SOTA) performance across tasks including In-hospital Mortality, Prolonged Length of Stay (PLOS), and Readmission.
  • Notable Gain: The largest improvements were seen in Readmission prediction, a task known for being difficult due to multifactorial causes.
  • Robustness: The model maintained strong performance across various patient subgroups (e.g., diabetes, CKD, cancer) and datasets.
• Phenotyping Prediction: In multi-label classification for predicting future diagnoses, DT-BEHRT achieved the best Macro-AUPRC, particularly for patients with multiple hospital visits (strong temporal dependencies).
• Ablation Studies:
  • Removing the DA module significantly hurt mortality prediction.
  • Removing the DP module significantly hurt readmission prediction.
  • The ACP pre-training task provided the most significant boost to mortality prediction, validating the importance of ontology alignment.

6. Significance and Conclusion

• Clinical Alignment: Unlike "black box" models, DT-BEHRT provides interpretable outputs that correspond to familiar clinical constructs (organ systems and disease timelines). Case studies show the model correctly identifies relevant respiratory conditions or cardiovascular progressions, aligning with physician reasoning.
• Scalability vs. Performance: While the model introduces computational overhead due to graph attention and multiple modules, the trade-off yields superior accuracy and interpretability, which are critical for clinical decision support.
• Future Work: The authors note that while they focused on diagnosis codes, future iterations could extend the specialized modeling approach to medications and procedures.

In summary, DT-BEHRT represents a significant step forward in EHR modeling by acknowledging the heterogeneity of medical data and integrating graph-based structural knowledge with sequential temporal modeling to create robust, interpretable patient representations.