Caduceus：双向等效长范围DNA序列建模.

0. TL; DR

本文提出了一个名为Caduceus [kə'dju:siəs]的DNA序列建模架构，旨在解决基因组序列建模中的挑战，包括长程依赖、基因组上下游区域的影响以及DNA的反向互补性（RC）。Caduceus基于Mamba模块，通过扩展为BiMamba（支持双向性）和MambaDNA（支持RC等变性），实现了首个双向、RC等变的长程DNA语言模型。实验表明，Caduceus在下游任务中优于现有模型，尤其是在长程变异效应预测任务中，性能超过10倍大的非双向或非等变模型。

1. 背景介绍

大规模序列建模在自然语言处理（NLP）、蛋白质组学和基因组学等领域取得了显著进展。然而，基因组序列建模面临独特挑战：细胞表型可能受到基因组上下游区域的影响，需要模型处理双向上下文；DNA由两条互补链组成，反向互补链携带相同信息，模型需要考虑这一特性；许多基因组学任务涉及长程相互作用，如距离基因100万碱基对的核苷酸可能对基因表达有显著调控作用。

现有的DNA序列建模方法在处理双向性和RC等变性方面存在不足。例如，传统的单向模型无法充分利用基因组的双向信息，而忽略RC等变性的模型可能在反向互补序列上表现不佳。因此，本文提出一种新的架构，通过引入双向和RC等变性，提升DNA序列建模的性能。

2. Caduceus 模型

本文基于Mamba模块，提出BiMamba和MambaDNA模块，分别支持双向性和RC等变性。Caduceus模型结合了这两个模块，能够处理长程依赖关系，并且在预训练和微调策略下表现出色。

BiMamba是Mamba模块的双向扩展。它通过将输入序列$X$及其反转版本$\text{reverse}(X)$分别通过Mamba模块处理，然后将反转序列的输出再次反转并加到前向输出上，实现双向建模。为了避免参数翻倍，BiMamba共享了输入和输出投影参数。

\[\text{BiMamba}(X) = \text{Mamba}(X) + \text{reverse}(\text{Mamba}(\text{reverse}(X)))\]

MambaDNA模块在BiMamba的基础上引入了RC等变性。它将输入序列沿通道维度拆分为两部分$X_{1:D/2},X_{D/2:D}$，一部分应用RC操作（$RC(x_{1:T}^{1:D})=x_{T:1}^{D:1}$），然后共享参数处理前向和RC序列。最后，将处理后的反向序列再次应用RC操作，并与前向输出沿通道维度拼接。

\[\text{MambaDNA}(X) = \text{concat}(\text{BiMamba}(X_{1:D/2}), \text{RC}(\text{BiMamba}(\text{RC}(X_{D/2:D}))))\]

Caduceus模型基于MambaDNA模块构建，包括RC等变的词嵌入模块$\text{Emb}{\text{RCe}}$和语言模型头$\text{LM}{\text{RCe}}$。Caduceus-PS版本通过参数共享实现RC等变性，而Caduceus-Ph版本则在下游任务中采用后处理拼接技术。

\[\text{Caduceus}(X) = \text{LM}_{\text{RCe}}(\text{MambaDNA}(\text{Emb}_{\text{RCe}}(X)))\]

Caduceus模型的优势包括：

3. 实验分析

本文使用人类参考基因组（HG38）进行预训练，下游任务包括基因组学基准测试、NT任务和变异效应预测任务。根据不同任务类型，使用准确率、F1分数、马修斯相关系数（MCC）和AUCROC等指标进行评估。

Caduceus模型在多个基因组学基准任务中表现出色，尤其是在需要长程建模的任务中。例如，在人类增强子分类任务中，Caduceus-Ph模型的准确率达到0.893，超过了其他模型。

在18个NT任务中，Caduceus模型在多个任务上超过了现有的基于Transformer的模型，尤其是在长程建模任务中。例如，在H3K9ac组蛋白标记预测任务中，Caduceus-Ph模型的MCC达到0.622，超过了其他模型。

Caduceus模型在预测变异对基因表达的影响任务中表现出色，尤其是在长程变异效应预测方面。例如，在距离转录起始位点（TSS）100kb以上的变异效应预测任务中，Caduceus-PS模型的AUCROC达到0.586，超过了其他模型。