插补单细胞甲基化组的CpG Transformer.

0. TL; DR

CpG Transformer 通过将轴向注意力 (axial attention) 与滑动窗口自注意力 (sliding window self-attention) 相结合，将强大的 transformer 神经网络架构应用于甲基化矩阵的插补。

在一系列广泛的 scBS-seq 和 scRRBS-seq 数据集上，CpG Transformer 的性能达到了现有最先进水平，始终优于 DeepCpG 和 CaMelia 等竞争方法。

1. 背景介绍

DNA 甲基化，特别是 CpG 甲基化，在基因表达调控、发育生物学和癌症发生等广泛的生物学过程中扮演着核心角色。为了在细胞水平上研究这一现象，过去十年中发展了多种单细胞 DNA 甲基化测序技术，如 scBS-seq 和 scRRBS-seq。这些技术以前所未有的分辨率揭示了细胞间的表观遗传模式异质性和动态性。

然而，由于单个细胞中可用的遗传物质较少，单细胞测序面临着一些在批量测序中不存在的挑战。最主要的问题是基因组 CpG 位点的覆盖率低，通常只有1%到20%。此外，每个位点被测序到的读数（reads）也较少，导致甲基化水平的测量充满噪声。因此，开发有效的插补 (imputation) 和去噪 (denoising) 技术对于充分释放单细胞甲基化组分析的潜力至关重要。

近年来，已经有一些计算方法被提出来用于插补单细胞 DNA 甲基化数据。这些方法大多建立在利用 CpG 位点内部（intra-sample）和细胞之间（intercellular）相关性的思想之上。例如：

然而，这些方法各有局限。在此背景下，作者从自然语言处理领域的最新进展中汲取灵感，特别是 BERT 模型及其掩码语言建模 (Masked Language Modeling, MLM) 任务。MLM 的本质是填空，这与数据插补的目标惊人地相似。在语言模型中，Transformer 神经网络因其能够通过自注意力 (self-attention) 机制学习输入中所有单词之间的相互作用而大放异彩。

尽管生物系统可以被优雅地表示为图，但大多数生物信息学中的深度学习实践并未反映这一现实。Transformer 通过自注意力机制模拟了图的结构，能够明确地推理每个输入如何被其他输入所影响。然而，标准的 Transformer 无法直接应用于高维的甲基化矩阵，因为其计算和内存成本与输入数量的平方成正比。

为了解决这个问题，作者引入了 CpG Transformer，它通过结合轴向注意力 (axial attention) 和滑动窗口自注意力 (sliding window self-attention)，将 Transformer 架构成功地应用于部分观测的甲基化矩阵，从而在各种数据集上实现了最先进的插补性能。

2. CpG Transformer 方法

CpG Transformer 的核心是将 Transformer 架构应用于二维的甲基化矩阵，并通过一种类似于协同过滤 (collaborative filtering) 的方式构建输入，其架构创新点在于结合了两种高效的注意力机制。

2.1 模型输入

CpG Transformer 的输入是一个三维张量 $H \in \mathbb{R}^{n \times m \times d_{model}}$，其中 $H_{i,j}$ 代表了甲基化矩阵中第 $i$ 个细胞（行）和第 $j$ 个 CpG 位点（列）的输入表示。这个表示 $H_{i,j}$ 是由三个不同嵌入的拼接经过线性组合得到的：

\[H_{i,j} = W \cdot [h^{CpG}_{i,j}; h^{cell}_i; h^{DNA}_j]\]
  1. CpG 嵌入 ($h^{CpG}_{i,j}$)：将 CpG 位点 $j$ 在细胞 $i$ 中的甲基化状态（未知 ?, 未甲基化 0, 或甲基化 1）嵌入为一个向量。
  2. 细胞嵌入 ($h^{cell}_i$)：为每个细胞索引学习一个独特的嵌入向量，用于区分不同的细胞。
  3. DNA 嵌入 ($h^{DNA}_j$)：将每个 CpG 位点周围1001 bp的 DNA 序列通过一个 CNN 网络（架构与 DeepCpG 类似）进行处理，得到一个代表该位点 DNA 上下文的嵌入向量。

