论文阅读笔记(八)——Knowledge Distillaton of DNABERT for Prediction of Genomic Elements

论文阅读笔记(八)——Knowledge Distillaton of DNABERT for Prediction of Genomic Elements

摘要

理解人类基因组编码的信息以及 DNA 序列每一部分的影响是我们社会的一个基本问题，这对于揭示常见疾病的机制至关重要。随着基因组学领域最新技术的发展，许多研究机构具备收集海量基因组数据的工具。然而，缺乏能够以生物学上可靠且高效的方式处理和分析这些数据集的工具。

许多深度学习解决方案已经被提出用于解决当前的基因组任务，但大多数时候，主要研究兴趣在于基础的生物学机制，而非预测指标本身的高分数。最近，深度学习的最新技术已转向大型Transformer模型，这些模型使用一种可以利用的注意力机制以实现可解释性。这些大型模型的主要缺点是它们需要大量的内存空间，并且推理时间较长，这可能使其在实际应用中的使用变得不可行。

在这项工作中，我们测试了知识蒸馏的适用性，以获得基因组研究人员可以轻松微调以解决其科学问题的更易用且表现相当的模型。DNABERT 是一种在 DNA 数据上预训练的transformer模型，我们采用两种策略进行蒸馏：DistilBERT 和 MiniLM。获得了四种不同大小的学生模型，并针对启动子识别进行了微调。它们在三个关键方面进行了评估：分类性能、可用性和预测的生物学相关性。通过视觉检查 TATA 启动子预测的注意力图来评估后者，预计这些序列中的 TATA 基序会有一个注意力峰值。

结果表明，确实可以获得在启动子识别任务中表现同样出色且明显更小的模型，在测试的两种技术之间没有显著差异。最小的蒸馏模型在所有评估的性能指标（准确率、F1 分数和 Matthews 相关系数）上下降不到 1%，推理速度提高了 7.3 倍，同时仅具有 DNABERT 参数的 15%。学生模型的注意力图显示，它们成功地学习了模仿 DNABERT 对 DNA 的一般理解。

简介

github

这些模型的一个最大资产和缺点是它们的巨大规模。虽然这使得解决更复杂的任务并获得更好的性能成为可能，但它也需要大量的时间和资源来进行训练和推理，这使得在实际应用中使用变得困难。在基因组学领域，这些问题因数据的性质而被放大。一方面，隐私保护规定导致注释基因组数据的稀缺性，限制了可以训练的模型规模。另一方面，一些基因组分析（例如识别遗传变异）需要处理数百万数据点，使得低推理时间成为模型可用性的关键特性。知识蒸馏（KD）是一组技术，通过训练较小且更可用的模型（称为学生模型）以执行与较大模型（称为教师模型）同样的任务。KD 在自然语言处理和计算机视觉领域展示了很好的效果，并且在缓解当前基因组学领域先进模型的问题方面展示了很大的潜力。

目的和研究问题

本论文的目的是验证在 DNABERT 上使用知识蒸馏（KD）的适用性，以获得显著更小且更可用的Transformer模型。将实现两种不同的 KD 技术，并在特定的基因组任务（启动子识别）上测试所得模型。这是最著名的基因组任务之一，可以利用已知的生物学机制来评估模型预测的生物学相关性。将评估和比较模型的性能、可解释性和可用性，并将原始 DNABERT 作为本项目的基准。

研究问题：

• RQ1：通过 KD 在 DNABERT 上训练的小型Transformer模型在启动子识别中的表现如何（评估准确率、F1 分数和 Matthews 相关系数）？
• RQ2：DNABERT 的注意力映射策略在多大程度上可以用作蒸馏模型预测的生物学解释工具？
• RQ3：使用不同技术蒸馏的模型在性能、可解释性和可用性方面有多大差异？

Method

