[Paper Review] Caduceus: Bi-Directional Equivariant Long-Range DNA Sequence Modeling

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

Genomics and Phylogenetic Studies1 references49 citations

TL;DR

The paper introduces Caduceus, a family of RC-equivariant, bi-directional long-range DNA language models built from BiMamba and MambaDNA blocks, yielding state-of-the-art performance on genomic tasks with up to 10x smaller models than competing Transformers.

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.

Motivation & Objective

• Motivate genomic sequence modeling that benefits from long-range context, bi-directionality, and reverse-complement symmetry.
• Develop efficient architectural modules that extend Mamba to BiMamba and MambaDNA for RC-equivariant DNA modeling.
• Create Caduceus foundation models and pre-/post-training strategies to maximize performance on genomics benchmarks.
• Demonstrate that RC-equivariant, bi-directional modeling yields superior performance on long-range variant effect prediction tasks.

Proposed method

• Extend the long-range Mamba block into BiMamba to enable bi-directional processing by sharing most projection weights between forward and reverse passes.
• Introduce MambaDNA to enforce reverse-complement (RC) equivariance by processing a sequence and its RC in a shared-parameter block and concatenating the outputs.
• Form Caduceus by combining BiMamba and MambaDNA with RC-equivariant embeddings and a RC-equivariant LM head, yielding RC-equivariant DNA foundation models.
• Provide two RC-equivariant modeling variants: Caduceus-PS (parameter sharing) and Caduceus-Ph (post-hoc conjoining for downstream tasks).
• Pre-train with masked language modeling on the human genome, leveraging RC symmetry to enable RC-equivariant predictions without RC data augmentation for Caduceus-PS.
• Evaluate on genomics benchmarks and long-range variant effect prediction, comparing to HyenaDNA and Transformer baselines.

Experimental results

Research questions

• RQ1 Can RC-equivariant, bi-directional architectures improve long-range genomic sequence modeling compared to unidirectional or non-equivariant models?
• RQ2 Does parameter sharing for bi-directionality (BiMamba) confer advantages in pre-training and downstream tasks at equivalent model sizes?
• RQ3 How does RC equivariance (MambaDNA) influence pre-training losses and downstream performance on variant effect and regulatory tasks?
• RQ4 Do RC-equivariant variants of Caduceus outperform significantly larger Transformer-based or non-equivariant models on long-range genomic tasks?

Key findings

• Caduceus models outperform prior SSM-based models of similar size on downstream genomics benchmarks.
• Caduceus-Ph often achieves the best overall performance among the evaluated variants on Genomics Benchmarks.
• Caduceus-PS, which enforces RC equivariance via parameter sharing and RC data-symmetric predictions, excels on long-range variant effect prediction tasks, beating larger baselines.
• On the Nucleotide Transformer Tasks, Caduceus-Ph matches or surpasses size-matched baselines and often outperforms HyenaDNA on histone and regulatory tasks, with some tasks favoring alternative approaches.
• On long-range variant effect prediction (distance to TSS), Caduceus-PS shows strong gains, even surpassing Enformer for very distant regulatory contexts.