iDNA-DAPHA:基于领域自适应预训练与层次注意力的甲基化预测通用框架.

0. TL; DR

iDNA-DAPHA是一个准确且通用的两阶段DNA 甲基化识别框架,核心创新在于其领域自适应预训练 (Domain-Adaptive Pretraining, DAP) 策略,该策略通过特征对齐来学习来自多个物种的各种类型甲基化序列的共同特征。随后,通过在特定任务上进行微调来捕捉任务特异性特征。此外,该框架还引入了层次化注意力 (Hierarchical Attention, HA) 机制来增强其表示能力。

在覆盖三种代表性 DNA 甲基化类型的17个基准数据集上,iDNA-DAPHA 的性能优于现有的先进方法。特别是在训练样本有限的场景下,其优势尤为明显。通过可视化的注意力分析,模型能够捕捉到保守的序列模式,并学习到更具判别力的表示,为理解甲基化机制提供了新的视角。

1. 背景介绍

DNA 甲基化是一种影响基因表达的表观遗传机制,在多种医学疾病(如癌症)的发生和发展中扮演着重要角色。目前,三种代表性的 DNA 甲基化类型——N4-甲基胞嘧啶 (4mC)5-羟甲基胞嘧啶 (5hmC)N6-甲基腺嘌呤 (6mA)——已在多种物种中被发现,并与不同的功能机制相关。准确识别这些甲基化位点对于分析和揭示其功能机制至关重要。

然而,由于环境因素、衰老等影响,DNA 甲基化会发生动态变化,这给准确预测带来了挑战。传统的湿实验技术,如简化代表性亚硫酸氢盐测序 (RRBS),成本高昂且耗时。因此,计算预测方法成为一个很好的选择,它不仅能降低成本,还能为实验研究提供有价值的指导。

近年来,机器学习和深度学习方法在 DNA 甲基化预测方面取得了显著成效。但许多现有模型要么是针对特定修饰或物种设计的,泛化能力有限;要么是通用的预测器,但依赖于手动设计的特征(如 iDNA-MS),适应性不强。

为了解决这个问题,一系列基于深度学习的通用预测器(如 iDNA-ABT, iDNA-ABF, MuLan-Methyl)被提出。它们大多采用基于 BERT 及其变体的预训练-微调范式,在通用的大规模无标签序列上进行预训练,然后在特定任务上进行微调。然而,这些方法往往忽略了不同物种甲基化数据之间可能存在的共享特征。同时,它们在建模长程依赖时只考虑了单粒度的令牌到令牌 (token-to-token) 注意力,忽略了更宏观的组到组 (group-to-group) 相关性,限制了模型的表示能力。

基于以上考虑,作者提出了一个名为 iDNA-DAPHA 的通用两阶段甲基化预测框架。该框架利用领域自适应预训练 (DAP) 来学习跨物种、跨甲基化类型的共享特征,并引入层次化注意力 (HA) 机制来增强模型的表示能力。

2. iDNA-DAPHA方法

iDNA-DAPHA 的核心是一个两阶段的深度学习框架,它巧妙地结合了领域自适应预训练 (DAP)层次化注意力 (HA)

2.1 基准数据集

作者使用了与之前多个通用预测器相同的基准数据集,以便进行公平比较。该数据集包含 4mC5hmC6mA 三种甲基化类型,共覆盖12个物种的17个不同数据集。每个数据集都被随机划分为1:1的训练集和测试集。

2.2 iDNA-DAPHA 框架

iDNA-DAPHA 的工作流程如图所示,主要分为两个阶段:

2.2.1 第一阶段:甲基化领域自适应预训练 (DAP)

DAP 的目标是让模型学习到跨越不同物种和甲基化类型的通用特征。

作者将17个不同甲基化数据集的训练集合并,构建了一个大规模的 DAP 数据集。这个数据集包含了丰富的、与各种甲基化位点相关的序列信息,并为来自同一原始数据集的序列分配相同的任务标签。

特征提取模块

特征提取模块模块采用双流架构 (dual-stream architecture),并融合了层次化注意力 (HA) 机制。借鉴 iDNA-ABF 的思想,作者使用 3-mers6-mers 两种尺度来对输入的 DNA 序列进行令牌化,并分别送入两个并行的子网络(B_upB_down)。

将两个子网络提取的特征进行拼接,得到融合后的特征表示。

\[F(S) = [B_{up}(S_{3-mers}), B_{down}(S_{6-mers})]\]

其中,$B_{up}$ 和 $B_{down}$ 都是包含了 HA 机制的 DNABERT 模型。

特征对齐模块

为了学习共享特征,作者引入了一个特征对齐模块,该模块通过对抗训练 (adversarial training) 的思想,使得不同任务的特征分布尽可能相似。任务分类器用于区分输入特征来自哪个甲基化任务(共17个)的分类器。

梯度反转层 (Gradient Reversal Layer, GRL):在特征提取模块和任务分类器之间,该层在前向传播时是一个恒等变换,但在反向传播时,它会将来自后续层的梯度乘以一个负的标量 $-\lambda$。

\[G_{\lambda}(x) = x, \quad \frac{dG_{\lambda}}{dx} = -\lambda E\]