研究过程的第一步是深入了解基因组学领域，这是一个具有独特概念和方法的特殊领域。一旦确定了项目范围并选择了基准模型 DNABERT，下一步是决定评估学生模型的具体基因组任务。选择的任务是启动子识别，这是一个已知一些基础生物学机制的基因组问题，可以评估模型预测的可解释性。接下来，获取并预处理数据，以满足任务和模型的要求。获取预训练的 DNABERT 模型，并对其进行微调以进行启动子识别任务。之后，实施了第一种知识蒸馏（KD）方法 DistilBERT。评估其性能以识别方法的不足，并选择具有克服这些不足潜力的替代方法。然后，实施了 MiniLM。最后，评估所有所得模型的性能、可解释性和可用性，并相互比较和与教师模型比较。图 2.1 显示了所描述方法的概述。

数据

在处理深度学习模型时，通常可以识别出两个不同的训练阶段。第一个是预训练阶段，模型学习输入数据的一般结构，在本例中是 DNA 序列不同部分之间的关系。通常，通过训练蒙面语言建模（MLM）任务来实现这一点，该任务包括输入一个序列并隐藏其中的一些部分，模型需要猜测这些部分。然后可以微调模型以解决特定任务，如启动子识别。使用预训练模型可以在微调阶段获得更好的性能，同时减少训练时间和数据量。这两个训练阶段需要不同的数据集，下面进行描述。

预训练数据集

在此阶段，目标是让模型从整体上理解基因组。为了构建数据集，从 UCSC 基因组浏览器 [11] 获得了人类参考基因组 GRCh38.p13 装配版本，以 fasta 格式提供（详见附录 A.1.4 了解更多关于参考基因组的详细信息）。

整个基因组被切割成没有重叠的部分，其中 50% 的序列固定长度为 510 个碱基对（bp），其余 50% 的序列长度从 10 到 510 bp 不等（见第 2.2.3 节）。检查并丢弃了含有核苷酸空缺（在参考基因组中为 “N”）的序列。遵循基因组学领域的最新方法，染色体 7 和 8 被保留用于验证和测试目的。最终训练数据集包含 6,353,024 个序列。

微调数据集

在预训练阶段之后，模型被微调以识别输入序列中的启动子。所有包含启动子的已知人类基因组序列均来自真核启动子数据库（EPDnew） [12]，使用其在线选择工具提取。作为正样本的序列紧邻转录起始位点（TSS）上游，提取了从 -249 到 +50 的区域，总长度为 300 bp。数据集包括 TATA 和非 TATA 启动子（详见附录 A.1.2 了解更多关于启动子的详细信息）。

EPDnew 是一个专门包含启动子 DNA 序列的数据库。然而，为了能够区分这种调控元件，微调数据集需要包括负样本，即不包含启动子的序列。仅使用随机 DNA 序列作为负样本是不够的，因为它们具有不同的结构，模型可以轻易区分它们而无需理解上下文 [13]。例如，对于所谓的 TATA 启动子，已知 TATA 序列大约位于 TSS 上游 25 bp 处。TATA 负样本通过提取参考基因组中不包含 “TATA” 基序的非启动子区域构建，这些非启动子区域在相同相对位置上不包含 TATA 基序。对于非 TATA 负样本，使用随机替换以避免存在偏差。图 2.2 说明了用于获取负非 TATA 启动子的方法。

和处理。接着，在第 2.3 节中描述了教师模型，随后在第 2.4 节中描述了学生模型的训练和知识蒸馏（KD）技术。最后，第 2.5 节详细介绍了考虑的评估指标。需要提醒读者的是，附录 A 包含了理解本章方法所需的所有背景知识。

所得数据集在正负序列方面完全平衡（表 2.1）。数据随机分为训练集、验证集和测试集，分别包含 80%、10% 和 10% 的序列。

启动子类型	正样本序列数	负样本序列数
TATA	3065	3065
非 TATA	26533	26533

2.2.3 数据格式化

本研究中使用的所有模型都是 BERT 架构的变种。这些模型期望输入是经过标记化的字符串序列（详见附录 A.3.1.1 了解更多关于标记化的详细信息）。采用的标记化策略是 DNABERT 提出的，处理 DNA 序列为具有 1 个滑动窗口的 k-元组（图 2.3）。生成的词汇表大小为 4^k 加上特殊标记。BERT 模型期望输入的第一个元素是特殊分类标记 [CLS]，并在句子末尾有一个分隔标记 [SEP]。由于输入大小为 512，第一个和最后一个标记是固定的，因此输入序列的最大长度可以是 510。

