scWGBS-GPT: 用于捕获单细胞全基因组亚硫酸氢盐测序中长程 CpG 依赖关系以增强表观遗传分析的预训练模型.
0. TL; DR
scWGBS-GPT是专为scWGBS（单细胞全基因组亚硫酸氢盐测序）数据在单CpG水平高分辨率分析而设计的生成式基础语言模型。通过结合CpG特殊词元设计、Mamba主干网络和交叉注意力头，能够高效处理超长序列，同时保留局部CpG相互作用和更广泛的基因组上下文。
模型在来自28种人和小鼠组织的100万个单细胞上进行预训练，能够重建稀疏的甲基化图谱，显著提升表观遗传分析的分辨率和准确性。在关键的生物医学应用中，它大幅超越了现有方法：改进细胞聚类以揭示先前未被识别的细胞异质性，增强轨迹推断以绘制精确的分化路径，并自动化识别与疾病进展和治疗靶点相关的关键表观遗传生物标志物。
1. 背景介绍
近年来，单细胞数据分析技术的快速发展彻底改变了人们对细胞异质性、谱系分化和表观遗传调控的理解。在这场变革的前沿是生成式语言模型，特别是基于Transformer的架构，它们已经改变了单细胞组学研究。诸如 scBERT、scGPT 和 scFoundation 等模型在单细胞RNA测序分析中展现了最先进的性能，在细胞聚类、注释和基因扰动预测等任务上表现出色。这些模型利用强大的深度学习技术提取单细胞数据中复杂的相互关系，显著提升了解读生物系统的能力。
然而，Transformer 的语言模型面临一个关键局限：其自注意力机制的计算复杂度随序列长度呈平方级增长，这限制了它们在处理长度超过几千个词元序列的数据集上的应用。这一局限使得许多重要的生物数据类型（如单细胞全基因组甲基化数据）无法充分受益于这些变革性方法。
单细胞全基因组亚硫酸氢盐测序是一种尖端技术，可在单核苷酸和单细胞分辨率下分析DNA甲基化，为全基因组表观遗传调控提供了无与伦比的见解。它揭示了细胞间的变异性，使人们能更深入地理解发育过程、细胞异质性和疾病相关的表观遗传变化。该技术通过对DNA进行亚硫酸氢盐处理以区分甲基化和未甲基化的胞嘧啶，然后进行高通量测序，生成每个细胞涵盖数十万到数百万个CpG位点的甲基化数据。这个全面的数据集蕴含着理解表观遗传调控的巨大潜力，但其巨大的长度和稀疏的覆盖度使得Transformer语言模型在计算上难以处理。
尽管前景广阔，但分析scWGBS数据面临着巨大的计算和分析挑战。序列的庞大长度（通常跨越数百万个CpG位点）及其稀疏、可变的覆盖度使得传统的机器学习方法力不从心。当前的方法通常将基因组划分为固定区域（如 10-kb 窗口）并计算区域甲基化率。虽然这种方法在计算上可管理，但它牺牲了单核苷酸精度，并且无法捕捉CpG位点之间错综复杂的依赖关系。此外，这些方法忽略了更广泛的基因组上下文，导致信息大量丢失，限制了其发现细微但关键的生物学机制的能力。
为应对这一挑战，作者开发了一个名为 scWGBS-GPT 的生成式基础语言模型，专门用于在单CpG分辨率下分析scWGBS数据。
2. scWGBS-GPT 方法
scWGBS-GPT 的核心架构融合了三个关键组件：CpG特殊词元设计、基于选择性状态空间模型的Mamba主干网络，以及交叉注意力头，旨在应对scWGBS数据超长、稀疏的挑战。
2.1 scWGBS Tokenizer（分词器）
scWGBS 分词器是将原始基因组数据转换为适合模型处理格式的关键组件。鉴于scWGBS数据包含大片基因组区域中每个CpG位点的甲基化水平，分词器的主要作用是将这些甲基化值编码为词元化序列，同时保留输入数据的空间和生物学特征。
分词器的核心思想是将DNA序列视为CpG词的字符串，每个CpG位点被视为一个词元。这种设计将甲基化和未甲基化状态整合到词元表示中，确保甲基化模式得以保留供下游处理。
具体分词过程如下：
- DNA序列分割：基因组数据被分割成包含连续CpG位点的窗口。
- 编码甲基化状态与上下文：对于窗口中的每个CpG位点 $x_i$，其甲基化状态被编码为一个二元值（0 或 1）。此外，周围的核苷酸（例如每个CpG位点两侧的 ±3 个碱基）被合并以形成具有上下文感知的词元。这产生了一个 k-mer 词元表示，其中每个CpG都伴有其局部序列背景。
- 特殊词元插入：在每个窗口的开头插入特殊的
[BOS]（序列开始）词元，在末尾插入[SEP]（分隔）词元。[SEP]词元的最终隐藏状态将用作交叉注意力机制中的查询向量，以聚合整个序列的全局上下文。
最终的词元化序列表示为：
\[\mathrm{Tokenized Sequence} = [[\mathrm{BOS}], \mathbf{t}_1, \mathbf{t}_2, \ldots, \mathbf{t}_n, [\mathrm{SEP}]]\]其中 $\mathbf{t}_i$ 代表第 $i$ 个CpG位点及其甲基化状态和背景信息。
2.2 Mamba Backbone（主干网络）
Mamba 是一种最先进的序列建模框架，专为高效处理超长序列而设计。与Transformer架构不同，Mamba 利用选择性状态空间模型来实现线性时间复杂度，同时保持强大的表征能力。
Mamba 的核心状态更新方程为：
\[\mathbf{h}_t = \mathbf{A}\mathbf{h}_{t-1} + \mathbf{B}\mathbf{x}_t\]其中：$\mathbf{h}_t$ 是时间步 $t$ 的隐藏状态，$\mathbf{A}$ 是控制序列记忆的学习转移矩阵，$\mathbf{B}$ 是输入变换矩阵，$\mathbf{x}_t$ 是时间步 $t$ 的输入。
scWGBS-GPT 中的 Mamba 主干由 8 个堆叠的 Mamba块组成，总参数量约为 1M。每个块的操作如下：
- 卷积输入投影：输入序列 $\mathbf{X}$ 首先通过深度可分离卷积层处理以提取局部CpG模式：$\mathbf{X}’ = \mathrm{Conv1D}(\mathbf{X})$。
- 状态空间模型演化：每个词元通过连续时间状态空间层传递，更新内部序列表示：\(\mathbf{H}_t = \mathbf{A}\mathbf{H}_{t-1} + \mathbf{B}\mathbf{X}_t'\)。
- 选择性门控机制：Sigmoid 门控函数动态调节新信息的影响：$\mathbf{G}_t = \sigma (\mathbf{W}_g\mathbf{X}_t)$。
- 线性投影与归一化：输出被投影到新的特征空间并归一化：$\mathbf{Y}_t = \mathrm{LayerNorm}(\mathbf{W}_g\mathbf{H}_t)$。
这种深度堆叠机制使模型能够捕获scWGBS数据中的短程CpG模式和长程调控依赖性。
2.3 Cross Attention（交叉注意力）
交叉注意力是一种强大的机制，允许模型在处理一个序列时关注输入序列的相关部分。在 scWGBS-GPT 中，作者利用交叉注意力来捕获整个基因组中CpG位点之间复杂的依赖关系。
具体机制如下：给定输入序列 $\mathbf{X} = [\mathbf{x}_1, \mathbf{x}_2, \ldots, \mathbf{x}_n]$，其中 $\mathbf{x}_i$ 代表序列中第 $i$ 个词元的隐藏状态。
[SEP] 词元的最终隐藏状态被用作交叉注意力的查询矩阵 $\mathbf{Q}$。来自 Mamba 主干的每个CpG位点的隐藏状态同时作为键矩阵 $\mathbf{K}$ 和值矩阵 $\mathbf{V}$。交叉注意力操作描述为：
注意力加权的隐藏状态 $\mathbf{H}_{\mathrm{att}}$ 用于下游层的进一步处理，为最终的基因组序列表征做出贡献。这种机制使模型能够有效地从稀疏或不完整的数据中捕获有意义的模式，并建模DNA甲基化模式中常见的长程相互作用。
2.4 预训练与微调策略
在预训练阶段，scWGBS-GPT 采用了一种下一个模糊词元预测策略。在每个基因组位置，模型同时预测甲基化和非甲基化词元，损失权重根据甲基化率 $\alpha$ 及其补数 $(1-\alpha)$ 动态调整，有效捕获了CpG甲基化的概率性本质。
在微调阶段，通过交叉注意力头跨所有输入词元获得的 [SEP] 词元的最终表征，可作为全面的单细胞嵌入。该嵌入整合了CpG位点之间的局部和长程依赖关系，便于高精度、高生物相关性地执行下游任务，如细胞注释、分化轨迹推断和疾病状态预测。
2.5 训练数据
为了利用语言模型的缩放定律，作者策划了一个包含来自人和小鼠的100万个单细胞甲基化谱的综合跨物种数据集。数据集包括了所有 28 种小鼠组织类型和各种人体组织的甲基化谱。其中，90万个单细胞甲基化谱来自公共数据存储库（40万来自小鼠大脑区域，50万来自人脑区域）。为了进一步增强组织多样性，作者整合了来自 20 多种不同组织和器官、专注于纯化细胞类型的批量WGBS数据，贡献了约 5万 个样本。平均而言，每个单细胞样本产生 100万 个CpG词，为训练提供了丰富的资源。
3. 实验分析
作者对 scWGBS-GPT 进行了系统性评估，涵盖了模型效率、细胞注释、跨组织泛化、可解释性以及批量数据反卷积等多个方面。
3.1 模型效率与细胞注释性能评估
作者首先在scWGBS数据背景下评估了三种主流序列建模框架：Transformers、Flash Attention 和 Mamba。
- 图a-b：在推理速度和内存使用方面，Mamba 显著优于其他框架，尤其适合处理可能超过数百万词元的超长scWGBS序列。
- 图c：在自定义的下一个模糊词元预测预训练任务中，Mamba 以更少的浮点运算实现了更低的困惑度，证明其在大规模表观基因组分析中具有卓越的可扩展性。
基于以上结果，作者选择 Mamba 作为 scWGBS-GPT 的主干网络用于细胞类型注释。评估使用了来自人脑组织的scWGBS数据集。
- 图d：混淆矩阵显示，模型在测试集上的最终注释准确率达到 94.4%，各种神经元亚型（如 L2/3-IT, MSN-D1, Vip）的召回率均超过 90%，部分亚型（如 ODC, ASC）准确率超过 95%，表明模型有效捕获了亚型特异性甲基化特征。
- 图f：比较了不同预训练细胞数量下各种 k-mer 表示的性能，观察到一个清晰的缩放定律：6-mer 在所有 k-mer 中提供了最平衡和稳健的表征。
- 图e：比较了三种获取单细胞嵌入的策略：1) 对所有CpG词元的隐藏状态取平均（准确率 59.2%）；2) 提取特殊
[SEP]词元的最终隐藏状态（准确率 61.4%）；3) 使用交叉注意力头自动加权信息最丰富的CpG k-mer（准确率 94.4%）。交叉注意力嵌入显著优于前两种方法。 - 图g-h：热图对比显示，基于平均的嵌入将大部分变异集中在早期特征维度，难以区分某些细胞类型；而基于交叉注意力的嵌入则清晰地区分了不同类型，反映了模型聚焦于生物学相关CpG位点的能力。
- 图i：UMAP 可视化显示，基于交叉注意力的嵌入形成了定义良好的簇。具有更密切生物学关系的细胞类型（如 L2/3, L5, L6 兴奋性神经元）在UMAP空间中彼此靠近，而差异较大的细胞群体则相距较远。
Mamba 架构为处理超长scWGBS CpG词提供了可扩展性，而交叉注意力机制则能有效突出关键的CpG词，两者协同作用实现了高精度的单细胞注释。
3.2 跨组织泛化与单细胞轨迹分析
为了评估模型的少样本学习和跨组织泛化能力，作者将预训练后的 scWGBS-GPT 微调并应用于三个未出现在预训练数据集中的生物系统：结直肠癌、肺疾病和造血干细胞分化。
- 结直肠癌进展：
- 图d：模型成功区分了原发性肿瘤、淋巴结转移、未经治疗的肝转移和治疗后肝转移，揭示了不同的单细胞簇。
- 图e：伪时间轨迹推断捕获了肿瘤演化的清晰发展路径，反映了从原发性CRC到转移状态的转变。
- 图f：UMAP投影进一步证实了模型区分不同CRC亚型的能力。
- 肺疾病异质性：
- 图g：模型推断的伪时间轨迹提示了从正常肺细胞到疾病状态的细胞状态连续体，突出了与疾病进展相关的表观遗传轨迹。
- 图h：肺细胞类型的UMAP投影清晰区分了健康肺细胞、肺结节周围细胞和病变细胞。
- 造血干细胞分化：
- 图i：从模型嵌入推断的伪时间轨迹捕获了从HSC到定型祖细胞和完全分化的免疫细胞类型的明确谱系层次结构。
- 图j：相应的UMAP可视化显示，细胞根据其分化状态组织，密切相关的免疫细胞类型聚集在一起。
scWGBS-GPT 仅需少量微调样本即可准确注释细胞类型、推断细胞状态转换和重建来自单细胞甲基化数据的分化轨迹，展现了强大的跨组织泛化能力。
3.3 可解释性分析与肿瘤特异性生物标志物识别
作者通过分析模型在不同CRC进展阶段的注意力权重分布来评估其可解释性。
- 图a-b：散点图显示，在正常结肠样本中，注意力权重在不同甲基化水平上保持较低；而在原发性肿瘤样本中，模型选择性地将高注意力权重分配给低甲基化的CpG。
- 图c-d：展示了2号染色体上甲基化率与注意力权重的基因组分布。正常结肠样本中，注意力权重低且均匀。原发性肿瘤样本中，模型对某些肿瘤特异性低甲基化区域赋予高注意力，但其模式与转移性肿瘤中观察到的不同。
- 图e-g：在所有转移性样本（LN, ML, MP）中，模型系统地赋予低甲基化CpG高注意力得分，特别是在 2号染色体 50,000,000 位置附近观察到强烈且一致的注意力信号，而这在正常和原发性肿瘤样本中均不存在。
scWGBS-GPT 不仅能识别肿瘤特异性低甲基化模式，还能基于不同的表观遗传特征区分原发性和转移性肿瘤。基于注意力的深度学习方法可用于从单细胞甲基化数据中识别转移特异性生物标志物。
3.4 无参考反卷积性能
作者利用预训练的 scWGBS-GPT 在混合不同细胞类型的批量数据上进行微调，以执行端到端的无参考反卷积。使用真实的10种神经元细胞单细胞甲基化数据模拟批量样本，保留细胞间和细胞内的异质性以及测序覆盖度和深度的变化。
比较了 scWGBS-GPT 与三种经典无参考反卷积方法（RefFreeCellMix, EDec, MeDeCom）的性能。结果显示，经典方法由于采样策略和测序噪音难以处理棘手的细胞内类型异质性，而 scWGBS-GPT 则表现出相当稳健的性能（更高的 $R^2$ 值）。
scWGBS-GPT 在高度异质性的批量甲基化谱的无参考反卷积任务中表现出色，为解决这一传统难题提供了新工具。
