PUMA: Learning a Mutation-Aware Vocabulary of Protein Units

Burak Suyunu; Özdeniz Dolu; Ibukunoluwa Abigail Olaosebikan; Hacer Karatas Bristow; Arzucan Özgür

cs.CLMar 11, 2025

PUMA: Learning a Mutation-Aware Vocabulary of Protein Units

Authors: Burak Suyunu, Özdeniz Dolu, Ibukunoluwa Abigail Olaosebikan, Hacer Karatas Bristow, Arzucan Özgür

Organizations: Department of Computer Engineering, Boaziçi University, 34342, Bebek, Istanbul, Turkey · C. Eugene Bennett Department of Chemistry, West Virginia University, Morgantown, West Virginia 26505, United States · Department of Pharmaceutical Chemistry, Istanbul Medipol University, Faculty of Pharmacy, 34815 Beykoz, Istanbul, Turkey · Istanbul Medipol University, Research Institute for Health Sciences and Technologies (SABITA), 34810 Beykoz, Istanbul, Turkey

Abstract

Modeling protein sequences as a language has made language models a powerful tool in computational biology, yet the language itself remains poorly understood. A key step toward understanding it is identifying its constituent units. In natural languages, morphemes can occur in multiple forms; similarly, in proteins, mutations can give rise to variations of a unit that persist through evolution, forming families of related units. We introduce PUMA (Protein Units via Mutation-Aware Merging), an algorithm that learns protein units from sequence and explores their mutational variants using substitution matrices, forming a genealogy of unit families. Our results show that mutations remaining within a PUMA family are more often benign than the substitution matrix alone predicts, and that PUMA genealogy improves molecular function representations compared to treating units independently. A case study of a unit family demonstrates relatedness beyond homology. PUMA achieves competitive performance on downstream tasks when used as a protein language model tokenizer. Moreover, collapsing units into families results in a smaller embedding table and faster training. Together, these results support PUMA as a biologically grounded protein vocabulary that organizes protein units into plausible families of mutational variants. The source code is available at https://github.com/boun-tabi-lifelu/PUMA.

Figures & tables

Figure 1: Overview of the PUMA framework. Sequence Space: PUMA operates directly on 1D amino acid sequences, segmenting them into contiguous Protein Units. PUMA Engine: An iterative cycle integrating frequency and molecular evolution — the most frequent pair of units is merged into a parent; substitutions of that parent are simulated from a substitution matrix; candidates are filtered by an alignment-score cut-off and a frequency cut-off relative to the parent; survivors enter the vocabulary. PUMA Genealogy maps the hierarchy and mutational divergence of units. Hierarchical parents merge to form a new unit. A PUMA family consists of a central Mutational Parent and its resulting Mutational Children (appearing as Siblings to one another). This genealogy ensures that units and their plausible mutational variants are recognized as related entities rather than distinct fragments.

Figure 2: Distribution of PUMA family sizes at vocabulary size 51,200, for four configurations each varying one parameter off the PUMA(BLOSUM62, 0.7, 0.05) reference: substitution matrix (PAM70), alignment cut-off (0.8), and frequency cut-off (f = 0). Legend gives median and mean family size and the number of families.

Figure 3: Benign rate of PUMA family substitutions ( $P$ events) against that of non-negative-scoring substitutions ( $S$ events) under the same matrix, at vocabulary size 51,200. Colors: alignment cut-off; shapes: substitution matrix; frequency cut-off fixed at 0.05.

Figure 4: Spearman correlation between c-TF-IDF-derived GO vectors and ESM-2-based vectors for Molecular Function, for PUMA, BPE (Standard, Graph-Aware) and k-mer baselines. Violins span all PUMA configurations at a given vocabulary size, each averaged over ten held-out GO-term draws; black and red bars are the median and mean. Vocabularies of 800 are omitted because there the correlation is indistinguishable from zero for every tokenizer ( $p\approx 1$ ).

Variant effects
	PLM-ft										Scratch
	AA	800			12,800			51,200			AA	800			12,800			51,200
Task	33	BPE	PUMA	PC	BPE	PUMA	PC	BPE	PUMA	PC	33	BPE	PUMA	PC	BPE	PUMA	PC	BPE	PUMA	PC
Fluorescence	.68	.68	.68	.68	.68	.68	.66	.68	.68	.66	.30	.57	.56	.52	.60	.63	.61	.65	.63	.62
AAV	.81	.81	.83	.81	.77	.78	.73	.64	.78	.71	.33	.34	.09	.23	.45	.58	.49	.42	.52	.58
GB1	.59	.39	.40	.47	.50	.44	.48	.52	.50	.42	.29	.38	.28	.37	.16	.15	.05	.40	.44	.37