Teacher Model

在所有蒸馏实验中选择的教师模型是 Ji 等人提出的 DNABERT [10]（图 2.4）。这是一种基于Transformer的模型，在 DNA 序列上进行了预训练，在包括启动子识别在内的多个基因组任务中表现出色。它还引入了一种利用注意力分数发现新基序的新方法（详见附录 A.1.3 了解更多关于基序的详细信息）。

DNABERT 遵循 BERTbase 结构：12 个Transformer层，每层有 768 个隐藏单元和 12 个注意力头。作者提供了不同版本的预训练 DNABERT 模型，接受 3-元组、4-元组、5-元组和 6-元组标记化的输入序列。然而，他们仅报告了 6-元组模型的结果，因为他们声称这是表现最好的模型。不同的模型都进行了启动子分类任务的微调。对注意力图的视觉评估支持了 6-元组模型是最合适的选择，因此选择此版本进行所有蒸馏。

2.4 学生模型

可以蒸馏不同类型的知识，这些知识在学生模型训练期间通过特定的误差函数进行建模。在本研究中，测试了两种 KD 技术。它们都被认为是离线方法，因为使用的是预训练教师。第一种是 Sanh 等人提出的 DistilBERT，这是一种基于响应的蒸馏算法，它是第一种用于蒸馏 BERT 模型的技术，并且仍然是最流行的技术之一。然而，它有一个主要缺点，即对学生模型的大小有约束。考虑了几种替代方法，即 Jiao 等人的 TinyBERT 和 Wang 等人的 MiniLM。这两种策略都是基于特征的方法，其中学生关注注意力权重，这似乎是确保可解释性的有前途的特性。与 DistilBERT 相比，它们在接受的学生架构方面也更具灵活性。然而，TinyBERT 需要更复杂的双重蒸馏训练方案，并且在自然语言任务中报告的性能低于 MiniLM。因此，选择 MiniLM 作为替代方法。表 2.2 显示了两种蒸馏技术的关键要素概述。

使用两种 KD 技术训练了不同维度的模型，如图 2.5 所示。学生模型具有与教师模型相同的架构，但层数减少且隐藏空间较小。

技术	蒸馏损失	初始化	约束
DistilBERT	软目标概率分布	嵌入输出	隐藏状态
MiniLM	自注意力分布	自注意力值关系	无

DistilBERT

在此设置中，学生模型试图同时学习如何解决预训练任务以及如何模仿教师的行为。因此，学生的训练目标是三重的（图 2.6）：

• 解决一个特定任务，在本例中是蒙面语言建模（MLM 损失）。
• 匹配教师的概率分布，用交叉熵函数建模（蒸馏损失）。
• 对齐最后一层的嵌入与教师的嵌入（余弦损失）。

这种训练目标由一个损失函数建模，该损失函数是上述三种损失的线性组合，其中 ${\alpha }_{mlm}$、${\alpha }_{dist}$ 和 ${\alpha }_{cos}$ 是可调超参数。

$$L={\alpha }_{mlm}{L}_{mlm}+{\alpha }_{dist}{T}^2{L}_{dist}+{\alpha }_{cos}{L}_{cos}$$

值得一提的是，蒸馏损失中使用的教师和学生的概率分布是通过 Hinton 等人引入 KD 概念时所指出的 softmax 温度函数计算的。

使用这种技术，学生模型的架构受限于与教师模型具有相同的隐藏维度。因此，DistilBERT 被配置为具有与 DNABERT 相同的结构，但层数减半。因此，DistilBERT 有 6 个变压器层，每层有 768 个隐藏单元和 12 个注意力头。模型从教师模型初始化，每隔一层取一次，使用每一个奇数层。

由于训练时间长，不可行进行超参数搜索，因此训练参数按照 DistilBERT 和 DNABERT 论文的指导原则进行定义。学习率设置为 4e-4，前 150k 训练步骤内线性预热，然后缓慢线性下降。批量大小设置为 64，并使用梯度累积步骤模拟非常大的批量（3200 个样本）。MLM、蒸馏和余弦损失的权重分别选择为 2、7 和 1。使用 Adam 优化器，权重衰减为 0.01。当损失曲线达到平台期后，训练停止，共进行两个完整的训练周期。

