SWAMamba: A Sliding Window Attention Mamba Framework for Predicting Translation Elongation Rates

Abstract

Translation elongation is essential for cellular proteostasis and is implicated in cancer and neurodegeneration. Accurately predicting the rate of ribosome elongation in each codon (also called ribosomal A site) on mRNA is important for understanding and modulating protein synthesis. However, predicting elongation rates is challenging due to the trade-off between capturing distal codon interactions and focusing on proximal codon effects at the A site. Approaches capturing distal codon interactions in the coding sequences (CDS) of mRNA fail to effectively differentiate critical regions (codons near the A site) due to insufficient effective mechanisms for focusing on these regions. Conversely, due to the limitations of models when handling long mRNA sequences, some methods simplify inputs by conditioning solely on proximal codons surrounding the A site, leading to the loss of important information from distal codons. To address this issue, we leverage Mamba's success in capturing long-range dependencies to enable the consideration of distant codons' impact on the A site. Additionally, we introduce a sliding window attention mechanism to emphasize the proximal codons around the A site during ribosome elongation. Building on these advancements, we present Sliding Window Attention Mamba (SWAMamba), a novel framework that simultaneously leverages both proximal and distal codon effects on the A site. We conduct comprehensive evaluations on ribosome data across four species and find that SWAMamba significantly outperformed current state-of-the-art methods in predicting translation elongation rates.