利用基于Mamba架构的表征改进纵向数据的患者亚型分类

# 基于Mamba架构表示提升纵向数据中的患者亚型分类
来源：https://arxiv.org/html/2606.28623
###### 摘要

利用电子健康记录（EHR）数据进行有效的患者亚型分类（也称分组或聚类）能极大推动精准医疗的发展。然而，对时间序列EHR数据集进行亚型分类因EHR数据自身固有的问题（如复杂性和不规则性）而具有挑战性。本研究提出了一种基于Mamba的自监督模型，该模型能学习有效的EHR表示，从而实现更优的患者亚型分类。我们在公开和私有的真实EHR数据集上对所提模型进行了评估，根据现有标签对数据进行分类，并根据模型学到的表示对患者进行亚型分类。通过大量实验，我们证明了模型的设计选择在预测性能上优于有竞争力的基线模型。此外，我们评估了多种聚类技术，以表明我们的发现为基于EHR模型的时间序列记录对患者进行亚型分类提供了有价值的见解¹¹。

¹¹我们的实现可在 https://github.com/healthylaife/triplet_mamba 获取。

患者亚型分类，自监督学习，电子健康记录，Mamba

††版权：cc
††ccs：计算方法 学习潜在表示

## 1. 引言

医生依赖包括化验结果、手术、用药和诊断在内的全面病史来确定合适的治疗干预措施。除了个体临床决策外，研究人员利用这些纵向记录和人口统计数据来识别队列模式、评估药物疗效和分析疾病预后。该领域的一个关键方法是患者亚型分类，它旨在识别具有共同疾病进展路径的患者群体，以解决固有的异质性并促进精准医疗。通过识别复杂纵向轨迹中的共同特征，这些模型能够揭示隐藏的疾病模式（Rajkomar 等人，2018；Li 等人，2020；Mottalib 等人，2023），从而显著增强治疗方案的个性化和疾病早期发作的预测（Lee 和 Van Der Schaar，2020；Fredriksen 等人，2024）。有效分类这些轨迹不仅仅是预测性的工作，更是通过靶向干预改善临床结局的功能性要求。

尽管该任务至关重要，但由于医疗数据本身具有缺失性、极端稀疏性和非规则采样的特点，对时间序列数据集进行亚型分类异常困难。虽然基于Transformer的模型在临床预测中取得了成功（Pang 等人，2021；Li 等人，2020；Rupp 等人，2023；Rasmy 等人，2021；Li 等人，2022；Labach 等人，2023；Poulain 和 Beheshti，2024），但其二次计算复杂度 $O(L^2)$ 造成了关键瓶颈（Vaswani 等人，2017）。这种缩放特性有效限制了上下文长度（即模型能够处理的病史量），迫使在截断有价值的患者记录或采用计算上昂贵的离散化方法（如分箱或插值）之间做出选择（Shukla 和 Marlin，2019；Che 等人，2018）。这些限制对亚型分类尤其不利，因为亚型分类依赖于捕捉长期医疗动态来识别患者状态的细微转变。

为解决这些挑战，我们提出了 **Triplet-Mamba**，一种专为纵向表示学习和增强患者亚型分类而设计的自监督架构。我们的模型从传统的Transformer转向基于Mamba的主干，利用选择性状态空间模型（SSM）实现线性时间复杂度 $O(L)$（Gu 和 Dao，2024）。这一转变是处理定义稳健临床表型所需的长患者轨迹、同时不超出内存限制的功能性必要条件（Gu 和 Dao，2024）。SSM主干充当高效的连续时间滤波器，特别适合EHR数据中固有的非均匀采样率。

与先前对数据进行离散化的方法不同，**Triplet-Mamba** 通过直接嵌入连续时间戳和测量值来原生处理非规则采样。通过使用自监督预测任务训练这些表示，模型学习到一个结构化的潜在空间，该空间捕捉真实的疾病进展，而非简单的预测性记忆。我们通过在公开（Silva 等人，2012；Johnson 等人，2020）和私有数据集上的大量实验证明，这种设计相比于有竞争力的基线模型（Tipirneni 和 Reddy，2022；Labach 等人，2023）能产生更连贯且临床相关的患者亚型。