然后，模型使用两种不同的范式进行启动子识别任务的训练。第一种策略是常规微调。第二种策略是在微调任务的基础上进行新的蒸馏步骤，其中教师模型是微调后的 DNABERT。此附加蒸馏的损失函数是上述蒸馏损失和启动子识别任务损失（交叉熵）的线性组合，如下所示：

$$L={\alpha }_{bce}{L}_{bce}+{\alpha }_{dist}{T}^2{L}_{dist}$$

MiniLM

这种技术被认为是任务无关的，因为目标函数不包括监督训练损失（例如 DistilBERT 中的 MLM）。相反，学生试图模仿教师模型最后一层的自注意力行为（详见附录 A.3.1.3 了解更多关于自注意力的详细信息）。该目标由以下公式建模，其中 $d_k$ 是自注意力头的维度，$Q$、$K$、$V$ 分别代表查询、键和值（图 2.7）。

• 自注意力分布（AT）：查询（$Q$）和键（$K$）的缩放点积。

$$AT=\text{softmax}\left(\frac{Q{K}^T}{\sqrt{d_k}}\right)$$

• 自注意力值关系（VR）：值（$V$）之间的缩放点积。

$$VR=\text{softmax}\left(\frac{V{V}^T}{\sqrt{d_k}}\right)$$

计算教师和学生之间的 Kullback–Leibler（KL）散度，并平均所有注意力头的自注意力分布和值关系。最终损失是这两个值的总和。

设置

所有模型均基于 HuggingFace 提供的框架构建。DNABERT、DistilBERT 的作者提供的代码以及 MiniLM 的非官方重新实现[17]的一部分代码作为基础，并进行了扩展以满足具体需求。所有学生模型使用一台具有 48GB VRAM 的 NVIDIA Quadro RTX 8000 GPU 进行蒸馏，使用一台具有 25GB VRAM 的 RTX TITAN GPU 进行微调。

生成的模型可以在常规 CPU 上轻松进行推理，无需特殊硬件。通过限制批量大小，蒸馏模型也可以在常规 CPU 上进行微调，代价是增加训练时间。

评估

蒸馏模型在下游任务的性能、注意力权重和可用性方面进行了评估和比较。DNABERT（教师模型）作为基准进行比较。

性能指标

所有模型均在启动子识别任务上进行了微调，这是一个二元分类问题。按照 DNABERT 的建议，使用三种不同的指标评估模型的性能。这些指标根据真阳性（TP）、真阴性（TN）、假阳性（FP）和假阴性（FN）定义。

1. 准确率。表示正确预测的比例相对于样本总数的比例。这对于给定的数据集完全平衡的情况下是一个合适的指标。其值范围从 0 到 1，后者表示完美的性能。

$$\text{Acc}=\frac{TP+TN}{TP+FP+TN+FN}$$

2. F1 分数。这是一个更复杂的指标，考虑了召回率和精度之间的平衡。即，找到正样本的能力和正确分类正样本的比例。值同样范围从 0 到 1。

$$\text{F1}=\frac{2\cdot TP}{2\cdot TP+FP+FN}$$

3. Matthews 相关系数（MCC）。统计测量，如果样本中正类和负类均正确分类，则得分较高。在这种情况下，得分范围从 -1 到 1，0 表示随机猜测。

$$\text{MCC}=\frac{TP\cdot TN-FP\cdot FN}{\sqrt{(TP+FP)(TP+FN)(TN+FP)(TN+FN)}}$$

注意力可视化

尽管不能仅依靠自注意力分数进行模型解释，但它们可以用于了解模型在分类序列时关注输入的哪些区域。DNABERT 的作者描述了一种用于核苷酸级别注意力评分的方法。此外，他们通过使用评分成功识别输入序列中的生物相关基序来验证该方法的相关性。

该方法包括计算每个 k-元组标记相对于分类标记 [CLS] 的注意力权重，将每个注意力头的权重相加。然后，通过平均包含给定核苷酸的所有 k-元组的值，将 k-元组分数转换为核苷酸分数。结果是每个核苷酸的一个值，表示注意力，可以轻松地可视化为热图。

