Source-linked AI summary
Prot2Token: A Unified Framework for Protein Modeling via Next-Token Prediction
Mahdi Pourmirzaei, Farzaneh Esmaili, Salhuldin Alqarghuli, Mohammadreza Pourmirzaei, Ye Han, Kai Chen, Mohsen Rezaei, Duolin Wang, Dong Xu
TL;DR
Protein-prediction models often require task-specific architectures, and autoregressive protein models lack robust control over interpretable predictions. Prot2Token uses a shared next-token framework to map diverse protein outputs, with initial analyses reporting synergistic gains from joint training.
Problem
Protein-prediction approaches often rely on task-specific models, while autoregressive protein models lack robust ways to guide outputs into human-interpretable formats.
Source: §1 · p. 2Method
Prot2Token uses a universal tokenization protocol and a task-token-guided autoregressive decoder conditioned on pretrained protein representations to predict diverse outputs.
Source: §1 · p. 2, §2.1 · p. 4Results
Takeaways & Limitations
Prot2Token offers a general-purpose sequence-generation alternative to specialized task heads for diverse protein outputs.
Source: §4 · p. 9Performance remains bounded by foundational-model biases and tokenizer fidelity; the 3D structure TM-score plateau of approximately 0.55 reflects the VQ-VAE reconstruction ceiling.
Source: §4.2 · p. 10
Abstract
from arXiv · showhide
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 .
1 Introduction
Protein prediction remains difficult to unify because models often require task-specific predictors and autoregressive protein models lack robust control over interpretable outputs. Prot2Token maps diverse tasks to a shared next-token format and demonstrates joint-task training and decoder pre-training for binding-site tasks.
1 Introduction
- Existing PLMs commonly require separately trained task-specific predictors, while autoregressive protein models lack robust ways to guide outputs into interpretable formats.These approaches are described as inefficient to scale and as limiting practical use of autoregressive PLMs.Source: §1 · p. 2
- Prot2Token maps diverse protein-prediction tasks into five categories for a general autoregressive predictor, spanning protein-level, residue-level, and protein-interaction tasks.Evaluated examples include phosphorylation-site prediction, ligand-binding sites, 3D structures, and mutation stability.Source: §1 · p. 2
- Initial analyses report synergistic performance improvements when related tasks are trained jointly.The framework inherently supports multi-task learning.Source: §1 · p. 3
- Decoder self-supervised pre-training significantly boosts binding-site prediction, and learned ligand-token relationships helped improve predictions for underrepresented ligands.The reported ligand-token analysis concerns protein-ligand binding-site prediction.Source: §1 · p. 3
1.1 Related work
- Recent unified protein systems show modality transfer, but still rely on prompt engineering for fine-grained control and lack benchmarks across the full standard task suite [23-25].Other approaches include task-specialized wrappers and a model restricted to PTM prediction.Source: §1.1 · p. 3
2 Method
Prot2Token combines pretrained encoders with a task-prompted autoregressive decoder, representing diverse protein predictions as output-token sequences. Its tokenization protocol covers several label types, and the evaluated tasks draw on established benchmarks and curated sources.
2.1 Prot2Token architecture
- Learned task tokens prompt decoder behavior, while projected encoder outputs provide context for sequential output generation.The decoder receives the task tokens and attends to the projected encoder representations.Source: §2.1 · p. 4
- A pretrained protein encoder and optional chemical encoder feed a shared autoregressive decoder, allowing diverse outputs to use one decoder in joint or independent training.Practical training combined fewer tasks than the architecture’s full multi-task potential because of computational costs.Source: §2.1 · pp. 3–4, Figure 2 · p. 5
2.2 Tokenization
- Prot2Token converts heterogeneous labels into standardized token sequences, bounded by beginning- and end-of-sequence tokens.Classification labels become discrete tokens, regression values are encoded digit by digit, sequence outputs align to input residues, and binding sites use sorted 1-based indices.Source: §2.2 · pp. 5–6, Figure 3 · p. 6
2.3 Datasets
- The evaluated datasets span protein-related regression, classification, binding-site, and sequence-to-sequence tasks from established benchmarks and curated sources.Sources include PEER, ProteinShake, CATH, AlphaFoldDB, and ProteinGym.Source: §2.3 · p. 6
3 Experiments
The experiments evaluate Prot2Token across protein-, residue-, and protein-protein-level tasks, using ESM2-650m protein encoding and optional decoder pre-training for some tasks. Results cover classification, regression, binding sites, structure prediction, and phosphorylation prediction.
3 Experiments
- Experiments initialized the protein encoder with ESM2-650m, used BARTSmiles only for ligand-affinity prediction, and varied unfrozen ESM2 layers by task.Source: §3 · p. 6
3.1 Classification
- Deeploc 2 performance significantly improved over the original method, and ER performance rose 7.5 percent with multi-task learning; EC and GO lacked Fmax, preventing direct comparisons.Source: §3.1 · p. 7
3.2 Regression
- Across regression tasks, Prot2Token outperformed PEER baselines; mutation stability improved over 51.5% against the best baseline, and fluorescence gained up to 5.6% with multi-task learning.Source: Table 4 · p. 7, §3.2 · p. 7, Table 7 · p. 8
3.3 Binding site
- With self-supervised decoder pre-training, Prot2Token achieved competitive protein-ligand binding-site accuracy across ligand types.Table 5 compares configurations using 10, 20, 30, or 41 auxiliary ligands, with a condition excluding self-supervised tasks during supervised training.Source: Table 5 · p. 8, §3.3 · p. 8
- Learned task-token embeddings correlated with biochemical properties, and using these relationships improved predictions for underrepresented ligands.Source: §3.3 · p. 8
3.4 Sequence-to-sequence
- Prot2Token-D predicts a typical 384-residue structure in 1–2 seconds on one A100, approximately 1,000x faster than AF2 with MSA (18–25 minutes).On CAMEO 2024, structure accuracy plateaus at ∼0.55 TM-score, aligning with the VQ-VAE tokenizer’s ∼0.60 TM-score reconstruction ceiling.Source: §3.4 · p. 8, Table 8 · p. 9
3.5 Other types
- Kinase phosphorylation was jointly trained with 20 self-supervised tasks; GPS 6.0 may appear artificially strong because its unpublished training split could overlap validation data.Substrate sequences longer than 1,280 amino acids were excluded, and combined substrate-kinase length was capped at 2,048 tokens.Source: §3.5 · p. 9
4 Discussion
Prot2Token maps diverse protein prediction outputs into a shared next-token framework, with evaluations supporting its use across task categories. The discussion highlights strong benchmark performance, multi-task benefits, faster structure inference, data and encoder limits, and directions toward generation.
4.1 Key Insights and Empirical Observations
- The unified next-token format performs comparably to or better than specialized or baseline models across representative task categories.The authors present this as evidence for a scalable generative framework that does not require bespoke architectures for each domain.Source: §4.1 · pp. 9–10
- Prot2Token outperformed the same-encoder ESM plus linear-probe baseline on nearly all regression benchmarks, except protein thermostability prediction.Digit-by-digit generation first establishes magnitude and then refines precision.Source: §4.1 · p. 10
- Combining 41 ligand-specific tasks yielded synergistic benefits, especially for data-scarce targets, with learned task tokens reflecting physicochemical correlations.The authors also report beneficial correlations when distinct prediction types are combined as auxiliary tasks.Source: §4.1 · p. 10
- For 3D structure prediction, inference is approximately 1000× faster than AlphaFold2, and Prot2Token exceeds its single-sequence version in quality and speed.The authors attribute efficiency to one fixed decoder step per amino acid, avoiding specialized models’ bi-directional recycling iterations.Source: §4.1 · p. 10
4.2 Limitations
- Highly uneven labels, including rare Fold classes and long-tailed binding-site indices, constrained evaluation and required validating categories rather than one all-task run.The authors characterize this sensitivity as a data limitation and report better performance when data are abundant.Source: §4.2 · p. 10
- Performance is bounded by foundational models: the 3D structure prediction TM-score plateau of approximately 0.55 reflects the VQ-VAE tokenizer’s reconstruction ceiling.The authors say closing these gaps requires higher-fidelity components rather than changes to the predictor architecture.Source: §4.2 · p. 10
4.3 Future Directions
- Future work includes higher-fidelity discrete 3D tokenizers, stochastic decoding for conformational diversity, and conditioning sequence generation on desired property tokens.The proposed generation direction would combine prediction and design within one model.Source: §4.3 · p. 11
5 Broader impact
The paper presents Prot2Token as a potentially accessible, computationally efficient framework for protein prediction, with possible applications in research and industry. It also calls for responsible validation, data practices, and safeguards as prediction and generation capabilities expand.
5 Broader impact
- Prot2Token’s unified architecture and speed improvements may broaden access to protein prediction and support real-time clinical or industrial applications.The paper names personalized medicine, real-time drug screening, and rapid biomarker discovery as possible applications.Source: §5 · p. 11
- The authors call for rigorous validation before deployment and safeguards against misuse of models that can predict or design biologically active proteins.They identify responsible data use and sharing, biosecurity, and collaboration among computational biologists, policymakers, and ethicists as considerations.Source: §5 · p. 11
A.1 Architecture
Prot2Token’s autoregressive transformer combines fused encoder context with decoder tokens, then generates task-specific labels one token at a time. Learnable task prompts select the prediction task without being penalized for reconstruction errors.
A.1 Architecture
- The decoder combines fused protein and optional chemical encoder context with a sequence of decoder input tokens.The fusion block adds learnable positional encodings to protein encoder outputs and applies a linear projection before decoding.Source: §A.1 · p. 16
- In standard decoding, a <BOS> token precedes the labels, causal attention uses prior tokens, and cross-attention accesses fused encoder features.Training accumulates negative log-likelihood loss over every decoder position.Source: §A.1 · p. 16
- For multi-task training, a learnable task token precedes labels and conditions decoder states, while its zero-weighted position is masked from loss.The remaining label positions can receive token-specific loss weights.Source: §A.1 · p. 16, Figure 4 · p. 16
- Task prompts let one decoder handle heterogeneous output formats and share parameters across tasks without duplicating specialized heads.The task token steers generation but is not penalized for reconstruction errors.Source: §A.1 · p. 16
- Inference uses greedy autoregressive decoding, selecting the most probable token at each step; stochastic strategies remain unexplored.The paper leaves evaluation of sampling strategies for future work.Source: §A.1 · p. 17
A.2 Tokenization
Prot2Token uses the respective pretrained encoders’ tokenizers for inputs and a unified discrete-token strategy for predicted labels. Output sequences have explicit beginning and ending tokens.
A.2 Tokenization
- Protein and chemical inputs use the tokenizers associated with their pretrained encoders, such as ESM2’s character-level tokenizer or BARTSmiles’ unigram tokenizer.The passage lists 33 ESM2 tokens and 1,025 BARTSmiles tokens, including special characters.Source: §A.2 · p. 17
- The output-label strategy converts diverse biological targets into discrete token sequences for consistent next-token prediction.Standardized outputs begin with <BOS> and end with <EOS>.Source: §A.2 · p. 17
- Output-label tokenization methods are organized according to the nature of the prediction task.Source: §A.2 · p. 17
A.2.1 Classification
Classification targets are encoded as discrete tokens, with the sequence format reflecting whether labels are exclusive, multiple, or hierarchical. The authors report that the preferred format differed between ER and EC tasks in early experiments.
A.2.1 Classification
- Multi-class targets use one token for the single correct class, including binary interaction labels such as Interacted.Examples include protein fold class, subcellular localization, and enzyme reaction categorization.Source: §A.2.1 · p. 17
- Multi-label targets concatenate a protein’s non-exclusive label tokens, typically in alphabetical order.The passage gives GO terms as an example of multiple labels assigned to one protein.Source: §A.2.1 · p. 17
- Hierarchical EC labels can be encoded level by level, while ER labels can use the same approach or one token per complete code.The authors report that hierarchical formatting reduced ER performance relative to multi-label formatting, while the opposite held for EC; they did not investigate this thoroughly.Source: §A.2.1 · pp. 17–18
A.2.2 Regression
Regression targets are represented as digit-by-digit token sequences rather than fixed-size bins. The authors report that reversing the digit order degraded performance, which they attribute to losing a coarse-to-fine generation bias.
A.2.2 Regression
- Digit-by-digit encoding represents each continuous value as a sequence of its sign, digits, and decimal point, with typically four decimal places used in training.Predictions are converted back to numerical form and de-normalized when target values were normalized.Source: §A.2.2 · p. 18
- Fixed-size binning can leave some bins sparse or empty when values are unevenly distributed, so Prot2Token adopts digit-by-digit encoding instead.The authors describe this representation as more granular and inherently balanced.Source: §A.2.2 · p. 18
- Reversing digit order degraded performance; the authors attribute this to losing left-to-right generation’s coarse-to-fine bias.Left-to-right generation predicts significant digits before lower-order precision.Source: §A.2.2 · p. 18
A.2.3 Sequence-to-sequence
Sequence-to-sequence tokenization produces one output token per input residue, so the target sequence length mirrors the protein length.
A.2.3 Sequence-to-sequence
- The output token sequence mirrors the input protein’s length, with a label corresponding to each residue.This format is used when predictions are made residue by residue.Source: §A.2.3 · p. 18
- Secondary-structure predictions encode each residue as a state such as α-helix (H), β-strand (E), or coil (C).An example target is {H, H, C, ...}.Source: §A.2.3 · p. 18
- For 3D structure prediction, a pretrained VQVAE encoder converts coordinates into discrete structural tokens, one for each amino acid.The passage gives 3D_i tokens as an example.Source: §A.2.3 · p. 18
A.2.4 Binding site
Binding-site outputs list the sorted 1-based residue indices involved in an interaction. The same index-based format is also used for self-supervised residue-position tasks.
A.2.4 Binding site
- Binding-site targets represent interacting residues directly as sorted 1-based indices.The passage applies this format to protein-ligand and protein-ion binding tasks.Source: §A.2.4 · p. 18
- A binding example with residues at positions 2, 3, 5, and 9 uses the target sequence {2, 3, 5, 9}.Source: §A.2.4 · p. 18
- The same index-based tokenization can train the decoder to predict positions of each standard amino acid, creating twenty self-supervised tasks.For example, serines at positions 2 and 5 in MSGLSNYT yield {2, 5}.Source: §A.2.4 · p. 18
A.2.5 Other types
Other task types include PTM prediction and outputs that combine different types. PTM outputs encode potential and confirmed modification sites, while combined outputs concatenate classification and regression tokens.
A.2.5 Other types
- This category covers PTM prediction and tasks combining different output types.Source: §A.2.5 · p. 19
- PTM outputs list 1-based potential sites, a <SEP> token, and confirmed positive sites; multi-sequence tasks still report substrate sites using kinase context.For phosphorylation, potential residues can include S, T, and Y.Source: §A.2.5 · p. 19
- TargetP 2.0 combines classification and regression by concatenating a localization-class token with a cleavage-site position.The example label is {sp, 96}, where sp denotes Signal Peptide.Source: §A.2.5 · p. 19
A.3 Dataset
This section describes task-specific datasets and evaluation setups spanning sequence properties, residue labels, interactions, structure prediction, regression, and localization. It also examines how Prot2Token’s task-token and auxiliary-task embeddings relate to biochemical and structural similarities and how auxiliary tasks may support prediction.
A.3.1 PEER benchmark
- PEER supplies five datasets for human PPI, secondary structure, fluorescence, mutation stability, and protein–ligand affinity; DeepLoc 2.0 is evaluated on an independent HPA localization test set.PEER uses homology-aware splits, while DeepLoc merges four folds for training and uses the fifth for validation.Source: §A.3.1 · p. 19, §A.3.2 · p. 20
A.3.3 PTMs
- PTM prediction uses reviewed Swiss-Prot sequences, length filtering, 40% identity clustering, and six modification types; the kinase-specific phosphorylation dataset contains 386 kinase types across 12 groups.The kinase dataset uses GPS 6.0 sites, CD-HIT clustering at 70% similarity, and separate training, validation, rare-group, and GPS test sets.Source: §A.3.3 · p. 20, §A.3.4 · pp. 20–21
A.3.5 Protein mutation stability
- Using ProteinGym’s supervised DMS substitution subset, the study followed its five-fold cross-validation protocol, conducting five independent training runs, each on a 300,000-sample subset because of computational constraints.Source: §A.3.5 · p. 21
A.3.7 3D structure prediction
- Using PPB-Affinity, the study retrieved protein sequences from RCSB PDB by protein name, retained samples with one receptor and one ligand chain, then split the final dataset 50%/20%/30% (765/180/270).Source: §A.3.8 · p. 22
A.3.9 Gene ontology
- The section also covers GO, enzyme-reaction and EC labels, DeepSF fold classes, and TargetP localization using sequence-identity-aware or homology-partitioned splits and tokenized labels.GO terms are leaf-level annotations split by 30% identity; ER sequences use 40% identity splits, EC labels are stratified by superfamily, and DeepSF covers 1,195 folds.Source: §A.3.9 · p. 22, §A.3.10 · p. 23, §A.3.11 · p. 23
A.3.14 Protein-ligand binding site
- Binding-site datasets use BioLip2 ligand annotations and PPIRef complexes, while ProteinShake structure similarity is evaluated with sequence input only under its Structure Split protocol.For PPIRef, contacting residues are defined by cross-chain centroid distances below 6 Å and splits are grouped by PDB ID; ProteinShake uses a 70% similarity threshold.Source: §A.3.14 · p. 23, §A.3.15 · p. 24, Figure 5 · p. 24
A.4.1 Classification
- Auxiliary tasks stabilized single-task fold classification and boosted Human PPI performance, although Human PPI tended to overfit; GO/EC evaluation used accuracy and F1 instead of Fmax.The GO tasks (BP, MF, and CC) and EC were trained jointly; the authors note that this prevented direct comparison with methods reporting Fmax.Source: §A.4.1 · pp. 25–26, Table 18 · p. 26
- Prot2Token recovered kinase-group membership from substrate sequences alone, achieving a macro-averaged F1 of 0.54 without kinase information at inference.The study also reports improved kinase-group clustering after encoder fine-tuning, supported by t-SNE, UMAP, and unsupervised clustering metrics.Source: §A.4.1 · pp. 26–28, Table 19 · p. 26
A.4.2 Regression
- Prot2Token improved HotProtein HP-S performance over HotProtein by roughly 2.5–3.0 pp in both correlation metrics.HP-S labels used the same min–max normalization as the other regression tasks.Source: §A.4.2 · p. 28, Table 23 · p. 28
A.4.3 Self-supervised pre-training of decoder
- Self-supervised decoder pre-training across 20 amino-acid position-prediction tasks improved phosphorylation accuracy, with broader amino-acid coverage boosting performance as auxiliary-sample volume increased.This label-free strategy was motivated by the randomly initialized decoder’s lack of inductive biases, especially for tasks with extensive label vocabularies.Source: §A.4.3 · p. 29, Table 24 · p. 29
A.4.4 Binding site
- Retaining some self-supervised samples during fine-tuning measurably improved supervised protein–ligand binding-site performance.Comparisons with other methods were constrained by unavailable or slow servers and possible overlap between their training data and the test sets.Source: §A.4.4 · pp. 31–32
A.4.5 Sequence-to-sequence
- Kinase and deaminase clustering reached ARI=0.5806 and ARI=0.7963, respectively, and Prot2Token achieved markedly higher CHI and ARI scores than other models overall.The paper also reports more distinct, interpretable embedding clusters aligned with biological labels than ESM2-650m and ESMC-600m.Source: §A.4.5 · p. 35, Table 28 · p. 35, Figure 16 · p. 36
- On CAMEO 2024, sequence-to-3D validation TM-score plateaued at ≈0.55, while the VQ-VAE tokenizer’s reconstruction upper bound was ≈0.60.The passage attributes the plateau to the tokenizer ceiling rather than insufficient Prot2Token optimization.Source: Figure 15 · p. 35
A.4.6 Other
- TargetP-2 experiments found that auxiliary tasks improved cleavage-site position and localization predictions, while retaining auxiliary samples also improved protein-kinase phosphorylation-site generalization.The TargetP-2 evaluation used five-fold cross-validation; kinase fine-tuning unfroze the encoder’s last eight blocks.Source: §A.4.6 · pp. 36–37, Table 29 · p. 36, Table 30 · p. 37
A.5 Protein-ligand binding site task tokens interpretation
- The ligand-token analysis compared clusters of the top 28 ligands by F1 score against clusters based on biochemical features to assess whether embeddings reflected real-world relationships.The broader framework uses a 640-dimensional task token per ligand in a multi-task binding-site prediction model.Source: §A.5 · p. 37, §A.5.1 · p. 37
A.5.1 Interpretation
- Ligand task-token embeddings aligned with biochemical-feature clusters, with ARI=0.447, NMI=0.614, and pairwise-accuracy=0.783.The optimal biochemical feature set included MolecularWeight, NetCharge, and RotatableBonds.Source: §A.5.1 · p. 38
- In the clustering comparison, task-token embeddings captured global relationships for ACO, ATP, FAD, GTP, NAD, and SAM, and local similarity for FE2+ and MN2+, but did not align HEM with the same ligands as feature clustering.The authors suggest that missing relationships may reflect unrepresentative ligand features or embeddings needing improvement.Source: §A.5.1 · p. 38
- Across all 41 ligands, task-token embeddings yielded ARI=0.259, NMI=0.333, and Pairwise-Accuracy=0.733; 21 ligands were represented globally, 13 locally, and 7 misrepresented.The authors attribute the metric drop to misaligned clusters from ligands with low F1 scores, while noting that meaningful information remained in their embeddings.Source: §A.5.1 · p. 38
A.5.2 Features pool creation
- The ligand feature pool comprised 26 descriptors selected using domain knowledge to describe physical, chemical, and structural properties, drawing on PubChem data and RDKit-computed descriptors.PubChem features were retrieved using Compound IDs and SMILES; RDKit converted SMILES strings into molecular objects for descriptor calculation.Source: §A.5.2 · pp. 38–39
A.5.3 Optimizing feature selection
- Feature selection compared k-means clusters of up to 13-feature combinations from 26 descriptors against high-F1 ligand embedding clusters, selecting the combination with the highest ARI.The reference embeddings were reduced from 640 dimensions to 27 principal components while retaining 99% of variance; the search evaluated approximately 39 million combinations.Source: §A.5.3 · p. 39
- The feature-combination search was also extended to all 41 ligands without the F1 threshold, and the authors report that embeddings retained significant information for lower-performing ligands.Tables 32 and 33 present the top three feature combinations for the 28-ligand and 41-ligand sets, respectively.Source: §A.5.3 · p. 39
A.5.4 Visualization
- For visualization, 640-dimensional embeddings for 41 ligands were projected to two dimensions using t-SNE and UMAP, with both methods configured to focus on local relationships.Both methods used a setting of 3 for perplexity or number of neighbors, respectively; reducing embeddings or biochemical features to 2D caused substantial information loss.Source: §A.5.1 · p. 38, §A.5.4 · p. 42
A.5.5 Multi-task learning effect
- Group 3, whose auxiliary ligands share a cluster with the targets, improves F1 more than the other groups in a comparison with similar token counts for Groups 2 and 3.The groups were used to evaluate multi-task learning on low-sample ligands GSH, CO, and AGS.Source: §A.5.5 · p. 43, Table 34 · p. 43