🤖 machine learning

HyMaTE: A Hybrid Mamba and Transformer Model for EHR Representation Learning

The paper introduces HyMaTE, a novel hybrid model that combines the linear-time efficiency of Mamba State Space Models with the channel-level mixing capabilities of Transformers to effectively address the challenges of sparsity, missing values, and long sequences in Electronic Health Record (EHR) data, demonstrating superior performance and interpretability in clinical prediction tasks.

Original authors: Md Mozaharul Mottalib, Thao-Ly T. Phan, Rahmatollah Beheshti

Published 2026-06-24

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Md Mozaharul Mottalib, Thao-Ly T. Phan, Rahmatollah Beheshti

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: The "Library" vs. The "Speedster"

Imagine a patient's Electronic Health Record (EHR) as a massive, chaotic library containing their entire medical history. This library has thousands of books (lab results, doctor notes, vital signs) written at irregular times, with many pages missing or torn out.

To predict what might happen to a patient next (like if they might get sick again), we need a computer model to read this library. However, current models have two major flaws:

The "Super-Reader" (Transformers): These are like brilliant scholars who can read every book in the library and understand how they all connect. But, they are slow. If the library gets too big, they get overwhelmed and run out of time (computational power). They also struggle to read the whole library at once, so they often have to skip the beginning or the end.
The "Speedster" (Mamba): These are like a fast runner who can sprint through the entire library in seconds without getting tired. They are incredibly efficient with long stories. However, they are a bit "one-track minded." They are great at reading the story in order (time), but they aren't very good at noticing how different types of books (like how a heart rate relates to a blood test) interact with each other at the same moment.

The Solution: HyMaTE (The Hybrid Librarian)

The authors of this paper built HyMaTE, a new model that acts like the perfect librarian by combining the best traits of both the "Super-Reader" and the "Speedster."

Think of HyMaTE as a two-step reading process:

Step 1: The Speed Run (Mamba): First, the model uses the "Speedster" (Mamba) to quickly scan through the patient's long, messy history. It grabs the general flow of the story and handles the long timeline efficiently without getting bogged down.
Step 2: The Deep Dive (Transformer): Once the story is scanned, the model switches to the "Super-Reader" (Transformer). It pauses to look closely at specific moments. It asks, "How does this heart rate relate to that blood pressure reading right here?" This step ensures it understands the complex relationships between different medical variables.

Finally, it uses a special "Fusion Attention" mechanism. Imagine a conductor in an orchestra. The Speedster played the instruments fast, and the Super-Reader understood the notes. The Conductor (Fusion Attention) blends them together to create a single, perfect symphony that represents the patient's health.

How It Learned: The "Fill-in-the-Blanks" Game

Before HyMaTE was ready to help doctors, it had to study. The paper describes a clever training method:

The Game: The model was given a patient's history but with many numbers hidden (masked).
The Task: It had to guess the missing numbers based on the rest of the story.
The Result: By playing this "Fill-in-the-Blanks" game over and over, the model learned to understand the patterns of human health without needing a human to label every single answer. This is crucial because in real life, we often don't have perfect, labeled data for every patient.

The Results: Winning the Race

The researchers tested HyMaTE against other top models on three different "libraries" (datasets):

PhysioNet 2012: A public dataset of ICU patients.
MIMIC-IV: A massive dataset from a major hospital.
Pediatric Weight Management: A private dataset of children.

The Outcome: HyMaTE won. It was better at predicting:

Who might die in the hospital (Mortality).
How long a patient would stay (Length of Stay).
If a patient would return to the hospital (Readmission).
If a child would lose weight on medication.

It beat the "Super-Readers" (Transformers) because it didn't get overwhelmed by long histories, and it beat the "Speedsters" (Mamba) because it understood the complex relationships between different medical tests.

The "Why" (Interpretability)

One of the coolest parts of HyMaTE is that it doesn't just give an answer; it explains why.

The Spotlight: Because of the "Conductor" (Fusion Attention) step, the model can highlight exactly which moments in the patient's timeline were most important.
- Example: In a case study, the model predicted a high risk of death for an 82-year-old. When the researchers looked at the "spotlight," they saw the model was intensely focused on the patient's blood pressure readings from a specific day. This matched what real doctors know is a critical warning sign.
The Scoreboard: The model also ranked which medical variables mattered most overall (like Heart Rate or Blood Urea Nitrogen), and these rankings matched established medical knowledge.

Summary

HyMaTE is a new AI tool that reads patient medical histories faster than the old "Super-Readers" and understands the complex connections between different medical tests better than the old "Speedsters." It learns by guessing missing data, performs better at predicting patient outcomes, and can show doctors exactly which parts of a patient's history led to its predictions.

Technical Summary: HyMaTE – A Hybrid Mamba and Transformer Model for EHR Representation Learning

Problem Statement

Electronic Health Records (EHRs) serve as a critical resource for analyzing patient health progression, documenting longitudinal histories including diagnoses, procedures, and laboratory results. However, modeling EHR data presents significant challenges due to its inherent complexity: it consists of long, multivariate sequences that are often sparse, irregularly sampled, and contain numerous missing values.

Existing deep learning approaches face specific limitations in this domain:

Transformer-based models: While successful in capturing long-range temporal dependencies, they suffer from quadratic computational complexity ( $O(N^2)$ ) and memory requirements relative to input sequence length. This restricts their ability to process extensive patient histories (spanning thousands of tokens) without truncation or quantization, which can degrade accuracy. Furthermore, standard Transformers often prioritize sequence mixing over channel mixing, potentially overlooking interactions between different event types (variables).
State Space Models (SSMs) / Mamba: These models offer linear-time sequence modeling and superior efficiency for long sequences. However, traditional Mamba architectures primarily focus on sequence mixing and lack robust mechanisms for channel mixing, which is essential for multivariate time series data where inter-variable relationships are critical.

Methodology

The authors propose HyMaTE (A Hybrid Mamba and Transformer Model for EHR Representation Learning), a novel architecture designed to synthesize the linear efficiency of SSMs with the contextual and channel-mixing strengths of Transformers.

1. Input Representation

EHR events are represented as triplets $(t, f, v)$ , consisting of observation time, feature/variable, and observed value. These triplets are embedded into a $\tau$ -dimensional vector by summing embeddings for the feature (lookup table), continuous value, and time (derived via Feed Forward Networks).

2. Model Architecture

The architecture follows a sequential processing pipeline:

Contextual Attention (Mamba Blocks): The embedded inputs pass through multiple Mamba blocks. These blocks utilize a discretized State-Space Model (SSM) to filter relevant information efficiently. The SSM unit maps input sequences to output sequences via an implicit hidden state, employing a global convolution kernel for parallel processing. This stage handles the efficient encoding of long sequences.
Local Context Extraction (Self-Attention): Following the Mamba encoding, self-attention layers are applied to extract local context and immediate dependencies within the sequence.
Fusion Self-Attention: The output from the Mamba blocks is passed to a Fusion Self-attention module. This layer applies a Feed Forward Network (FFN) to each contextual embedding and uses a softmax function to calculate attention weights ( $\alpha_i$ ). The final time-series embedding ( $z_T$ ) is computed as a weighted sum of the contextual embeddings.
Static Variable Integration: Static demographic variables are embedded separately via an FFN with two hidden layers to produce a static embedding ( $z_d$ ).
Final Representation: The temporal embedding ( $z_T$ ) and static embedding ( $z_d$ ) are concatenated to form the final patient representation ( $e_E$ ).

3. Training Strategy

HyMaTE employs a semi-supervised two-phase training approach:

Self-Supervised Pre-training: The model is pre-trained on a larger dataset using a masked forecasting task. It learns to predict future variable values based on observed history, utilizing a Masked Mean Squared Error (MSE) loss to handle missing values. This phase allows the model to learn robust representations from limited labeled data.
Supervised Fine-tuning: The pre-trained model is fine-tuned on specific downstream clinical prediction tasks (e.g., mortality, length of stay) by attaching task-specific heads to the final embeddings.

Key Contributions

Hybrid Architecture: HyMaTE strategically combines Mamba blocks for efficient linear-time sequence encoding with Transformer-based attention mechanisms to address the channel-mixing limitations of pure SSMs and the computational bottlenecks of pure Transformers.
Fusion Mechanism: The introduction of a Fusion Self-attention layer allows for the synthesis of local context (from self-attention) and global sequence information (from Mamba) into a unified representation.
Interpretability: The model explicitly incorporates interpretability through the analysis of attention weights from the fusion layer and the coefficients of the final predictive dense layers, providing both temporal and global feature importance insights.
Semi-Supervised Learning: The integration of a self-supervised forecasting pre-training phase enables the model to learn high-quality representations from scarce labeled EHR datasets, a critical advantage in healthcare settings.

Experimental Results

HyMaTE was evaluated on three diverse datasets: PhysioNet Challenge 2012 (ICU mortality), MIMIC-IV (mortality, length of stay, readmission), and a private Pediatric Weight Management dataset.

Performance: HyMaTE consistently outperformed baseline models, including SeFT, STraTS, EHR-Mamba, and DuETT.
- On MIMIC-IV mortality prediction, HyMaTE achieved an AUROC of 0.907 and AUPRC of 0.648, surpassing the second-best model (EHR-Mamba) by 0.026 in AUROC.
- On PhysioNet 2012, it achieved an AUROC of 0.868.
- On the Pediatric weight-loss task, it achieved an AUROC of 0.752, significantly outperforming EHR-Mamba (0.704).
Ablation Studies: Removing key components resulted in performance degradation:
- Removing the input embedding caused a notable drop, highlighting the importance of triplet representation.
- Removing self-attention or attention fusion led to identical performance drops, confirming the interdependence of these components for effective feature blending.
- Removing self-supervised pre-training significantly impacted performance, particularly on the pediatric dataset, validating the necessity of the semi-supervised approach.
- Removing Mamba blocks drastically reduced performance, confirming the value of linear-time sequence processing.
Scalability: The model demonstrated robustness in handling long sequences (up to 8,654 observations), with minimal performance decline even at longer sequence lengths.

Significance and Claims

The paper claims that HyMaTE offers a scalable and generalizable solution for real-world healthcare applications by effectively balancing computational efficiency with representational richness. By overcoming the quadratic complexity of Transformers and the channel-mixing limitations of Mamba, HyMaTE provides a unified representation of EHR data that captures both long-range dependencies and fine-grained inter-variable relationships.

The authors emphasize that the model's interpretability—derived from attention weights and dense layer coefficients—aligns with established clinical findings (e.g., identifying NIMAP and Heart Rate as critical factors in mortality), thereby enhancing transparency and potential clinical utility. The work positions HyMaTE as a step forward in EHR representation learning, capable of handling the complexities of longitudinal medical data while maintaining the efficiency required for processing extensive patient histories.