PTM-Mamba: A PTM-Aware Protein Language Model with Bidirectional Gated Mamba Blocks

Introduction: Proteins serve as the workhorses of living organisms, orchestrating a wide array of vital functions. Post-translational modifications (PTMs) of their amino acids greatly influence the structural and functional diversity of different protein types and uphold proteostasis, allowing cells to swiftly respond to environmental changes and intricately regulate complex biological processes. To this point, efforts to model the complex features of proteins have involved the training of large and expressive protein language models (pLMs) such as ESM-2 and ProtT5, which accurately encode structural, functional, and physicochemical properties of input protein sequences. However, the over 200 million sequences that these pLMs were trained on merely scratch the surface of proteomic diversity, as they neither input nor account for the effects of PTMs. In this work, we fill this major gap in protein sequence modeling by introducing PTM tokens into the pLM training regime. We then leverage recent advancements in structured state space models (SSMs), specifically Mamba, which utilizes efficient hardware-aware primitives to overcome the quadratic time complexities of Transformers. After adding a comprehensive set of PTM tokens to the model vocabulary, we train bidirectional Mamba blocks whose outputs are fused with state-of-the-art ESM-2 embeddings via a novel gating mechanism. We demonstrate that our resultant PTM-aware pLM, PTM-Mamba, improves upon ESM-2's performance on various PTM-specific tasks. PTM-Mamba is the first and only pLM that can uniquely input and represent both wild-type and PTM sequences, motivating downstream modeling and design applications specific to post-translationally modified proteins.

Materials and

Methods: We employ recent innovations in state space modeling (SSMs), specifically the Mamba model, which utilizes hardware-aware primitives to overcome the quadratic time complexity of the standard attention mechanism (Gu and Dao, 2023). For training, rather than relearning the semantics and evolutionary statistics of protein sequences, we curate a comprehensive, experimentally-validated dataset of post-translationally modified protein sequences, and leverage a fusion module that dynamically combines the embeddings of PTM-Mamba and ESM-2 (Lin et al. 2023) in a gated mechanism. Our resultant PTM-Mamba model is the first and only pLM that can input both wild-type and modified protein sequences to derive unique and accurate latent embeddings, enabling downstream PTM-aware protein modeling.PTM-Mamba is a bidirectional pLM that represents both amino acids and PTM tokens. The model builds on top of the state-of-the-art, structured SSM, Mamba (Gu and Dao, 2023), with an added bidirectional module and gated embedding fusion module. To preserve comprehension of regular amino acids, we train PTM-Mamba as a head to the state-of-the-art ESM-2-650M model (Lin et al. 2023), where wild-type amino acid tokens are passed into ESM-2-650M to retrieve its output embeddings and PTM tokens are converted into tokens for ESM-2-650M input. Sequences are finally fed into the embedding layer of the PTM-Mamba, which naturally processes both wild-type and PTM tokens. To join the ESM-2-650M and PTM-Mamba embeddings, we propose a novel gating mechanism, where the two embeddings are concatenated and filtered via a sigmoid-activated linear gate to produce a final output representation.

Results, Conclusions, and Discussions: We sought to evaluate PTM-Mamba’s ability to predict non-histone acetylation sites as compared to ESM-2-650M, by training a per-residue binary classifier on the recently-curated Non-Histone Acetylation Collection (NHAC) dataset (Meng et al., 2023). Our results demonstrate that PTM-Mamba outperforms ESM-2-650M on all classification metrics over the entire test dataset (Figure 1B), highlighting the added benefit of PTM representations to ESM-2-650M embeddings and motivating usage of PTM-Mamba for additional PTM site prediction tasks. We next turned our attention to phosphorylation. As a clear demonstration of its expressive PTM-aware latent representations, for all classification metrics, PTM-Mamba strongly outperforms all other models, including ESM-2-650M (Figure 1C).
Finally, we show that modified residue tokens belonging to specific PTM classes, such as phosphorylation and acetylation, are spatially proximal in two dimensions via t-SNE (Figure 2A). Furthermore, based on token embedding distribution, it is clear that PTM-Mamba strongly attends to PTM residue tokens, compared to less spatially diverse wild-type tokens, indicating that the model focuses effort on augmenting ESM-2 embeddings with PTM residue information (Figure 2B). Finally, to assess the functional plausibility of PTM-Mamba embeddings in a zero-shot fashion, we analyzed the model’s output logits for specified amino acids in randomly downselected wild-type sequences. We show that the largest logits correspond to PTMs that are capable of being installed on the wild-type amino acid (Figure 2C).

In this work, we introduce the first pLM trained to represent both wild-type and post-translationally modified protein sequences, capturing PTM-specific effects in its latent space. To our knowledge, PTM-Mamba is the only model that enables input of protein sequences with PTM-altered residues specifically encoded at inference time. We envision that such a PTM-aware pLM will enable the design of binders that recognize specific post-translational protein states. As such, next steps include augmenting our PepMLM architecture, which leverages a novel peptide masking strategy on ESM-2 embeddings (Chen et al with PTM-Mamba embeddings, enabling binder design to specific modified proteoforms. We envision that our platform, now incorporating PTM information, will impart a level of targeting specificity that is currently intractable within standard drug discovery pipelines.

Acknowledgements (Optional):