对于大多数基因组任务，很难确定模型是否关注 DNA 序列的生物相关部分，因为与分类任务相关的具体基序是未知的。在 TATA 启动子中，已知 “TATA” 基序大约出现在 TSS 上游 25 bp 处。因此，预计模型将特别关注输入的这些区域。使用 TATA 启动子的注意力热图来可视化和主观评估模型生成生物学可解释预测的能力。

模型可用性

执行蒸馏的目标是获得显著更小的模型，可以更快地训练并具有较短的推理时间。即，无需特殊硬件或不合理的长预测时间即可在实际应用中使用的模型。在这方面，评估了参数数量和推理时间，并与教师模型 DNABERT 进行比较。还监控了每个学生模型的蒸馏和微调时间。

3. 结果

性能指标

表 3.1 显示了在启动子识别任务上训练的不同学生模型的性能指标。报告的值是五次使用不同种子进行微调运行的平均结果。为了便于比较，表中还包括仅在下游任务（即无预训练）上训练的模型的性能。此外，由于启动子数据集是平衡的，因此“愚蠢”分类器的准确率为 50%。为了便于视觉解释，将这些结果绘制在图 3.1 中。

模型	准确率 (%)	F1 分数 (%)	MCC (%)
DNABERT	97.0 ± 0.3	97.0 ± 0.3	94.1 ± 0.5
DistilBERT	96.9 ± 0.2	96.9 ± 0.2	94.0 ± 0.4
DistilBERT (D)	96.7 ± 0.1	96.7 ± 0.1	93.6 ± 0.2
MiniLM	97.1 ± 0.2	97.1 ± 0.2	94.6 ± 0.4
MiniLM small	96.7 ± 0.1	96.7 ± 0.1	93.4 ± 0.1
MiniLM mini	96.5 ± 0.1	96.6 ± 0.1	93.2 ± 0.1
无预训练	92.9 ± 0.1	92.9 ± 0.1	85.8 ± 0.2

可以看出，所有蒸馏模型在性能上非常相似，相互之间以及与教师模型相比。它们的准确率、F1 分数和 MCC 分别为 0.97、0.97 和 0.94。相比之下，直接在启动子任务上训练的模型在所有三个指标上的表现较低：分别为 0.93、0.93 和 0.86。

注意力可视化

图 3.2 显示了不同训练模型在分类含有 TATA 启动子的序列时的注意力景观。需要注意的是，不同模型的注意力强度值略有不同。这种可视化选择的原因是我们不关注绝对注意力值，而是关注模型在输入序列中更加关注的部分。

DNABERT（教师模型）一致地在 TSS 上游 25 bp（坐标 -25）附近给出更多注意力，同时对序列的其余部分给予较低但一致的注意力。所有蒸馏模型显示出类似的注意力景观，在 TATA 框周围具有较高的值。

比较两个 DistilBERT 模型（图 3.2(b) 和 ©），使用附加蒸馏步骤训练的模型相比常规微调策略训练的模型，在 TATA 区域有更高的关注度。对于相同架构但使用不同蒸馏技术的模型（图 3.2(b)/© 和 (d)），MiniLM 显示出略微分散的注意力景观。如果我们关注 MiniLM 家族的模型（图 3.2(d) 到 (e)），随着模型大小的减少，景观变得更加陡峭，输入序列的不同部分出现了高峰。特别是，迷你版表现出较少定义的注意力焦点，在序列的开头有显著的峰值。最后，仅为启动子识别任务训练的模型（无预训练）没有显示出任何特定的注意力模式。

3.3 模型可用性

表 3.2 提供了蒸馏模型在大小、推理时间和训练时间方面的可用性概述。报告的推理时间对应于在包含 185 个批次、每批次 32 个序列的完整测试数据集上进行启动子识别的平均运行时间。需要注意的是，为了公平比较，所有模型的训练设置（主要是 GPU 容量、批量大小和训练期间执行验证的频率）保持不变。

对于教师模型 DNABERT，使用了已经预训练的模型。然而，值得一提的是，作者报告了在 8 台 NVIDIA 2080Ti GPU 上预训练时间为 25 天。

