LEMON-ZEST: Evolution-Informed Tokenization for Efficient Protein Language Modeling

Biswajit Banerjee; Claudia Alvarez Carreno; Anton S. Petrov

cs.LGSep 29, 2026

LEMON-ZEST: Evolution-Informed Tokenization for Efficient Protein Language Modeling

Authors: Biswajit Banerjee, Claudia Alvarez Carreno, Anton S. Petrov

Organizations: Georgia Institute of Technology · University College London

Abstract

Protein Language Models (PLMs) have made remarkable progress following scaling laws established in natural language processing across sequence- and structure-based tasks, yet the potential of tokenization remains underexploited. Unlike human language, proteins preserve structure despite extensive sequence variation a property standard tokenization strategies fundamentally fail to capture. We introduce ZEST (Zoned Encoding of Sequence Traits), an evolution-informed vocabulary derived from conserved regions of multiple sequence alignments. ZEST allows embedding domain-level biological priors directly at the tokenization stage rather than learning them implicitly through scale. ZEST natively compresses sequences to an average token length of 4 residues, enabling our model to process 4,000 residues within a standard 1024-token context window. Building on this, we present LEMON (Layered Extraction of Molecular Ordering from Nature), a compact 200M-parameter sequence-based model for detection of remote homology between protein sequences trained on a single H100 GPU for one week. Despite its modest size, LEMON outperforms state-of-the-art models ranging from 600M to 3B parameters. Our results demonstrate that evolution-informed tokenization can substitute for massive parameter scaling, opening a new direction for efficient, biologically-grounded protein representation learning. All code, model weights, and results are publicly available under the MIT license.

Figures & tables

Figure 1: Overview of the proposed architecture and methodology. (A) The biological prior-injected vocabulary generation process. (B) The dual-head Y-architecture model design. (C) A reduced-dimensionality projection of the latent representations of protein domains, illustrating clear spatial separation between alpha helices, beta sheets, and their combinations.

Figure 2: Multi-seed analysis of vocabulary coverage across one million sequences across epochs.

Figure 3: (a) AUROC computed on SCOP dataset plotted against Sequence Identity Threshold; each point corresponds to an evaluation where the retrieval pool is filtered to proteins sharing at most t% sequence identity with the query (x-axis). Dashed lines indicate models requiring 3Di structural tokens as input (SaProt, ProstT5) and a dotted line represents HHblits. (b) Mean fold-level AUROC averaged across all six thresholds in comparison to model size. HHblits is a non-parametric, HMM profile-based comparison.

Model	Params	Time	CATH S20				SCOPe				SCOP
			Architecture		Topology		Fold		Superfamily		Fold		Superfamily
			AUROC	mAP	AUROC	mAP	AUROC	mAP	AUROC	mAP	AUROC	mAP	AUROC	mAP
LEMON	200M	5m	0.825	0.356	0.898	0.320	0.903	0.347	0.959	0.555	0.904	0.287	0.948	0.421
ProtTucker	3,000M	9m	0.783	0.250	0.874	0.293	0.876	0.263	0.971	0.693	0.855	0.212	0.961	0.541
SaProt	650M	$\times$ †	0.619	0.143	0.658	0.067	0.717	0.075	0.859	0.204	0.670	0.055	0.778	0.062
ProtTrans	3,000M	9m	0.608	0.126	0.699	0.091	0.701	0.092	0.914	0.386	0.710	0.085	0.889	0.233

Table 1: Complete benchmark results across all datasets and classification levels. The Superfamily level measures detection of proteins sharing the same fold. Time = total GPU wall-clock on single H100. † Additionally requires ESMFold structure prediction as a pre-processing step.

Figure 4: Multi-seed test time augmentation effect.

Tokenizer	ROC-AUC	mAP	fold ROC	$T_{\text{pre}}$ (mins)	$T_{\text{ft}}$ (mins)
ZEST (32K)	0.739 $\pm$ 0.017	0.068 $\pm$ 0.002	0.749 $\pm$ 0.014	304.7	985.5
BPE (32K)	0.729 $\pm$ 0.008	0.036 $\pm$ 0.002	0.743 $\pm$ 0.008	278.0	1030.0
Char (25)	0.708 $\pm$ 0.011	0.059 $\pm$ 0.002	0.721 $\pm$ 0.010	650.7	1043.1