2.2 CpG Transformer 架构

标准的自注意力机制对于一个 $n \times m$ 的矩阵，其计算复杂度高达 $O(n^2 m^2)$，这在处理数百万个 CpG 位点的甲基化数据时是不可行的。为了解决这个问题，CpG Transformer 巧妙地结合了两种注意力机制：

2.2.1 轴向注意力 (Axial Attention)

轴向注意力将二维的注意力计算分解为两个独立的一维注意力操作：

通过这种方式，计算复杂度从 $O(n^2 m^2)$ 大幅降低到 $O(mn(n+m))$。

2.2.2 滑动窗口自注意力 (Sliding Window Self-Attention)

考虑到基因组上邻近的 CpG 位点通常是相关的，作者进一步优化了行注意力。

在执行行注意力时，每个查询（query）只关注其邻近的一个固定窗口（window size $w$）内的键（keys）。这进一步将计算复杂度从 $O(mn(n+m))$ 降低到 $O(mn(n+w))$。

2.2.3 模型结构

CpG Transformer 由四个相同的 Transformer 层堆叠而成。每个层包含三个子层：

  1. 列自注意力层：全连接的注意力，用于捕捉细胞间的相互作用。
  2. 行滑动窗口自注意力层：在一个大小为41的窗口内进行注意力计算，用于捕捉邻近 CpG 位点间的相互作用。
  3. 前馈网络层：一个位置无关的全连接前馈网络。

每个子层周围都使用了残差连接和层归一化。

2.2.4 训练目标

作者借鉴了 BERT 中的掩码语言建模 (Masked Language Modeling, MLM) 思想来设计训练目标。这本质上是一种去噪自编码 (denoising autoencoding) 的过程。

在训练的每个批次（batch）中，随机地将一部分已观测到的位点替换为特殊的 ? 令牌。此外，为了让模型不仅能插补，还能去噪，将另一部分待掩码的位点随机替换为 0 或 1。模型的目标是根据被破坏的输入，恢复这些被掩码或随机化位点的原始甲基化状态。使用交叉熵损失来优化模型。

3. 实验分析

作者在五个公开的单细胞甲基化数据集上，将 CpG Transformer 与 DeepCpG 和 CaMelia 进行了全面的性能比较。

3.1 插补性能

结果显示，在所有五个数据集上，CpG Transformer 在 ROC AUC 和 PR AUC 两个指标上均一致地优于 DeepCpG 和 CaMelia。

消融研究表明，移除了 CpG 嵌入后，模型性能下降最大，表明 CpG 位点间的依赖关系是预测的关键。移除了细胞嵌入后，性能下降次之，凸显了利用细胞异质性的重要性。移除 DNA 嵌入和位置编码的影响相对较小。

作者分析了测序深度和局部稀疏度对性能的影响。CpG 位点被测序到的读数（reads）越多（标签噪声越小），或者其邻域内观测到的位点越多（局部稀疏度越低），模型的预测性能就越高。热图显示，局部稀疏度是对预测性能影响最大的因素。

3.2 模型解释

作者使用了积分梯度 (Integrated Gradients) 方法来解释模型的预测，该方法通过计算预测输出相对于输入特征的梯度来衡量每个输入的贡献度。

3.3 迁移学习 (Transfer learning)

由于 CpG Transformer 的通用性，作者预期它能够很好地迁移到其他数据集。

作者进行了一项实验，将在 MBL 数据集上训练好的模型权重，用于初始化一个在 HCC 数据集上训练的新模型。结果显示，使用迁移学习的模型仅需50个训练步骤就能达到与随机初始化模型在1000个步骤后相当的性能，收敛速度提升了约20倍。

即使在没有任何微调的情况下，迁移过来的模型也能达到95.62%的 ROC AUC，这表明 Transformer 的权重已经学习到了通用的、能够识别细胞身份的表示。

这些结果表明，CpG Transformer 对于计算资源有限的研究人员来说，是一个非常实用且高效的工具。