可以看出，推理和微调时间随着模型大小的减小而减少。除了这一点，对于具有相同架构的模型，DistilBERT 使用的蒸馏技术比 MiniLM 使用的时间更少。专注于 DistilBERT，当训练下游任务时，常规微调比第二次蒸馏步骤更快。

模型	蒸馏时间	微调时间	推理时间	参数数量
DNABERT	44s (1.0x)	-	-	110M (1)
DistilBERT (D)	22s (2.0x)	36.7h (2)	0.37h (3)	66M (0.60)
DistilBERT	22s (2.0x)	36.7h (2)	0.20h (3)	66M (0.60)
MiniLM	22s (2.0x)	42.4h (2)	0.20h (3)	66M (0.60)
MiniLM small	13s (3.4x)	42.6h (4)	0.12h (3)	22M (0.20)
MiniLM mini	6s (7.3x)	39.0h (4)	0.11h (3)	17M (0.15)

讨论

知识蒸馏的影响

评估结果表明，所有蒸馏的学生模型在启动子识别任务中表现大致相同。即使是 MiniLM mini，尽管其参数仅为 DNABERT 的 15%，但仍然表现良好（表 3.1）。相比之下，对于未经任何预训练方案训练的模型，F1 分数下降了 4%，MCC 下降了 8%。这些结果支持了这样一种观点，即可能不需要处理使用像 DNABERT 这样大的模型的不便，因为如果适当训练，较小的模型也可以获得类似的性能。值得一提的是，每当层数或隐藏空间减半，推理时间也会减半。最小的模型具有 DNABERT 总参数的八分之一，其推理时间提高了 7.3 倍（表 3.2）。

尽管获得了很好的性能结果，但在基因组学领域，大多数时候研究重点在于驱动调控元素的基础生物学机制，而不是对这些元素的识别本身。例如，基因组学家可能更感兴趣的是发现启动子序列中一致存在的模式，而不是正确识别它们。不同模型注意力图的可视化清楚地显示了这一点。在所有蒸馏模型中，注意力权重的模式非常接近 DNABERT，在 TSS 上游 25 bp 附近显示出一致的注意力峰值（图 3.2）。这一发现表明，尽管学生模型的大小不同，但它们仍能有效捕捉到与教师模型相同的生物学相关性。这进一步验证了知识蒸馏在保持模型解释性方面的有效性。

蒸馏技术的比较

虽然所有蒸馏技术在性能上都相当接近，但在可用性和训练效率方面存在显著差异。DistilBERT 由于其简单的蒸馏过程，比 MiniLM 更快。然而，MiniLM 提供了更大的灵活性，可以在不牺牲性能的情况下设计更小的学生模型。

DistilBERT (D) 在使用附加蒸馏步骤后，显示出在生物学相关性的注意力分数方面的微小改进。这一发现表明，在某些情况下，额外的蒸馏步骤可能有助于进一步提高模型的解释性。

模型可用性的影响

显著减小模型大小和推理时间是实现基因组学研究中的实时应用的关键。在这方面，所有学生模型在不显著损失性能的情况下，都显示出了很好的改进。特别是 MiniLM mini，在推理时间和参数数量方面的改进最为显著（表 3.2）。

未来的研究方向

虽然本研究表明知识蒸馏可以显著减小模型大小并提高可用性，但仍有一些未来的研究方向值得探索。例如，可以进一步优化蒸馏过程，以实现更好的性能和解释性。此外，还可以探索其他基因组任务，如基因预测或调控元件识别，以评估知识蒸馏在更广泛的基因组学应用中的有效性。

5. 结论

在本研究中，我们验证了在 DNABERT 上使用知识蒸馏（KD）技术的可行性，并展示了通过两种不同的 KD 策略获得的学生模型在启动子识别任务中的出色表现。我们的结果表明，通过 KD 技术，可以获得显著更小且更可用的变压器模型，同时保持其高性能和生物学相关性。这为基因组学领域提供了更高效的工具，有助于更好地理解和分析 DNA 序列。

未来工作。未来的研究可以进一步优化 KD 策略，探索其他基因组任务中的应用，并结合其他深度学习技术以提高模型的性能和可解释性。