Caduceus: Bi-Directional Equivariant Long-Range DNA Sequence Modeling

Yair Schiff; Chia-Hsiang Kao; Aaron Gokaslan; Tri Dao; Albert Gu; Volodymyr Kuleshov

Source-linked AI summary

Caduceus: Bi-Directional Equivariant Long-Range DNA Sequence Modeling

Yair Schiff, Chia-Hsiang Kao, Aaron Gokaslan, Tri Dao, Albert Gu, Volodymyr Kuleshov

arXiv:2403.03234v2Mar 5, 2024q-bio.GNcs.LG

TL;DR

Genomic modeling requires long-range, bidirectional context while accounting for DNA’s reverse-complement symmetry. The paper introduces BiMamba and MambaDNA, then builds Caduceus foundation models with them. Caduceus outperforms prior long-range models, including models up to 10x larger on long-range tasks.

Problem
Genomic tasks require upstream and downstream context, reverse-complement modeling, and interactions spanning up to 1 million base pairs.
Method
The paper extends long-range Mamba with parameter-efficient bidirectionality in BiMamba and reverse-complement equivariance in MambaDNA, using MambaDNA to build Caduceus.
Results
Caduceus outperforms previous long-range models and models up to 10x larger that do not use bi-directionality or equivariance, including on long-range variant effect prediction.
Takeaways & Limitations
Caduceus provides bidirectional, long-range, reverse-complement-equivariant DNA foundation models for biologically relevant genomic prediction tasks.
Takeaways & Limitations
The authors note that, like machine-learning and language models generally, the work may provide societal benefits but is subject to misuse.

Abstract

from arXiv · show

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 10x larger models that do not leverage bi-directionality or equivariance.

1. Introduction

Genomic modeling must capture bidirectional and long-range regulatory context while respecting DNA’s reverse-complement symmetry. Caduceus addresses these challenges with BiMamba and MambaDNA components and achieves strong performance across long-range genomic tasks.

Motivation: Genomic sequences require bidirectional context, reverse-complement modeling, and long-range interactions that can extend up to 1 million base pairs.These challenges arise because regulatory effects may depend on upstream and downstream regions and distant genomic elements.
Architectural components: MambaDNA further adds reverse-complement equivariance and can replace sequence modules in supervised or self-supervised genome-analysis architectures.It builds on the long-range Mamba block, which handles sequences of over hundreds of thousands of nucleotides without attention’s quadratic cost.
Caduceus models: Caduceus uses MambaDNA to create the first family of RC-equivariant bidirectional long-range DNA language models.The models are developed as DNA foundation models using pre-training and fine-tuning strategies for genomic prediction tasks.
Architectural components: BiMamba extends Mamba for bidirectional sequence modeling, while sharing projection parameters to improve parameter and hardware efficiency.The extension processes forward and reverse sequences and combines their outputs.
Results: Caduceus outperforms models up to 10x larger that lack bi-directionality or equivariance on long-range tasks.It also outperforms existing attention- and SSM-based models on a long-range variant-effect prediction task.

2. Background

DNA sequence modeling must capture long-range context and account for reverse-complementary strands, while SSMs provide efficient mechanisms for long sequences. Mamba adds input-dependent selection to improve expressivity beyond linear-time-invariant formulations.

DNA Terminology: DNA tasks require bi-directional context because phenotypes can depend on upstream and downstream genomic base pairs.
DNA Terminology: Reverse-complement equivariance is important because both DNA strands carry equivalent information and either strand may be sequenced.
Structured State Space Models: Structured State Space Models map input sequences to outputs through an intermediate state representation and are effective for long-range modeling.
Structured State Space Models: Zero-order-hold discretization converts continuous SSM parameters using A, B, and a time-scale parameter ∆.
Structured State Space Models: Linear-time-invariant SSMs can be unrolled as convolutions for efficient parallel training, but this efficiency limits adaptation to specific inputs.
Structured State Space Models: Mamba introduces a selective SSM whose B, C, and ∆ parameters depend on the input sequence.