我们的主要贡献可总结如下：

- • 一种新颖的基于Mamba的架构 **Triplet-Mamba**，它满足了广泛上下文长度的要求，以有效捕捉纵向医疗动态。
- • 为该架构设计了输入表示，嵌入连续时间戳和测量值，保留了细粒度信息。
- • 一种新颖的患者亚型分类流程，通过自监督预测任务学习患者时间序列表示，并将时间序列分类为有意义的亚型。

### 关于机器学习在医疗保健背景下的通用见解

本研究表明，在医疗保健中实现有效且高效的患者亚型分类关键取决于学习可扩展的、时间上富有表现力的表示，这些表示要与纵向EHR数据的结构对齐。通过优先考虑线性时间序列模型和自监督目标而非仅依赖任务特定的监督，本研究证明长而规则的患者轨迹可以在不牺牲计算可行性或临床保真度的情况下进行建模。一个核心见解是，由原则性的事件级别编码和自监督预测支持强大的表示学习，能够产生跨数据集和聚类方法更连贯、更稳定、更具临床意义的患者亚型的潜在空间。更广泛地说，这项工作强调患者分层的进步源于同时考虑效率、时间结构以及对亚型质量的评估，从而强化了亚型分类作为医疗机器学习中一个基础性的、由表示驱动的问题的地位。

## 2. 相关工作

##### 从Transformer到状态空间模型

Transformer架构在自然语言处理中的成功促使其被应用于临床数据，像STraTS（Tipirneni 和 Reddy，2022）和DuETT（Labach 等人，2023）等模型通过设计专门的输入表示和注意力机制取得了重大进展。尽管强大，但基于Transformer的模型的一个关键瓶颈是自注意力机制的二次计算复杂度（Vaswani 等人，2017），这限制了能够高效处理的患者历史长度（Rajkomar 等人，2018；Wen 等人，2023；Li 等人，2022）。在医疗保健领域，这一限制可能很重要，因为更长的上下文通常能提高临床预测性能（Wornow 等人，2025）。

这一挑战推动了亚二次架构的发展，最著名的是状态空间模型（SSMs）（Gu 等人，2020；Fu 等人，2023；Poli 等人，2023；Dao 等人，2023）。最近引入的Mamba（Gu 和 Dao，2024），一种选择性SSM，实现了与Transformer竞争的性能，但具有线性时间复杂度，使其对非常长的序列非常高效。Mamba在医疗保健领域的潜力正在被积极探索（Mottalib 等人，2025）。例如，EHRMamba（Fallahpour 等人，2024）将架构调整为EHR预测任务，而ClinicalMamba（Yang 等人，2024）已被发展为在大型临床笔记语料库上预训练的基础模型。我们的工作建立在这一势头之上，利用Mamba的效率从长患者轨迹中学习。

##### 为患者亚型分类学习表示

这些强大表示的一个主要临床应用是患者亚型分类——根据疾病进展将患者无监督地分组为有临床意义的表型。有效的亚型分类可以为精准医疗提供信息，并揭示复杂疾病的新见解。最近的计算方法（Gorla 等人，2025；Jiang 等人，2023）已在阐明和验证新表型方面发挥重要作用，使得基于静态横截面快照和动态纵向轨迹进行患者分层成为可能。

多种深度学习技术已被应用于此问题。早期工作使用递归神经网络，如时间感知LSTM（T-LSTM）（Baytas 等人，2017），为聚类学习患者表示。其他方法利用自编码器（Ienco 和 Interdonato，2020），有时通过注意力机制增强（Li 等人，2025，2023），从时间序列信息中学习用于亚型分类的相似性度量。更近期的研究还探索了多模态数据；例如，M-ClustEHR（Bampa 等人，2024）采用多模态自编码器从多样的EHR数据中学习用于脓毒症亚型分类的全面表示。X-DEC模型验证了深度嵌入聚类模型及其通过X形变分自编码器集成混合数据类型的适应性（de Kok 等人，2024）。Med-ROAR（Baskett 等人，2025）通过将数据编码为多级离散表示来分层划分数据，采用改进的自注意力机制来识别同质患者。

