🤖 machine learning

Improving Patient Subtyping on Longitudinal Data using Representations from Mamba-based Architecture

This paper proposes a self-supervised Mamba-based model that learns effective representations from irregular longitudinal electronic health record data to significantly improve patient subtyping and precision medicine outcomes compared to existing baseline models.

Original authors: Md Mozaharul Mottalib, Rahmatollah Beheshti

Published 2026-06-30

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Md Mozaharul Mottalib, 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 Picture: Sorting Patients Like a Librarian

Imagine a hospital as a massive library where every patient has a "book" of their medical history. These books are messy. Some pages are missing, some are written in different languages, and the entries are scattered randomly across time (like a patient getting a blood test on Monday, an X-ray three weeks later, and a check-up six months after that).

Doctors need to group these patients into "clubs" or subtypes based on how their diseases progress. This helps them treat people more precisely. However, sorting these messy, irregular books is incredibly hard for computers.

The Problem: The "Quadratic" Bottleneck

The paper explains that previous computer models (called Transformers) were like librarians who tried to read every single page of every book against every other page to find similarities.

The Analogy: If you have 10 books, it's easy to compare them all. But if you have 1,000 books, the librarian has to do a million comparisons. If you have 10,000 books, they have to do 100 million comparisons. This is called quadratic complexity ( $O(L^2)$ ).
The Result: To avoid the computer crashing, these old models had to throw away long patient histories or chop them into tiny, rigid chunks (like forcing a long story into a 50-word summary). This meant they missed the subtle, long-term patterns of how a disease actually develops.

The Solution: The "Mamba" Train

The authors built a new model called Triplet-Mamba. Instead of a librarian comparing every page to every other page, they used a Mamba architecture.

The Analogy: Think of Mamba as a high-speed train moving through the patient's history. It doesn't stop to compare every single station to every other station. Instead, it moves linearly (one step at a time), remembering the important context as it goes. This is linear complexity ( $O(L)$ ).
Why it matters: It can handle a patient's entire life history (even if it spans years) without the computer running out of memory. It's fast and efficient.

How It Handles Messy Data: The "Triplet" System

Real medical data is irregular. A patient might get a temperature reading at 2:00 PM and a blood test at 4:30 PM the next day. Old models often forced this data into a neat grid (like a spreadsheet), which loses the exact timing.

Triplet-Mamba treats every piece of data as a Triplet:

Time: When it happened.
Feature: What it was (e.g., "Blood Pressure").
Value: How much (e.g., "120").

The Analogy: Instead of forcing a messy diary into a calendar, the model reads the diary exactly as written: "On Tuesday, I felt dizzy (Time: Tuesday, Feature: Dizziness, Value: High)." This preserves the exact, messy reality of the patient's life.

The Training: Learning Without a Teacher

The paper describes a Self-Supervised approach.

The Analogy: Imagine teaching a student to understand a story by having them read a book with the last page missing and asking them to guess what happens next. They aren't being told the answer by a teacher; they are learning the structure of the story by trying to predict the future.
In the paper: The model looks at a patient's past records and tries to predict their future lab results. By doing this millions of times, it learns a deep, internal map of how diseases progress. It learns the "shape" of the disease, not just memorizing labels.

The Results: Better Grouping

The researchers tested this new model on real-world data (including ICU patients and children with weight management issues) and compared it to the best existing models (like STraTS and DuETT).

Better "Clubs": When they used the model's internal map to group patients, the groups were much tighter and more distinct.
- The Metric: They used a score called the Silhouette Score. Think of this as a "grouping quality" score. Triplet-Mamba got the highest scores, meaning the patients in each group were very similar to each other and very different from patients in other groups.
Better Predictions: The model was also slightly better at predicting outcomes (like mortality or weight loss) than the competition.
Stability: The groups the model found were stable. If you gave the model slightly different data, it still found the same groups. This suggests the groups are real patterns, not just random noise.

The Takeaway

The paper claims that by switching from "quadratic" models (which get slow and messy with long data) to "linear" Mamba models, and by treating data as flexible triplets rather than rigid grids, they created a system that understands long, messy patient histories much better. This leads to more accurate ways of grouping patients, which is the first step toward personalized medicine.

What the paper does NOT claim:

It does not claim this is currently being used in hospitals to treat patients today.
It does not claim the model can read medical text (like doctor's notes) or look at X-rays yet (it currently focuses on numbers and timestamps).
It does not promise that this will cure diseases, only that it improves the grouping and prediction tools used by researchers.

Technical Summary: Improving Patient Subtyping on Longitudinal Data using Representations from Mamba-based Architecture

Problem Definition
Effective patient subtyping—grouping patients with shared disease progression pathways—is critical for precision medicine but remains challenging due to the inherent nature of Electronic Health Record (EHR) data. Clinical records are characterized by missingness, extreme sparsity, and irregular sampling intervals. While Transformer-based models have shown success in clinical prediction, their quadratic computational complexity ( $O(L^2)$ ) creates a bottleneck for processing long patient histories. This limitation forces researchers to truncate valuable data or adopt computationally prohibitive discretization methods (e.g., binning), which can obscure subtle transitions in a patient's state that are vital for identifying robust clinical phenotypes.

Methodology: Triplet-Mamba
To address these challenges, the authors propose Triplet-Mamba, a self-supervised architecture designed for longitudinal representation learning. The model transitions from traditional Transformer backbones to a Mamba-based architecture utilizing a selective State Space Model (SSM). This shift enables linear-time complexity ( $O(L)$ ), allowing the model to process extensive patient trajectories without exceeding memory constraints.

The core components of the methodology include:

Triplet Input Representation: Unlike prior methods that discretize time, Triplet-Mamba handles irregular sampling natively. It embeds each observation as a triplet $(t, f, v)$ , representing time, feature, and value. These are mapped into a latent space using an Initial Triplet Embedding module, which combines learned embeddings for the variable, continuous value, and continuous time via Feed Forward Networks (FFNs).
Contextual Encoding (Mamba Blocks): The initial embeddings pass through a Contextual Triplet Embedding module consisting of stacked Mamba blocks. These blocks utilize RMS normalization, linear projections, and the SSM to filter relevant information, acting as an efficient continuous-time filter suited for non-uniform sampling rates.
Fusion Self-Attention: A Fusion Self-attention module aggregates the contextual embeddings from the Mamba blocks to generate a fixed-length time-series representation. This is concatenated with static demographic embeddings to form the final patient representation.
Training Strategy: The model employs a two-phase training approach:
- Self-Supervised Pre-training: The model learns representations via a forecasting objective on a larger dataset, predicting masked future values using a masked Mean Squared Error (MSE) loss. This allows the model to learn structured latent spaces capturing genuine disease progression without relying solely on task-specific labels.
- Fine-tuning: The learned representations are adapted for downstream tasks (binary classification or subtyping) using a dense layer with sigmoid activation and cross-entropy loss.

Key Contributions
The paper outlines three primary contributions:

Novel Architecture: The introduction of Triplet-Mamba, which leverages the efficiency of Mamba to capture long-range healthcare dynamics that exceed the context limits of Transformer models.
Specialized Input Representation: The design of an embedding scheme that preserves fine-grained temporal information by directly embedding continuous timestamps and measured values, avoiding the information loss associated with discretization.
Subtyping Pipeline: A complete pipeline that utilizes self-supervised forecasting to learn patient time-series representations, which are subsequently categorized into meaningful subtypes using standard clustering techniques (K-means, Spectral Clustering, GMM).

Experimental Results
The authors evaluated Triplet-Mamba on three datasets: the public PhysioNet 2012 (mortality prediction), MIMIC-IV (mortality prediction), and a private Pediatric Weight Management dataset.

Predictive Performance: Triplet-Mamba achieved state-of-the-art or competitive results in binary classification tasks. It outperformed baselines (STraTS, EHR-Mamba, DuETT) on MIMIC-IV and the Pediatric dataset, achieving a 1.7% increase in AUROC on MIMIC-IV and a 2.8% increase on the Pediatric weight-loss task. On PhysioNet 2012, it performed comparably to the best baseline (DuETT).
Subtyping Quality: In unsupervised clustering tasks, representations learned by Triplet-Mamba consistently yielded higher Silhouette scores across all datasets and clustering methods compared to baselines. For the Pediatric dataset, where ground truth was defined by obesity classes, Triplet-Mamba achieved the highest Adjusted Rand Index (ARI), indicating superior alignment with clinical classifications.
Stability: Bootstrap stability analysis demonstrated that the subtypes discovered by Triplet-Mamba were highly consistent (mean Jaccard Index of 0.814), significantly outperforming baselines and exceeding the threshold for stable clinical phenotyping.
Ablation Studies: Removing the triplet embedding, the fusion attention layer, or the self-supervised pre-training phase resulted in consistent performance declines, confirming that the synergy between these components is essential for the model's success.

Significance and Claims
The authors claim that Triplet-Mamba demonstrates that effective patient subtyping relies on learning scalable, temporally expressive representations that align with the structure of longitudinal EHR data. By prioritizing linear-time sequence models and self-supervised objectives, the study shows that long, irregular patient trajectories can be modeled without sacrificing computational feasibility or clinical fidelity.

The work emphasizes that robust representation learning, supported by principled event-level encodings and self-supervised forecasting, yields latent spaces that enable more coherent and clinically meaningful patient subtypes. The authors position subtyping not merely as a predictive exercise but as a foundational, representation-driven problem in healthcare machine learning. They note that while the current implementation focuses on lab readings and vital signs, the architecture is designed to be extensible to include procedures, diagnoses, and other modalities in future work.