Table 1: Downstream scores on the thirteen PETA tasks. One split per task (GB1: two-vs-rest; Meltome: mixed; remote homology: fold holdout); all splits are in Tables S7 and S8 . Spearman’s

\rho

for the variant-effect tasks and Meltome, accuracy otherwise; higher is better. Means over seeds. AA: single-residue baseline (33 tokens). PUMA and PUMA_PC (PC): BLOSUM62, 0.7, 0.05. PLM-ft: pretrained encoder, fully fine-tuned. Scratch: single-layer encoder trained on each task without pretraining.

Panel of Figure 6 — Figure S1: Mutated ratio statistics for PUMA models. Each data point is the interquartile range (IQR) over the PUMA models trained with that alignment cut-off: 20 models (4 substitution matrices $\times$ 5 $f$ values) at $a=0.8$ and $0.9$ , and 15 at $a=0.7$ , where PAM250 is not trained. (a) IQR of the ratio of mutated units to all units in the vocabulary. (b) IQR of the ratio of observed mutated units to all observed units in the dataset (after segmentation).

Figure S2: Mean vocabulary identity (proportion of shared units) between RANDOM, BPE, and PUMA for vocabulary size of 51200.

Figure S3: Distribution of PUMA family sizes at vocabulary size 12,800 for the four configurations of Figure 2 .

Panel of Figure 9 — Figure S4: (a) Mean unit lengths of PUMA and BPE vocabularies at different sizes. (b) Mean observed unit lengths in PUMA and BPE vocabularies at different sizes. Observed units are units that have been observed after segmenting the dataset.

Figure S5: Vocabulary overlap (in green) shows number of units that exist in both vocabularies divided by vocabulary size. Shared unit usage (in yellow and blue) is the ratio of total number of times overlapping units are used divided by total number of units that are used (by PUMA and BPE respectively) in a segmented dataset.

Figure S6: Difference in Spearman correlation between alignment score cut-offs of 0.7 and 0.9 for Molecular Function, shown separately for standard and graph-aware c-TF-IDF. Each observation is a pair of PUMA configurations differing only in the cut-off, with each configuration summarized by its mean over ten held-out GO-term draws; positive values favor the 0.7 cut-off. Black and red bars show median and mean.

Tokenizer	33	.8k	1.6k	3.2k	6.4k	12.8k	25.6k	51.2k
PLM-ft
AA	0.671 $\pm$ .009
BPE		0.654 $\pm$ .009	0.656 $\pm$ .009	0.643 $\pm$ .011	0.642 $\pm$ .011	0.630 $\pm$ .008	0.635 $\pm$ .011	0.623 $\pm$ .031
PUMA (B62, 0.7)		0.651 $\pm$ .007	0.643 $\pm$ .015	0.636 $\pm$ .017	0.637 $\pm$ .012	0.646 $\pm$ .009	0.636 $\pm$ .011	0.639 $\pm$ .010
PUMA (B62, 0.8)		0.656 $\pm$ .013	0.646 $\pm$ .016	0.647 $\pm$ .008	0.643 $\pm$ .010	0.633 $\pm$ .012	0.642 $\pm$ .010	0.627 $\pm$ .022
PUMA (PAM70, 0.7)		0.653 $\pm$ .010	0.648 $\pm$ .018	0.653 $\pm$ .012	0.648 $\pm$ .012	0.632 $\pm$ .011	0.636 $\pm$ .011	0.637 $\pm$ .012

Table S1: Task-balanced mean scores for every tokenizer and PUMA configuration. Splits are averaged within each task and then across the 13 tasks, so tasks with several splits do not dominate. Deviations are seed standard deviations aggregated the same way. The three PUMA configurations track one another and BPE across the whole vocabulary range in both arms, so the conclusions of the main text do not depend on the configuration reported there. Differences between neighboring cells are frequently smaller than the seed standard deviation (median 0.007 in PLM-ft and 0.018 in Scratch) and should not be read as an ordering.

Variant effects
	PLM-ft							Scratch
