Caduceus: Bi-Directional Equivariant Long-Range DNA Sequence Modeling (2403.03234v2)

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 LLMs, 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.

• The paper introduces the Caduceus model, which extends the MambaDNA block to capture bi-directional context, reverse complement equivariance, and long-range dependencies in genomic sequences.
• It presents two variants—Caduceus-PS with parameter sharing for RC equivariance and Caduceus-Ph using post-hoc conjoining for downstream tasks.
• Empirical results demonstrate that Caduceus outperforms state-of-the-art models on genomic benchmarks, notably enhancing variant effect prediction.

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

In the rapidly evolving field of genomic sequence modeling, addressing the challenges of modeling DNA sequences introduces unique hurdles compared to traditional sequence modeling tasks. These challenges include the necessity for bi-directional context modeling, reverse complement (RC) equivariance, and the handling of long-range dependencies inherent in genomic data. Addressing these issues, Schiff et al. propose a novel architecture in their recent work, entitled "Caduceus: Bi-Directional Equivariant Long-Range DNA Sequence Modeling."

Introduction to the Challenges in DNA Sequence Modeling

Modeling genomic sequences differs significantly from natural language processing or even protein sequence modeling due to several distinct characteristics of DNA:

• Bi-directionality: Cellular phenotypes are influenced by base pairs both upstream and downstream in the genome, necessitating models that can leverage bi-directional context.
• Reverse Complementarity: DNA's double-stranded structure means that either strand can functionally represent the same genetic information, albeit in reverse and with complementary base pairs. This property is critical for accurate DNA sequence modeling.
• Long-Range Dependencies: Many genomic functions are regulated by elements that may be located far from the genes they control, necessitating models capable of capturing long-range dependencies.

The Caduceus Architecture

Addressing these specificities, Schiff et al. introduce the MambaDNA block, an extension of the long-range Mamba block. This novel block supports bi-directionality and incorporates RC equivariance, making it suited for genomic sequence modeling. The paper introduces two versions of their proposed model, Caduceus, built upon the MambaDNA block:

• Caduceus-PS, incorporating parameter sharing to enforce RC equivariance through the architecture, allowing for RC equivariant LLM pre-training.
• Caduceus-Ph, leveraging post-hoc conjoining, a method ensuring RC invariance in downstream tasks, effectively acting as a potent alternative to inherent RC equivariance.

Empirical Evaluation and Findings

Upon evaluation, Caduceus demonstrates compelling performance advantages over existing models:

• Performance on Downstream Benchmarks: Caduceus outperforms previous state-of-the-art models on a range of genomic benchmarks, particularly on tasks necessitating long-range modeling.
• Variant Effect Prediction (VEP): In tasks predicting the phenotypic effect of genetic mutations, Caduceus, especially the PS variant, exhibited superior performance. Its ability to model long-range dependencies appears to offer significant advantages in recognising the regulatory impacts of distant genetic variations.

Implications and Future Directions

The introduction of Caduceus marks a significant advancement in the field of genomic sequence modeling, addressing key challenges unique to DNA sequences with innovative architectural modifications. The model’s performance highlights the importance of bi-directionality and RC equivariance in capturing the complex regulatory mechanisms encoded within the genome.

Future research directions could explore the extension of Caduceus’s architecture to other biological sequences, such as RNA, or investigate its applicability in more specific genomics tasks, such as chromatin accessibility prediction. Furthermore, the model’s adaptability to other sequence modeling domains outside genomics presents an exciting avenue for broader applications.

Conclusion

Schiff et al.'s Caduceus model introduces a powerful tool for genomic sequence modeling, adept at handling bi-directionality, RC equivariance, and long-range dependencies. Its superior performance on challenging genomic tasks underscores the potential of carefully designed model architectures to advance our understanding of complex biological systems.