Table 2: SCOPe (10%-identity) retrieval benchmark across three seeds with identical architecture and effective batch size of 1024; the total times are averaged across three runs of each category.

Task	Metric	LEMON (200M)	ESM-2 (650M)	ProtTucker (3B)
Circular Permutation	AUROC	0.779	0.476	0.720
	AUPRC	0.285	0.081	0.238
EC Number	F-max	0.821	0.800	0.878
GO Mol. Function	F-max	0.639	0.603	0.635
	AUPR	0.625	0.584	0.616
Domain Boundary	NDO	0.758	0.771	0.750

Table 3: Zero-shot generalization across four protein tasks. All methods use frozen embeddings with cosine-similarity

k

-NN label transfer (

K\!=\!15

) or pairwise classification—no task-specific fine-tuning. The table shows that at a fraction of its parameter size, LEMON matches or exceeds SOTA zero-shot performance.

Explore similar work

Protein Language Models Diverge from Natural Language: Comparative Analysis and Improved Inference
Feb 24, 2026Anna Hart, Chi Han, Jeonghwan Kim +2Protein Language Models Protein Representation Learning
Siebel School of Data Science and Computing, University of Illinois Urbana-Champaign, 201 N Goodwin Ave, Urbana, IL 61801, USA · DOE Center for Advanced Bioenergy and Bioproducts Innovation, University of Illinois Urbana-Champaign, 1206 W Gregory Drive, Urbana, IL 61801, USA · Chemical and Biomolecular Engineering, University of Illinois Urbana-Champaign, 600 S Matthews Ave, Urbana, IL 61801, USA
Reverse Distillation: Consistently Scaling Protein Language Model Representations
Mar 8, 2026Darius Catrina, Christian Bepler, Samuel Sledzieski +1Protein Language Models Protein Representation Learning
Department of Computer Science, Duke University, Durham, NC, USA · Depts of Biostats. and Bioinformatics; and Cell Biology, Duke University, Durham, NC, USA
PUMA: Learning a Mutation-Aware Vocabulary of Protein Units
Mar 11, 2025Burak Suyunu, Özdeniz Dolu, Ibukunoluwa Abigail Olaosebikan +2Protein Language Models Protein Representation Learning
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 +1

Feb 24, 2026J/K move · Enter open · S save

Protein Language Models Diverge from Natural Language: Comparative Analysis and Improved Inference

Anna Hart, Chi Han, Jeonghwan Kim, Huimin Zhao+1

Siebel School of Data Science and Computing, University of Illinois Urbana-Champaign, 201 N Goodwin Ave, Urbana, IL 61801, USA · DOE Center for Advanced Bioenergy and Bioproducts Innovation, University of Illinois Urbana-Champaign, 1206 W Gregory Drive, Urbana, IL 61801, USA · Chemical and Biomolecular Engineering, University of Illinois Urbana-Champaign, 600 S Matthews Ave, Urbana, IL 61801, USA

Protein Language Models Protein Representation Learning Molecular Property Prediction

Open paper AskPDF arXiv

Modern Protein Language Models (PLMs) apply transformer-based model architectures from natural language processing to biological sequences, predicting a variety of protein functions and properties. However, protein language has key differences from natural language, such as a rich functional space despite a vocabulary of only 20 amino acids. These differences motivate research into how transformer-based architectures operate differently in the protein domain and how we can better leverage PLMs to solve protein-related tasks. In this work, we begin by directly comparing how the distribution of information stored across layers of attention heads differs between the protein and natural language domain. Furthermore, we adapt a simple early-exit technique-originally used in the natural language domain to improve efficiency at the cost of performance-to achieve both increased accuracy and substantial efficiency gains in protein non-structural property prediction by allowing the model to automatically select protein representations from the intermediate layers of the PLMs for the specific task and protein at hand. We achieve performance gains ranging from 0.4 to 7.01 percentage points while simultaneously improving efficiency by over 10 percent across models and non-structural prediction tasks. Our work opens up an area of research directly comparing how language models change behavior when moved into the protein domain and advances language modeling in biological domains.

Abstract

Explore similar work

Protein Language Models Diverge from Natural Language: Comparative Analysis and Improved Inference

Reverse Distillation: Consistently Scaling Protein Language Model Representations

PUMA: Learning a Mutation-Aware Vocabulary of Protein Units