Selective State Space Models Outperform Transformers at Predicting RNA-Seq Read Coverage

Area of Science:

• Computational Biology and Genomics
• Machine Learning in Bioinformatics

Background:

• Transformer models are foundational for gene expression prediction from DNA, but their training is time-consuming and costly.
• Alternative models inspired by signal processing, including state-space models (Mamba), Fourier transforms (Hyena), and wavelet transforms (MultiResNet), have emerged to address these limitations.

Purpose of the Study:

• To evaluate alternative machine learning architectures as replacements or complements to attention mechanisms for predicting gene expression.
• To compare the performance of convolutional, attention, Hyena, Mamba, and striped-architecture models in functional genomics tasks.

Main Methods:

• Developed the 'bilby' software library in Python and Jax/Flax, offering various convolutional, attention, and state-space models.
• Conducted supervised multi-task learning experiments on gene expression data.
• Compared model performance using prediction accuracy metrics (Pearson R, r^2) and downstream SNP classification on a GTEx eQTL dataset.

Main Results:

• Convolutional models incorporating bidirectional Mamba layers showed small but consistent prediction accuracy improvements (3-4% Pearson R, 1-2% r^2) over convolution-attention models.
• The highest gains were observed with a striped architecture combining Mamba and attention layers.
• Hyena models were not competitive, and MultiResNet was too slow; Mamba's accuracy gains did not significantly improve SNP classification performance.

Conclusions:

• Selective state-space models, particularly Mamba and striped Mamba architectures, demonstrate potential for improving gene expression prediction in functional genomics.
• Further research into Mamba-based models is warranted for these tasks, despite current limitations in downstream SNP classification.
• The 'bilby' library and trained models are publicly available for reproducible research.