Prot2Token: A Unified Framework for Protein Modeling via Next-Token Prediction

Mahdi Pourmirzaei, Farzaneh Esmaili, Salhuldin Alqarghuli, Mohammadreza Pourmirzaeioliaei, Ye Han, Kai Chen, Mohsen Rezaei, Duolin Wang, Dong Xu

Keywords: protein language model, large language model, protein prediction, bioinformatics, next token prediction, generative ai, autoregressive language model, chemical language model, multi-task learning

TL;DR: We introduce a task-agnostic tokenization approach for nest token predictio to solve various types of protein prediction tasks.

Abstract: The diverse nature of protein prediction tasks has traditionally necessitated specialized models, hindering the development of broadly applicable and computationally efficient Protein Language Models (PLMs). In this work, we introduce Prot2Token, a unified framework that overcomes these challenges by converting a wide spectrum of protein-related predictions—from sequence-level properties and residue-specific attributes to complex inter-protein interactions—into a standardized next-token prediction format. At its core, Prot2Token employs an autoregressive decoder, conditioned on embeddings from pre-trained protein encoders and guided by learnable \texttt{task tokens}, to perform diverse predictions. This architecture uniquely facilitates multi-task learning, enabling general-purpose decoders to generalize across five distinct categories. We present extensive experimental validation across a variety of benchmarks, demonstrating Prot2Token's predictive power in different types of protein-prediction tasks. In 3D structure prediction, Prot2Token delivers substantial speedups (up to $\sim$1000$\times$ faster than AlphaFold2 with MSA on the same hardware) while, across other numerous tasks, matching or surpassing specialized methods. Beyond that, we introduce an auxiliary self-supervised decoder pre-training approach to improve spatially sensitive task performance. Prot2Token thus offers a step towards standardizing biological prediction into a generative interface, promising to accelerate biological discovery and the development of novel therapeutics.

Primary Area: applications to physical sciences (physics, chemistry, biology, etc.)

Submission Number: 9875