在训练过程中,任务分类器会努力最小化任务分类损失 $\mathcal{L}_{fa}$,以更好地区分不同任务的特征。而特征提取模块由于接收到的是被反转的梯度,其参数更新方向会与任务分类器相反,即最大化任务分类损失。这个对抗过程会迫使特征提取模块学习到那些在不同任务间不可区分的、通用的共享特征。

甲基化预测模块

与特征对齐模块并行的是一个标准的甲基化预测模块,它负责预测输入序列是否被甲基化。其损失函数为标准的交叉熵损失 $\mathcal{L}_{mp}$。

总体优化目标

DAP 阶段的总体损失函数是任务分类损失和甲基化预测损失之和:

\[\mathcal{L}_{total} = \mathcal{L}_{fa} + \mathcal{L}_{mp}\]

2.2.2 第二阶段:任务特异性微调 (Fine-tuning)

DAP 阶段获得了一个具有良好通用特征表示能力的预训练模型后,作者在每个具体的甲基化预测任务上对其进行微调。微调阶段的流程与第一阶段类似,但有两点主要区别:

  1. 移除了特征对齐模块,因为此时的目标是学习任务特异性的判别特征。
  2. 只使用目标任务的数据集进行训练。

2.3 层次化注意力 (Hierarchical Attention, HA)

传统的 Q-K-V 自注意力只捕捉了单个令牌之间的相关性。为了克服这一局限,作者引入了 HA 机制,它在每个 Transformer 编码器层内同时建模令牌到令牌 (token-to-token)组到组 (group-to-group) 的相关性。

通过一个可学习的聚合操作 Agg(X)(使用深度可分离卷积实现),将 Q, K, V 向量的局部邻域聚合成组代理 $Q’, K’, V’$。在这些组代理上执行标准的自注意力计算,以捕捉组间的依赖关系。

\[\text{Attn}(Q', K', V') = \text{softmax}\left(\frac{\text{Agg}(Q)\text{Agg}(K)^T}{\sqrt{d_k}}\right)\text{Agg}(V)\]

最终,将传统的令牌到令牌注意力和组到组注意力的输出通过一个线性层进行融合,得到 HA 的最终输出。

\[\text{HA-Attn}(Q, K, V) = F_{linear}[\text{Attn}(Q, K, V), \text{Attn}(Q', K', V')]\]

通过这种方式,HA 机制增强了模型的整体表示能力。

3. 实验分析

作者通过广泛的实验,验证了 iDNA-DAPHA 及其关键组件的有效性。

3.1 与现有方法的性能比较

作者在17个基准数据集上将 iDNA-DAPHA 与多种先进的通用预测器(iDNA-MS, iDNA-ABT, iDNA-ABF, MuLan-Methyl, StableDNAm)进行了比较。

ACC 指标上,iDNA-DAPHA 在所有17个数据集上均取得了最优性能。在 MCC 指标上,它在16个数据集上排名第一,在剩下的一个数据集上与最优性能相当。在 AUC 指标上,它在13个数据集上排名第一。

特别是在训练样本最少的两个数据集 D1 (4mC_C.equisetifolia)D13 (6mA_R.chinensis) 上,iDNA-DAPHA 的性能提升尤为显著,ACC 分别提升了3.82%和1.82%,MCC 分别提升了9.27%和4.01%。

这些结果充分证明了 iDNA-DAPHA 的卓越性能和可靠性,尤其是在小样本场景下。

3.2 消融研究

作者通过消融实验,探究了 DAPHA 这两个核心组件的贡献。

这些结果表明,DAPHA 都是有效的设计,其中 DAP 的贡献更为突出。

3.3 双尺度 vs. 四尺度 k-mer 表示策略的评估

作者比较了使用双尺度(3-mer, 6-mer)和四尺度(3, 4, 5, 6-merk-mer 表示策略的性能。结果显示,四尺度策略虽然能带来轻微的性能提升,但其计算成本(参数量、训练和推理时间、GPU 内存)几乎翻倍。因此,考虑到性能和效率的权衡,双尺度策略是更实用和高效的选择。

3.4 通过基序相似性分析和注意力可视化来解释 iDNA-DAPHA

作者使用 iDNA-DAPHA 的注意力机制来提取序列中的重要区域,并与已建立的基序发现工具 STREME 的结果进行比较。

注意力热图的可视化显示,模型的高注意力区域与这些已识别的基序高度吻合,增强了模型的可解释性。在三个不同物种的数据集上,iDNA-DAPHA 发现的基序与 STREME 发现的基序高度相似,P-value 都非常显著。这证明了模型发现的模式是具有生物学意义的。

3.5 领域自适应预训练对 DNA 甲基化预测的影响

为了直观地展示 DAP 的效果,作者使用 t-SNE 对通用预训练模型和经过 DAP 的模型所学习到的特征表示进行了可视化。

在三个数据集上,通用预训练模型的特征空间中,正负样本的聚类都比较分散和重叠。经过 DAP 之后,模型学习到的特征表示形成了更紧凑、边界更清晰的聚类,正负样本被很好地区分开来。

这个对比清晰地表明,DAP 能够引导模型学习到更具判别力的特征表示,从而提升了下游任务的预测准确性。