Task	.8k	1.6k	3.2k	6.4k	12.8k	25.6k	51.2k	.8k	1.6k	3.2k	6.4k	12.8k	25.6k	51.2k
Fluorescence	.00	.00	.00	.00	.00	.00	.00	$-$ .01	$-$ .06	$-$ .02	+.04	+.04	$-$ .02	$-$ .03
AAV	+.02	$-$ .13	$-$ .06	$-$ .03	+.01	$-$ .01	+.14	$-$ .26	+.25	$-$ .01	$-$ .12	+.13	+.02	+.10
GB1	$-$ .02	$-$ .05	$-$ .04	$-$ .03	.00	.00	.00	$-$ .04	$-$ .02	+.02	$-$ .08	$-$ .02	$-$ .03	$-$ .01
Stability	.00	+.02	+.01	.00	+.06	$-$ .01	$-$ .02	+.02	+.03	$-$ .05	$-$ .03	+.13	$-$ .02	$-$ .04

Table S2: Per-task difference between PUMA and BPE. Each cell is PUMA

-

BPE for one task and vocabulary size (seed means, averaged over splits within a task); positive values favor PUMA. PUMA: BLOSUM62, 0.7, 0.05. PLM-ft: pretrained encoder, fully fine-tuned. Scratch: single-layer encoder trained on each task without pretraining.

PUMA vocab size	PUMA_PC	Reduction	PUMA_PC	Reduction	PUMA_PC	Reduction
	BLOSUM62, 0.7		BLOSUM62, 0.8		PAM70, 0.7
800	389	2.1 $\times$	531	1.5 $\times$	404	2.0 $\times$
1,600	513	3.1 $\times$	896	1.8 $\times$	579	2.8 $\times$
3,200	730	4.4 $\times$	1,446	2.2 $\times$	873	3.7 $\times$
6,400	1,044	6.1 $\times$	2,168	3.0 $\times$	1,315	4.9 $\times$
12,800	1,481	8.6 $\times$	3,375	3.8 $\times$	2,051	6.2 $\times$

Table S3: Vocabulary contraction under parent collapse. A sequence is tokenized with the original PUMA tokenizer and each token is then remapped to the mutational parent of its family; singletons are unchanged. Because most PUMA tokens are rare mutational children of a shared parent, the trained embedding table shrinks sharply, and increasingly so with vocabulary size. The contraction is strongest for the configurations that spend most of their budget on mutational variants.

Tokenizer	33	.8k	1.6k	3.2k	6.4k	12.8k	25.6k	51.2k
Wall-clock hours
AA	18.18
BPE		11.82	11.12	10.78	10.77	11.11	15.56	16.97
PUMA		11.91	11.41	11.18	11.08	11.31	15.45	16.85
PUMA_PC		11.84	11.22	10.86	10.55	10.24	10.10	9.94
Speedup vs. AA

Table S4: Pretraining wall-clock cost. Hours for one epoch of masked-language-model pretraining over UniRef50 (2024-06), and the corresponding speedup relative to the single-residue baseline. PUMA and PUMA_PC rows are averaged over the configurations of Table S1 with a recorded runtime (one to three per vocabulary size), whose individual times differ by at most 6% at any vocabulary size. Every multi-residue tokenizer is faster than AA, but BPE and full PUMA reach a minimum near 3,200–6,400 and slow again beyond it, because the masked-language-model objective projects onto the full vocabulary at every masked position. Parent collapse leaves the segmentation unchanged while shrinking that projection, and is the only vocabulary whose pretraining time falls monotonically across the whole range.

Tokenizer	33	.8k	1.6k	3.2k	6.4k	12.8k	25.6k	51.2k
Compression $\bar{k}$ (residues per token)
AA	1.00
BPE		1.93	2.13	2.30	2.45	2.62	2.80	3.02
PUMA		1.89	2.07	2.24	2.39	2.55	2.74	2.94
PUMA_PC		1.89	2.07	2.24	2.39	2.55	2.74	2.94
PLM-ft — speedup vs. AA

Table S5: Compression and downstream wall-clock cost. Top block: mean compression

\bar{k}

, residues encoded per token, averaged over the 18 downstream task/splits. Parent collapse relabels an unchanged segmentation, so PUMA and PUMA_PC compress identically and differ only in the size of the embedding and output layers. Bottom block: fine-tuning speedup relative to AA, computed per task/split as the ratio of the seed-mean AA runtime to the seed-mean tokenizer runtime and aggregated as a geometric mean over the 18 task/splits. PUMA and PUMA_PC rows are averaged over the three configurations. Fine-tuning carries no vocabulary-sized output layer, so its cost is set by sequence length alone: downstream speedups rise with vocabulary and never invert as they do in pretraining.

