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

TL;DR

Prot2Token is a versatile, unified framework that converts various protein prediction tasks into a next-token prediction format, enabling multi-task learning, faster predictions, and broad applicability across protein modeling tasks.

Contribution

It introduces Prot2Token, a novel unified autoregressive framework that standardizes diverse protein prediction tasks into a single generative model, improving efficiency and generalization.

Findings

Achieves up to 1000x speedup in 3D structure prediction compared to AlphaFold2.

Matches or surpasses specialized methods across multiple protein prediction tasks.

Demonstrates effective multi-task learning and improved performance with self-supervised pre-training.

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 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 1000x 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. The code is available at https://github.com/mahdip72/prot2token .