[论文解读] Caduceus: Bi-Directional Equivariant Long-Range DNA Sequence Modeling

Yair Schiff, Chia-Hsiang Kao+4 more|PubMed|Mar 5, 2024

Genomics and Phylogenetic Studies参考文献 1被引用 49

一句话总结

该论文提出 Caduceus，一族 RC-equivariant, bi-directional long-range DNA language models 构建自 BiMamba 与 MambaDNA 模块，在基因组任务上达到 state-of-the-art 的表现，模型规模比竞争 Transformer 小多达 10x。

ABSTRACT

Large-scale sequence modeling has sparked rapid advances that now extend into biology and genomics. However, modeling genomic sequences introduces challenges such as the need to model long-range token interactions, the effects of upstream and downstream regions of the genome, and the reverse complementarity (RC) of DNA Here, we propose an architecture motivated by these challenges that builds off the long-range Mamba block, and extends it to a BiMamba component that supports bi-directionality, and to a MambaDNA block that additionally supports RC equivariance. We use MambaDNA as the basis of Caduceus, the first family of RC equivariant bi-directional long-range DNA language models, and we introduce pre-training and fine-tuning strategies that yield Caduceus DNA foundation models. Caduceus outperforms previous long-range models on downstream benchmarks; on a challenging long-range variant effect prediction task, Caduceus exceeds the performance of <mml:math xmlns:mml="http://www.w3.org/1998/Math/MathML"><mml:mrow><mml:mn>10</mml:mn> <mml:mi>x</mml:mi></mml:mrow> </mml:math> larger models that do not leverage bi-directionality or equivariance. Code to reproduce our experiments is available here.

研究动机与目标

• 为基因组序列建模提供对长程上下文、双向性和反向互补对称性有益的动机。
• 开发高效的架构模块，将 Mamba 扩展到 BiMamba 与 MambaDNA，用于 RC-equivariant DNA 建模。
• 创建 Caduceus 基础模型及预/后训练策略，以最大化在基因组基准测试上的表现。
• 证明 RC-equivariant、双向建模在长程变体效应预测任务上具有更优表现。

提出的方法

• 将长程 Mamba 块扩展为 BiMamba，以通过前向与反向传递共享大多数投影权重实现双向处理。
• 引入 MambaDNA，通过在共享参数块中对一个序列及其 RC 进行处理并连接输出来实现反向互补 (RC) 等变性。
• 通过将 BiMamba 与 MambaDNA、RC-equivariant 嵌入和 RC-equivariant LM 头相结合来形成 Caduceus，产生 RC-equivariant 的 DNA 基础模型。
• 提供两种 RC-equivariant 建模变体：Caduceus-PS（参数共享）和 Caduceus-Ph（下游任务的事后拼接）。
• 使用人类基因组的掩码语言模型进行预训练，利用 RC 对称性实现 RC-equivariant 的预测，而无需对 Caduceus-PS 进行 RC 数据增强。
• 在基因组基准测试和长程变体效应预测上进行评估，并与 HyenaDNA 和 Transformer 基线进行比较。

实验结果

研究问题

• RQ1 RC-equivariant、双向架构是否能在长程基因组序列建模方面优于单向或非等变模型？
• RQ2 在等效模型规模下，BiMamba 的参数共享双向性在预训练和下游任务中是否具有优势？
• RQ3 RC 等变性 (MambaDNA) 如何影响预训练损失和下游在变异效应和调控任务上的表现？
• RQ4 Caduceus 的 RC-equivariant 变体是否在长程基因组任务上显著超越更大 Transformer 基线或非等变模型？

主要发现

Task	CNN	HYENADNA	MAMBA	CADUCEUS W/O EQUIV.	CADUCEUS-PH	CADUCEUS-PS
MOUSE ENHANCERS	0.715 ±0.087	0.780 ±0.025	0.743 ±0.054	0.770 ±0.058	0.754 ±0.074	0.793 ±0.058
CODING VS INTERGENOMIC	0.892 ±0.008	0.904 ±0.005	0.904 ±0.004	0.908 ±0.003	0.910 ±0.003	0.915 ±0.003
HUMAN VS WORM	0.942 ±0.002	0.964 ±0.002	0.967 ±0.002	0.970 ±0.003	0.973 ±0.001
HUMAN ENHANCER COHN	0.702 ±0.021	0.729 ±0.014	0.732 ±0.029	0.741 ±0.008	0.747 ±0.004	0.745 ±0.007
HUMAN ENHANCER ENSEMBL	0.744 ±0.122	0.849 ±0.006	0.862 ±0.008	0.883 ±0.002	0.893 ±0.008	0.900 ±0.006
HUMAN REGULATORY	0.872 ±0.005	0.869 ±0.012	0.814 ±0.211	0.871 ±0.007	0.872 ±0.011	0.873 ±0.007
HUMAN OCR ENSEMBL	0.698 ±0.013	0.783 ±0.007	0.815 ±0.002	0.818 ±0.003	0.828 ±0.006	0.818 ±0.006
HUMAN NONTATA PROMOTERS	0.861 ±0.009	0.944 ±0.002	0.933 ±0.007	0.933 ±0.006	0.946 ±0.007	0.945 ±0.010

• Caduceus 模型在相似规模的基于 SSM 的模型上在下游基因组基准测试中表现更优。
• Caduceus-Ph 在所评估的变体中通常在 Genomics Benchmarks 上获得最佳整体表现。
• Caduceus-PS 通过参数共享实现 RC 等变性并进行 RC 数据对称预测，在长程变异效应预测任务中表现出色，超越更大基线模型。
• 在 Nucleotide Transformer 任务中，Caduceus-Ph 与尺寸匹配的基线相当或超越，且在组蛋白和调控任务上通常优于 HyenaDNA，但某些任务更偏向替代方法。
• 在长程变异效应预测（距离 TSS 的距离）方面，Caduceus-PS 展现出强劲提升，甚至在非常远的调控背景下超过 Enformer。