[Paper Review] Bio-xLSTM: Generative modeling, representation and in-context learning of biological and chemical sequences

Niklas Schmidinger, Lisa Schneckenreiter+8 more|arXiv (Cornell University)|Nov 6, 2024

Machine Learning in Bioinformatics5 citations

TL;DR

Bio-xLSTM tailors the xLSTM architecture to DNA, proteins, and SMILES, enabling long-context generative modeling, rich representations, and in-context learning with linear-runtime sequence modeling.

ABSTRACT

Language models for biological and chemical sequences enable crucial applications such as drug discovery, protein engineering, and precision medicine. Currently, these language models are predominantly based on Transformer architectures. While Transformers have yielded impressive results, their quadratic runtime dependency on the sequence length complicates their use for long genomic sequences and in-context learning on proteins and chemical sequences. Recently, the recurrent xLSTM architecture has been shown to perform favorably compared to Transformers and modern state-space model (SSM) architectures in the natural language domain. Similar to SSMs, xLSTMs have a linear runtime dependency on the sequence length and allow for constant-memory decoding at inference time, which makes them prime candidates for modeling long-range dependencies in biological and chemical sequences. In this work, we tailor xLSTM towards these domains and propose a suite of architectural variants called Bio-xLSTM. Extensive experiments in three large domains, genomics, proteins, and chemistry, were performed to assess xLSTM's ability to model biological and chemical sequences. The results show that models based on Bio-xLSTM a) can serve as proficient generative models for DNA, protein, and chemical sequences, b) learn rich representations for those modalities, and c) can perform in-context learning for proteins and small molecules.

Motivation & Objective

• Motivate and develop long-context language models for biological and chemical sequences beyond Transformer-based architectures.
• Adapt xLSTM into domain-specific variants (DNA-xLSTM, Prot-xLSTM, Chem-xLSTM) and enable generation, inpainting, and ICL.
• Evaluate Bio-xLSTM on genomics, protein, and chemical sequence tasks, comparing to state-of-the-art baselines.
• Demonstrate RC-equivariance for DNA and assess downstream performance on classification and design tasks.
• Showcase in-context learning capabilities and domain-conditioned generation without fine-tuning.

Proposed method

• Extend xLSTM with sLSTM and mLSTM blocks tuned for biological and chemical sequences.
• Develop three domain-specific variants: DNA-xLSTM, Prot-xLSTM, Chem-xLSTM with context windows and RoPE for long-range dependencies.
• Implement modeling modes: causal language modeling (CLM), masked language modeling (MLM), fill-in-the-middle (FIM), and in-context learning (ICL).
• Incorporate reverse-complement (RC) equivariance via post-hoc conjoining (PH) or parameter sharing (PS).
• Train DNA-xLSTM on the human genome with RC-equivariant variants and compare to HyenaDNA, Mamba, DNA-Mamba, and Transformers.
• Train Prot-xLSTM with homology-aware, alignment-free inputs using FIM on unaligned homologs and evaluate for generation and variant fitness prediction.
• Train Chem-xLSTM for unconditional SMILES generation and domain-conditioned ICL, evaluating on SMILES realism.
• Assess long-context capabilities with RoPE and context sizes up to 262k tokens for proteins and up to 32k for DNA.

Figure 1: Overview of Bio-xLSTM. Top left: xLSTM for natural language processing tasks. Top right: Considered modeling approaches for biological sequences: masked language modeling, equivariance to reverse complementary sequence, and in-context learning. Bottom left: DNA-xLSTM models are trained on

Experimental results

Research questions

• RQ1 Can Bio-xLSTM variants effectively model long biological and chemical sequences with linear memory scaling?
• RQ2 Do DNA-xLSTM, Prot-xLSTM, and Chem-xLSTM provide competitive or superior performance to domain-specific Transformers and SSM-based models on their respective tasks?
• RQ3 Are RC-equivariant designs (PH/PS) beneficial for DNA modeling and downstream tasks?
• RQ4 Can Prot-xLSTM perform homology-conditioned in-context learning for generative design and residue-level predictions?
• RQ5 Can Chem-xLSTM enable domain-conditioned in-context learning for molecule generation without fine-tuning?

Key findings

• DNA-xLSTM outperforms Transformers, Mamba, and HyenaDNA on CLM and MLM pre-training with 2M parameters on the human genome.
• DNA-xLSTM-2M (PH/PS) matches or exceeds baselines on downstream genomic classification tasks across 12 of 18 tasks at under 2M parameters.
• Prot-xLSTM-102M achieves superior perplexity and generation quality in homology-conditioned protein generation, outperforming ProtMamba and Transformer++ baselines especially at longer contexts.
• Prot-xLSTM-102M outperforms ProtMamba-107M despite using fewer total training tokens, indicating efficient long-context learning.
• Chem-xLSTM achieves lowest Fréchet ChemNet Distance (FCD) and competitive perplexity in unconditional SMILES generation, indicating realistic chemical output.

Figure 2: Pre-training of 2M-parameter DNA models on the human reference genome (GRCh38). Models are trained at single-nucleotide resolution with a context length of 1024 bases. Left: causal language modeling . Learning curves display NTP loss ( $\downarrow$ ) on a test set, plotted against the numb