Journal of Cheminformatics 2024

PETA

A systematic study of how sub-word tokenization — BPE and Unigram versus per-amino-acid vocabularies — shapes protein pretraining and transfer to downstream applications.

Yang Tan, Mingchen Li, Ziyi Zhou, Pan Tan, Huiqun Yu, Guisheng Fan, Liang Hong

Paper Code
PETA overview: sub-word tokenization for protein transfer learning
PETA compares per-residue and sub-word protein tokenizers under matched pretraining and fine-tuning protocols.

Why tokenization matters

Most protein language models tokenize sequences one amino acid at a time. That choice is simple and biologically natural, but it ignores a lesson from NLP: sub-word units can compress repetitive motifs, change context length, and alter what a Masked LM learns to represent.

PETA asks a practical question for transfer learning: when you hold architecture and pretraining budget roughly fixed, how do BPE / Unigram vocabularies of different sizes compare to a classic per-AA tokenizer on real downstream protein tasks?

What PETA provides

The project (also referred to as Venus-PETA in the codebase) releases a family of Masked LM checkpoints and tokenizers on Hugging Face under AI4Protein/deep_*, spanning:

  • Per-AA baseline — AI4Protein/deep_base (vocab size 33)
  • BPE — vocab sizes 50 → 3200 (deep_bpe_*)
  • Unigram — vocab sizes 50 → 3200 (deep_unigram_*)

Sub-word models merge frequent residue patterns into multi-letter tokens (e.g. SLG, AK, PF), shortening tokenized length relative to character-level encoding while keeping special start/end tokens for Transformer training. 

from transformers import AutoTokenizer, AutoModelForMaskedLM

tokenizer = AutoTokenizer.from_pretrained("AI4Protein/deep_bpe_3200")
model = AutoModelForMaskedLM.from_pretrained("AI4Protein/deep_bpe_3200")

sequence = "MSLGAKPFGEKKFIEIKGRRM"
print(tokenizer.tokenize(sequence))
# ['M', 'SLG', 'AK', 'PF', 'GE', 'KK', 'FI', 'EI', 'KG', 'RR', 'M']

Downstream evaluation

Beyond releasing pretrained models, PETA packages a fine-tuning harness (peta/train.py) for transfer experiments: choose a dataset split, pooling head (e.g. attention1d), precision, and whether to fine-tune the full backbone or only the head. Benchmark dataset archives and split definitions live with the ProteinPretraining repository.

The takeaway is evaluative rather than a single “best tokenizer forever”: sub-word schemes change sequence compression and inductive bias, and their effect on transfer depends on the task and vocabulary size — PETA supplies the controlled comparison to measure that trade-off.

Citation

If you use PETA, please cite:

@article{tan2024peta,
  title={PETA: evaluating the impact of protein transfer learning with sub-word tokenization on downstream applications},
  author={Tan, Yang and Li, Mingchen and Zhou, Ziyi and Tan, Pan and Yu, Huiqun and Fan, Guisheng and Hong, Liang},
  journal={Journal of Cheminformatics},
  volume={16},
  number={1},
  pages={92},
  year={2024},
  publisher={Springer},
  doi={10.1186/s13321-024-00884-3}
}