# PUMA: Learning a Mutation-Aware Vocabulary of Protein Units

arXiv: https://arxiv.org/abs/2503.08838 · HTML: https://mlsanity.com/paper/2503.08838/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 3 of its 26 references that are in the map.
- First posted on arXiv 11 Mar 2025.

## Builds on (most cited first, 3 of 3)

- [Neural Machine Translation of Rare Words with Subword Units](https://mlsanity.com/paper/1508.07909/) (2015)
- [BERTopic: Neural topic modeling with a class-based TF-IDF procedure](https://mlsanity.com/paper/2203.05794/) (2022)
- [Linguistic Laws Meet Protein Sequences: A Comparative Analysis of Subword Tokenization Methods](https://mlsanity.com/paper/2411.17669/) (2024)

## Built on by (most cited first, 1 of 1)

- [LEMON-ZEST: Evolution-Informed Tokenization for Efficient Protein Language Modeling](https://mlsanity.com/paper/2609.37675/) (2026)

## Abstract (arXiv)

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.