HybriDNA：一种混合Transformer-Mamba2的长序列DNA语言模型.
0. TL; DR
本文介绍了一种名为 HybriDNA 的新型 DNA 语言模型，它采用混合 Transformer-Mamba2 架构，能够高效处理长达 131kb 的 DNA 序列并保留单核苷酸分辨率。
HybriDNA 在 33 个 DNA 理解数据集上取得了 SOTA 性能，并能生成具有特定功能的合成 cis-调控元件。实验表明，随着模型参数从 300M 增加到 3B 和 7B，性能持续提升，符合扩展规律。
HybriDNA 的提出为 DNA 研究和应用提供了有力工具，有助于深入理解 “生命语言” 并推动基因治疗和生物技术发展。
1. 背景介绍
DNA 作为生命遗传密码，其序列中单个核苷酸的差异可能对生物功能产生重大影响。对 DNA 进行建模面临两大挑战：一是如何高效处理超长 DNA 序列并保留单核苷酸分辨率；二是如何在生成任务和理解任务上均表现出色。
目前的 DNA 基础模型主要分为两类：仅编码器的 Transformer 架构（如 DNABERT2）在理解任务上表现良好，但在生成任务上有局限性；仅解码器的架构（如 HyenaDNA）在生成任务上有潜力，但在理解任务上不如编码器模型。
为应对上述挑战，作者提出了 HybriDNA，一种基于混合 Transformer-Mamba2 架构的仅解码器 DNA 语言模型。该模型结合了注意力机制和选择性状态空间模型的优势，既能高效处理长序列，又能关注序列中的细节信息。
2. HybriDNA 模型
2.1 模型结构
HybriDNA 采用仅解码器的序列到序列架构，将 HybriDNA Mamba2 块和 HybriDNA Transformer 块按 7:1 的比例交替排列。
Mamba2 块的核心是状态空间对偶性（SSD）层，其处理输入序列的递推公式为：
\[h_t = A_t h_{t-1} + B_t x_t \\ y_t = C_t^T h_t\]其中，$h_t$ 表示隐藏状态，$A_t$ 是状态转移矩阵，$x_t$ 是输入，$B_t$ 将输入映射到隐藏空间，$C_t$ 将隐藏状态映射到输出 $y_t$。SSD 层通过将 $A_t$ 简化为标量乘以单位矩阵（$A_t = a_t I$）来提高效率：
\[h_t = a_t h_{t-1} + B_t x_t\]对于多维输入，SSD 层采用多头设计，每个头独立处理输入的不同子集，类似于 Transformer 中的多头注意力机制。
Transformer 块则采用标准的 Transformer 解码器块，利用自注意力机制捕获序列中每个位置与其他位置之间的关系。
2.2 预训练阶段
HybriDNA 在大规模多物种基因组数据上进行预训练，使用下一个核苷酸（token）预测（NTP）目标。预训练数据集包含 845 个物种，共 1600 亿个核苷酸。
采用基础级分词策略，将每个核苷酸（A、C、T、G）编码为一个单独的 token，以保留基因组数据的自然结构并进行细致的解释和特征提取。
为增强模型对长基因组范围的泛化能力，预训练阶段采用多阶段预热过程。从 8192 token 的上下文长度开始训练，然后逐步增加到 32768 token 和 131072 token。
2.3 下游微调阶段
针对理解任务，HybriDNA 引入了回声嵌入（echo embedding）技术。其核心思想是重复序列有助于将后续元素的上下文信息编码到嵌入中。具体操作是将输入序列 $x$（如AACG） 复制一遍形成 “回声” 输入 $x_{echo}$（如AACGAACG），然后提取序列后半部分的嵌入，并将其输入分类头生成预测概率分布。
针对生成任务，HybriDNA 引入特定的提示 token 来编码任务特定指令。这些提示 token 被添加到现有的核苷酸词汇表中，并在嵌入层中随机初始化，嵌入层扩展以包含新添加的 token ID。HybriDNA 自回归地预测每个核苷酸 token $x_t$，条件是所有先前的提示 token，这些 token 指定任务特定要求。
3. 实验分析
作者训练了 3 个不同参数规模的 HybriDNA 模型（3 亿、30 亿和 70 亿参数），所有模型均采用一致的预训练策略。随着模型规模的增大，训练损失和验证损失均呈下降趋势，表明更大规模的模型在捕获复杂的基因组模式方面具有优势。
作者用 HybriDNA 与其他 5 个最先进的 DNA 基础模型进行了对比，包括 NT-500M-human、NT-2.5B-MS、DNABERT-2、HyenaDNA-medium-160k 和 Caduceus-Ph131k。
3.1 短距离理解基准测试
3.1.1 GUE 基准测试
GUE 基准测试包含 28 个数据集，涵盖 9 个任务，输入长度从 70 到 512 个碱基对。HybriDNA 在多个任务上表现出色，例如在人类基因组的剪接位点检测（SS）任务中，HybriDNA-7B(E) 达到了 90.12 的 MCC 值。
3.1.2 BEND 基准测试
BEND 基准测试评估了染色质可及性预测、组蛋白修饰预测和 CpG 甲基化预测任务。HybriDNA-7B 在染色质可及性预测任务上达到了 0.84 的 AUROC 值，在 CpG 甲基化预测任务上达到了 0.93 的 AUROC 值。
3.2 长距离基因组基准测试
LRB 基准测试评估模型对长距离基因组范围内依赖关系的理解能力。在因果 eQTL 变异效应预测任务中，HybriDNA-300M 模型在不同上下文长度下均表现出较好的性能，随着上下文长度从 8k 增加到 131k，AUROC 值从 0.71 提升到 0.74。
3.3 合成 cis-调控元件设计
3.3.1 人类增强子生成
HybriDNA 在人类增强子生成任务中表现出色，为特定细胞系生成高活性增强子序列，并且生成的序列具有更高的多样性。在 HepG2 细胞系中，Top-1 活性达到了 7.3，均值活性达到了 5.4。
3.3.2 酵母启动子生成
在酵母启动子生成任务中，HybriDNA 同样优于基线模型 HyenaDNA，在复杂培养基和定义培养基中的 Top-1 活性分别达到了 18.2 和 17.6，均值活动分别达到了 15.0 和 13.5，多样性达到了 30.7。
3.4 计算效率
HybriDNA 在训练阶段的吞吐量显著高于标准 Transformer 模型，尤其是在处理超过 32,000 token 的上下文长度时。例如，在 49,000 token 的上下文长度下，HybriDNA 的吞吐量约为 Transformer 的 3.4 倍。在 GPU 内存使用方面，HybriDNA 也表现出更高的效率，能够在 A100 GPU 上处理长达 65,000 token 的上下文，而标准 Transformer 模型会出现显存不足的问题。