我们的模型 **Triplet-Mamba** 综合了这些研究线索。我们利用Mamba的架构效率来有效建模长且不规则的患者历史，并结合专门设计的三元组输入表示来捕捉细粒度的时间细节。通过用自监督目标训练这些强大的表示，我们旨在产生比先前方法更连贯、更具临床意义的患者亚型。

## 3. 方法

##### 问题定义

考虑一个与典型EHR患者记录对应的数据集结构。对于每个患者，记录包含一个对应于非规则患者观测（如生命体征和化验结果）的事件时间序列，以及一组不随时间变化的静态变量（如人口统计信息）。这可以表示为一个稀疏的非规则时间序列数据集，格式为 $\mathcal{D} = \{ ( \mathbf{s}^k, \mathcal{T}^k, y^k ) \}_{k=1}^N$，包含 $N$ 个标记样本，其中第 $k$ 个样本包含一个 $n_s$ 维的静态向量 $\mathbf{s}^k \in \mathbb{R}^{n_s}$，一个可变长度 $n_k$ 的多变量时间序列 $\mathcal{T}^k = ( t_1^k, t_2^k, \cdots, t_{n_k}^k )$，以及一个对应预定义任务（如死亡率或疾病患病率）的二元标签 $y^k \in \{0,1\}$。

每个事件 $t_i^k$ 是一个观测三元组，正式定义为三元组 $(t, f, v)$，其中 $t \in \mathbb{R}_{\geq 0}$ 表示时间，$f \in \mathcal{F}$ 表示特征或变量，$v \in \mathbb{R}$ 表示观测值，例如，[72.9 天, 体重, 153 磅]。省略第 $k$ 个患者的 $k$，一个长度为 $n$ 的多变量时间序列 $\mathcal{T}$ 定义为 $n$ 个观测三元组的集合，即 $\mathcal{T} = \{ (t_i, f_i, v_i) \}_{i=1}^n$。底层时间序列变量集合 $\mathcal{F}$ 可能包括生命体征（如体温）、化验测量（如血红蛋白水平、体重等）以及输入/输出事件（如液体摄入量和尿量）。因此，对于患者 $p$，目标是根据 $(\mathbf{s}^p, \mathcal{T}^p)$（静态变量 $\mathbf{s}^p$ 和多变量时间序列 $\mathcal{T}^p$）预测 $y^p$。

(Figure 1 标题) 图 1. **Triplet-Mamba** 架构

### 3.1. 模型架构

我们模型 **Triplet-Mamba** 的架构如图 1 所示。我们将输入定义为一组观测三元组，其中输入中的每个观测三元组使用 **初始三元组嵌入** 模块进行嵌入。初始三元组嵌入随后通过一个 **上下文三元组嵌入** 模块，该模块利用 Mamba 块为每个三元组编码上下文。然后，**融合自注意力** 模块通过自注意力机制结合这些上下文嵌入，生成输入时间序列的嵌入，该嵌入与静态（人口统计）嵌入连接，并通过预测层进行最终预测。我们将亚型分类作为从学到的表示中衍生出的下游任务引入。

#### 3.1.1. 初始三元组嵌入

**初始三元组嵌入** 将输入序列 $\mathcal{T} = \{ (t_i, f_i, v_i) \}_{i=1}^n$ 中的每个 $n$ 三元组映射为一个 $\tau$ 维向量 $e_i \in \mathbb{R}^\tau$，其中 $\tau$ 是映射变量维度的超参数。向量 $e_i$ 通过将所有嵌入求和为一个单一嵌入来计算，使得 $e_i = e_i^f + e_i^v + e_i^t \in \mathbb{R}^\tau$。