Long-context Protein Language Modeling Using Bidirectional Mamba with Shared Projection Layers

TL;DR

This work addresses the limitations of Transformer-based protein LMs in handling long-context and graph-context information. It introduces LC-PLM, a long-context protein language model based on bidirectional Mamba with shared projection layers (BiMamba-S), trained with MLM on UniRef data and exhibiting improved scaling and length extrapolation over ESM-2. A graph-contextual variant LC-PLM-G is then trained using random walks on PPI graphs to infuse biomolecular interaction information into token-level representations, yielding gains in remote homology, protein function, and PPI prediction, as well as enhancing structure modeling with LMFold. The results demonstrate that structured state-space models can deliver efficient long-context representations, and that coupling sequence modeling with graph context substantially improves downstream biological tasks, with promising implications for protein design and interaction-aware analysis.

Abstract

Self-supervised training of language models (LMs) has seen great success for protein sequences in learning meaningful representations and for generative drug design. Most protein LMs are based on the Transformer architecture trained on individual proteins with short context lengths. Such protein LMs cannot extrapolate to longer proteins and protein complexes well. They also fail to account for the underlying biological mechanisms carried out by biomolecular interactions and dynamics i.e., proteins often interact with other proteins, molecules, and pathways in complex biological systems. In this work, we propose LC-PLM based on an alternative protein LM architecture, BiMamba-S, built upon selective structured state-space models, to learn high-quality universal protein representations at the amino acid token level using masked language modeling. We also introduce its graph-contextual variant, LC-PLM, which contextualizes protein-protein interaction (PPI) graphs for a second stage of training. LC-PLM demonstrates favorable neural scaling laws, better length extrapolation capability, and up to 30% and 16% improvements on protein downstream tasks compared to Transformer-based ESM-2 when trained with 100B and 1T tokens, respectively. LC-PLM-G further trained within the context of PPI graphs shows promising results on protein structure and function prediction tasks. Our study demonstrates the benefit of increasing the context size with computationally efficient LM architecture (e.g., structured state space models) in learning universal protein representations and incorporating molecular interaction contexts contained in biological graphs.

Figures (18)

• Figure 1: Our model enables long-context capability, length extrapolation ability, better neural scaling law, and interaction context-aware inference.
• Figure 2: BiMamba-S block. The forward and reverse modules share the linear projection layers. The normalized input will be reversed along the sequence dimension before being fed in. The output of the reversed will be flipped back and then added to the forward's output.
• Figure 3: The illustration of graph-contextual protein language modeling (LC-PLM-G). The positive paths are sampled with random walks on the graph and the same number of negative paths are sampled from disconnected node pairs randomly. The sampled paths will be transformed into multi-protein sequences composed of AA tokens and special graph identifier tokens. The BiMamba-S encoder is trained using MLM, the same as in the first-stage training.
• Figure 4: Evaluation loss across different model sizes for LC-PLM and ESM-2, showing that LC-PLM has a better scaling behavior when increasing the number of model parameters.
• Figure 5: Evaluation loss comparison between BiMamba and BiMamba-S over training steps.