Group	Task	Description	Metric	Splits	Split names
Variant effects	Fluorescence	Log-fluorescence of GFP mutants; trained on low-order mutants and tested on higher-order ones.	Spearman $\rho$	1	default
	AAV	Fitness of AAV2 capsid (VP1) variants mutated in positions 561–588; trained on variants with at most two mutations and tested on the rest.	Spearman $\rho$	1	two-vs-many
	GB1	Binding fitness of GB1 variants at four epistatic positions. Splits train on variants with at most two or three mutations, or on low-fitness variants, and test on the rest.	Spearman $\rho$	3	two-vs-rest † , three-vs-rest, low-vs-high
	Stability	Stability score of designed mini-proteins; tested on single-mutant neighbors of the most stable designs.	Spearman $\rho$	1	default
	SolMut CS / BLAT / LGK	Change in solubility upon mutation (SoluProtMutDB) for chalcone synthase, TEM-1 $\beta$ -lactamase and levoglucosan kinase; one task per protein, random 8:1:1 split.	Spearman $\rho$	1 each	default (each)
Protein properties	Meltome	Thermostability from mass-spectrometry melting curves of proteins from 13 species.	Spearman $\rho$	2	mixed split † , human

Table S6: Downstream tasks. Task definitions, splits and metrics follow the PETA benchmark [ 24 ] ; we use 13 of its tasks, evaluated on 18 task/splits. The three SolMut tasks are listed together because they share a definition. † Split reported in Table 1 for tasks with more than one split. ∗ Multi-label accuracy, averaged over the ten compartments. All metrics are higher-is-better.

Explore similar work

ProtSent: Protein Sentence Transformers
May 7, 2026Dan Ofer, Oriel Perets, Michal Linial +1Language Modeling Transformer Architectures
Department of Computer and Information Science Ben-Gurion University of the Negev
Better Protein Function Prediction by Modeling Survivorship Bias
May 7, 2026Zhongmou Chao, Poompol Buathong, Ekaterina Selivanovitch +2Survival Predictions
Smith School of Chemical and Biomolecular Engineering, Cornell University, USA · Center for Applied Mathematics, Cornell University, USA · School of Operations Research and Information Engineering, Cornell University, USA
ProtLingo: Efficient Protein Language Modeling via Conditional Memory and Expert Routing
Sep 7, 2026Mingrui Li, Sixian Shen, Minzhang Li +4Language Modeling
ShanghaiTech University

May 7, 2026, cs.LGJ/K move · Enter open · S save

ProtSent: Protein Sentence Transformers

Dan Ofer, Oriel Perets, Michal Linial, Nadav Rappoport

Department of Computer and Information Science Ben-Gurion University of the Negev

Language Modeling Transformer Architectures

Open paper AskPDF arXiv

Protein language models (pLMs) produce per-residue representations that capture evolutionary and structural information, yet their mean-pooled sequence embeddings are not explicitly trained to reflect functional, evolutionary or structural similarity between proteins. We present Protein Sentence Transformers (ProtSent), a contrastive fine-tuning framework for adapting PLMs into general-purpose embedding models. ProtSent trains with MultipleNegativesRankingLoss across five protein-pair datasets: Pfam families, structurally derived hard negatives, AlphaFold DB structural pairs, and StringDB protein--protein interactions, and Deep Mutational Scanning data. We evaluate on 23~downstream tasks using frozen embeddings with a k-nearest-neighbor probe to measure embedding neighborhood quality. On ESM-2 150M, ProtSent improves 15 of 23 tasks, with gains of +105% on remote homology detection, +17% on variant effect prediction, and +19.9% Recall@1 on SCOPe-40 structural retrieval. The 35M variant improves 16 of 23 tasks with +40.5% on remote homology and +15.5% Recall@1 on SCOPe-40. Contrastive fine-tuning restructures the embedding space to better capture protein function and structure, without any task-specific supervision. We release the models, public data, and training recipe and code.

Abstract

Explore similar work

ProtSent: Protein Sentence Transformers

Better Protein Function Prediction by Modeling Survivorship Bias

ProtLingo: Efficient Protein Language Modeling via Conditional Memory and Expert Routing