Saved in:
Bibliographic Details
Main Authors: Wang, Yingheng, Wang, Zichen, Sadeh, Gil, Zancato, Luca, Achille, Alessandro, Karypis, George, Rangwala, Huzefa
Format: Preprint
Published: 2024
Subjects:
Online Access:https://arxiv.org/abs/2411.08909
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866912304741744640
author Wang, Yingheng
Wang, Zichen
Sadeh, Gil
Zancato, Luca
Achille, Alessandro
Karypis, George
Rangwala, Huzefa
author_facet Wang, Yingheng
Wang, Zichen
Sadeh, Gil
Zancato, Luca
Achille, Alessandro
Karypis, George
Rangwala, Huzefa
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2411_08909
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Long-context Protein Language Modeling Using Bidirectional Mamba with Shared Projection Layers
Wang, Yingheng
Wang, Zichen
Sadeh, Gil
Zancato, Luca
Achille, Alessandro
Karypis, George
Rangwala, Huzefa
Biomolecules
Machine Learning
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.
title Long-context Protein Language Modeling Using Bidirectional Mamba with Shared Projection Layers
topic Biomolecules
Machine Learning
url https://arxiv.org/abs/2411.08909