0. TL; DR

本文介绍了一种名为 HyenaDNA 的基因组基础模型，该模型基于隐式卷积的 Hyena 架构，能够处理长达 100 万个单核苷酸的序列，显著提升了基因组序列建模的上下文长度和分辨率。HyenaDNA 在多个基因组任务上表现出色，包括在单核苷酸分辨率下识别调控元件和在长序列物种分类任务中有效解决问题。此外，HyenaDNA 还展示了在基因组学中首次使用上下文学习的能力，无需更新预训练模型权重即可简单适应新任务。

1. 背景介绍

基因组序列蕴含着大量关于基因调控、蛋白质合成及其他细胞属性的信息。研究人员提出基因组基础模型，旨在从无标签的基因组数据中学习可泛化的特征，进而通过微调用于识别调控元件等下游任务。然而，基于 Transformer 的基因组模型因注意力机制的二次方缩放限制，通常只使用 512 到 4096 个标记作为上下文（不足人类基因组的 0.001%），极大地限制了对 DNA 中长程相互作用的建模能力。此外，这些方法依赖于标记器或固定 k-mer 来聚合有意义的 DNA 单元，丢失了单核苷酸分辨率，而单核苷酸变异（SNPs）可能会完全改变蛋白质功能。

近期，基于隐式卷积的 Hyena 大语言模型在质量上与注意力相匹配，同时允许更长的上下文长度和更低的时间复杂度。本文利用 Hyena 的新长程能力，提出了 HyenaDNA，一个预训练于人类参考基因组的基因组基础模型，其上下文长度可达 100 万个标记，是之前基于密集注意力模型的 500 倍提升。HyenaDNA 在序列长度上呈亚二次方缩放（训练速度比 Transformer 快 160 倍），使用单核苷酸标记，并且在每一层都具有全局上下文。

2. HyenaDNA 模型

Hyena 操作符是 HyenaDNA 的核心组件，它结合了长卷积和逐元素门控层。给定输入 $x \in \mathbb{R}^L$，一个 Hyena 操作符可以定义为：

\[(x_1, x_2, v) \mapsto H(x_1, x_2)v \\ H(x_1, x_2) = D_{x_2} T h D_{x_1}\]

其中 $x_1, x_2, v$ 是输入的投影，$T_h \in \mathbb{R}^{L \times L}$ 是由神经网络输出的可学习的长卷积滤波器，$D_{x_1}, D_{x_2} \in \mathbb{R}^{L \times L}$ 是由 $x_1, x_2$ 构建的矩阵，用于逐元素门控。投影是通过应用密集线性层和短卷积得到的。

HyenaDNA 模型是一个解码器架构，由多个 Hyena 操作符和前馈神经网络堆叠而成。

由于直接在超长序列上训练会影响训练稳定性，本文引入了一种序列长度预热调度器，逐步增加序列长度。该调度器从长度为 64 的序列开始，每轮将窗口长度翻倍，同时保持全局批量大小不变。通过这种方式，每个连续阶段的迭代将包含更多的标记，确保调度器也可以作为批量大小预热的一种形式。

HyenaDNA 采用了一种新型的软提示技术，将可学习的标记（最多 32k）直接注入到输入序列中，从而在不更新预训练模型的情况下实现竞争性的下游结果。这种软提示方法为基因组学中的标准微调提供了一种更简单和更灵活的替代方案。

3. 实验分析

3.1 在人类基因组上的预训练

HyenaDNA 在人类参考基因组上进行预训练，使用下一个核苷酸（标记）预测任务。实验结果表明，随着上下文长度的增加，预训练期间的困惑度（perplexity）有所改善。然而，这需要更多的训练时间和标记。对于较浅的模型，过长的上下文可能导致困惑度恶化（增加）。因此，增加上下文可以作为一种新的正则化维度。

3.2 单核苷酸分辨率

为了评估单核苷酸分辨率性能，HyenaDNA 在多个短序列（<5k 核苷酸）的下游基准测试中进行了标准微调。在 GenomicBenchmarks 数据集上，HyenaDNA 在 7 个数据集上超越了现有 SotA 基线，并在人类增强子识别任务上提高了 20% 的准确率。在 Nucleotide Transformer 的 18 个数据集上，HyenaDNA 使用参数量少得多的模型，在 12 个数据集上达到了 SotA。