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

Transformers are the basis for many state-of-the-art machine learning tools, including those for predicting gene expression data from DNA sequence. The considerable time and cost of training transformer models has motivated development of alternative approaches inspired by ideas from the signal-processing literature, such as state-space models (Mamba), Fourier transforms (Hyena), and wavelet transforms (MultiResNet). To evaluate these methods as potential replacements (or complements) for attention, we developed a software library bilby, implemented using Python and Jax/Flax, providing convolutional, attention, bidirectional Hyena, bidirectional Mamba, and striped-architecture models for supervised multi-task learning in functional genomics. We report a comparison of these architectures, testing several hyperparameters and variations, and reporting performance statistics for the withheld test set as well as downstream SNP classifiers. Relative to models comprising convolution and attention layers (implemented in Python and TensorFlow via the Baskerville library used by the Borzoi software), models comprising convolutional, bidirectional Mamba, and (optionally) attention layers achieve small but consistent improvements in prediction accuracy, for roughly comparable training times and parameter counts, when averaged across all output tracks and data splits (a proportional increase of 3-4% in Pearson R, and 1-2% in r