3. Bi-Directional & RC-Equivariant Mamba

The paper extends Mamba from causal processing to bi-directional modeling and then adds reverse-complement equivariance through shared parameters. The resulting MambaDNA block is claimed to satisfy the desired equivariance property without substantial extra memory.

Bi-Directional Mamba: BiMamba processes the original and reversed sequences, flips the reversed output back, and adds it to the forward output.
Bi-Directional Mamba: Shared projection weights prevent the bi-directional design from doubling the module’s parameter count.
RC-Equivariant Mamba: MambaDNA applies a Mamba or BiMamba block to a sequence and its reverse complement using shared parameters.
RC-Equivariant Mamba: The RC-equivariant module splits channels, applies the shared sequence operator to the relevant representations, and recombines them along the channel dimension.
RC-Equivariant Mamba: Theorem 3.1 states that the MRCe,θ operator satisfies the desired reverse-complement equivariance property.
RC-Equivariant Mamba: MambaDNA does not entail significant additional memory because its wrapped sequence operator is completely shared across the forward and reverse-complement applications.

4. Caduceus

Caduceus is a bi-directional DNA language-model family that enforces reverse-complement equivariance through parameter sharing or downstream post-hoc conjoining. Its architecture combines equivariant embeddings, MambaDNA or BiMamba blocks, and an equivariant language-model head.

Caduceus: Caduceus has two variants: Caduceus-PS uses parameter sharing, while Caduceus-Ph uses post-hoc conjoining to maintain reverse-complement equivariance.
Caduceus-PS: Caduceus-PS combines an RC-equivariant token embedding module, BiMamba wrapped in MambaDNA blocks, and an RC-equivariant language-model head.
Caduceus-PS: Caduceus-PS predictions for a reverse-complement sequence correspond to reversed and complemented predictions for the original sequence.
Caduceus-PS: Caduceus-PS uses masked language modeling without reverse-complement augmentation because its predictions are inherently symmetric under the operation.
Caduceus-PS: For downstream tasks, Caduceus-PS splits final hidden states and averages the two splits to enforce reverse-complement invariance.
Caduceus-Ph: Caduceus-Ph uses BiMamba blocks, reverse-complement augmentation during pretraining and training, and post-hoc conjoining during inference.

5. Experiments

Experiments evaluate Caduceus across pre-training, regulatory benchmarks, and long-range variant effect prediction, showing gains from Mamba, parameter sharing, RC equivariance, and long-range modeling.

Pre-training: Mamba attains lower cross entropy loss than HyenaDNA at comparable model sizes and sequence lengths during human-genome pre-training.
Pre-training: Deeper BiMamba models with projection parameter sharing achieve better MLM pre-training loss than shallower untied bidirectional models at similar parameter counts.
Pre-training: RC equivariant language modeling produces better MLM pre-training loss, which is biologically relevant to downstream variant effect prediction.
Genomic Benchmarks: Caduceus models attain the best performance across all Genomic Benchmarks annotations, with Caduceus-Ph performing best overall.
Nucleotide Transformer Tasks: Caduceus-Ph beats attention-based methods with orders of magnitude more parameters on 8 of 18 Nucleotide Transformer prediction tasks.
Variant Effect Prediction: Caduceus models consistently outperform HyenaDNA, while Caduceus-PS surpasses Nucleotide Transformer v2 and Enformer for variants more than 100k bps from the nearest TSS.

6. Related Work

Related work spans long-context DNA language models, reverse-complement training, and bidirectional sequence modeling. Caduceus builds on these directions while addressing limitations in context directionality and RC handling.

DNA Language Models: Transformer-based DNA language models are constrained by quadratic scaling, with reported maximum context sizes of roughly 12,000 bps.
DNA Language Models: Sparse attention and dilated convolutions extend receptive fields, but the cited models use either substantially larger contexts or only 512-bp training contexts.
DNA Language Models: HyenaDNA scales to sequences up to 1 million bps but is unidirectional and not inherently robust to reverse-complement inputs.
Reverse Complement Training for DNA: Prior reverse-complement approaches include data augmentation, RC parameter sharing, post-hoc conjoining, and contrastive self-supervised pre-training.
Bidirectional Sequence Modeling: Bidirectional pre-training originated with simultaneous forward and backward recurrent models and later appeared in Transformer and state-space-model approaches.

7. Conclusion

The paper introduces Mamba-based architectural innovations for bidirectional and RC-equivariant sequence modeling and presents Caduceus as a DNA foundation model with broad downstream performance.

Architectural innovations to Mamba enable bidirectional and reverse-complement-equivariant sequence modeling.
Caduceus outperforms comparably sized unidirectional Hyena-based models and much larger Transformer-based models across biologically relevant tasks.
Variant effect prediction for gene expression is identified as a particularly notable application of Caduceus.

Impact Statement

The impact statement frames the work as advancing machine learning while acknowledging potential societal benefits and misuse.

The paper’s stated goal is to advance the field of machine learning.
The authors identify potential societal benefits and possible misuse as consequences of the work.

B. Proof of Theorem 4.1

The proof establishes RC equivariance compositionally: it verifies the embedding and model operators, then uses closure under composition. The argument also shows that reversing sequence positions reverses the outputs of the underlying language model.

RC equivariance is preserved under composition of two RC-equivariant sequence operators.This closure property supports proving equivariance for the composed Caduceus operators.
The proof reduces Caduceus-PS equivariance to verifying each operator in the composed model, including the embedding and recurrent components.The embedding operator is shown first, while the recurrent operator follows from Theorem 3.1 and induction.
Reversing the input sequence reverses the language-model output because the model applies the same weight transformation independently across sequence elements.The proof then focuses on equivariance with respect to channel flipping.

C. Pre-training

Pre-training uses human-reference-genome data, long sequence segments, and Mamba-specific optimization settings. Bidirectional models use masked language modeling, while most models additionally use reverse-complement augmentation.

34,021 training segments are extended to length 1,048,576, covering the genome and totaling around 35 billion nucleotide tokens.The data are based on splits from the previous Enformer study.
All Mamba-based models use a learning rate of 8e−3 and maintain 2^20 tokens per batch across sequence lengths.For example, batch size is 1,024 at sequence length 1,024 and 8 at sequence length 131k.
Models are trained with cosine decay and ADAM, using β1=0.95 and β2=0.9.
Bidirectional models mask 15% of tokens using the BERT recipe: 80% [MASK], 10% random tokens, and 10% unchanged.
HyenaDNA is re-pretrained at sequence lengths 1,024, 32k, and 131k using its original hidden dimensions and depths.The HyenaDNA learning rate is set to 6e−4, with other pre-training details matched to the Mamba-based models.

D.1. Genomics Benchmark

The Genomics Benchmark evaluation uses cross-validation, early stopping, and matched fine-tuning procedures across CNN, HyenaDNA, Mamba, and Caduceus models. The section also specifies the variant-effect data construction and SVM-based evaluation protocol.

Five-fold cross-validation uses 90/10 train/validation splits, with early stopping and 10 fine-tuning epochs.This protocol is used to maintain true train and test splits.
The CNN baseline uses an embedding layer, convolutional layers with 16, 8, and 4 channels, batch normalization, max-pooling, and two fully connected layers.
Caduceus and Mamba models use batch size 256, four layers, and approximately matched parameter counts through hidden dimensions 118 and 128, respectively.
Variant labels mark SNPs with causal probability above .9, while chromosomes 9 and 10 form the held-out test set.
Model inputs are 131k base pairs for Caduceus and HyenaDNA, compared with 12k for Nucleotide Transformer and 196k for Enformer.
SVM classifiers with RBF kernels are trained separately by distance-to-nearest-TSS strata, and mean AUROC with one standard deviation is reported across five repetitions.Each distance bucket uses 5,000 randomly selected training points.