LEMON: Layered Extraction of Molecular Ordering from Nature
LEMON is a protein sequence encoder trained with hierarchical contrastive
learning for family and fold similarity search, using the ZEST tokenizer
(Zoned Encoding of Sequence Traits). Submitted anonymously for double-blind
peer review.
Architecture
Component Details Encoder 24-layer transformer, 768-d, 12 heads, SwiGLU FFN (ff_mult=4), RoPE with linear scaling Pooling Learned-query multi-head attention aggregator → 768-d sequence vector Projector Bottleneck MLP (768 → 768 → 384-d), L2-normalised output Tokenizer ZEST 32K (Zoned Encoding of Sequence Traits) — greedy max-match trie over biochemically-substitutable amino-acid n-gram clusters Context 1,024 tokens (linear RoPE scaling for longer sequences) Dropout 0.04
Parameter breakdown (203.72M total):
Module Params Core transformer 194.52M Attention aggregator 2.95M Profile expansion head 2.41M Global position embedding 2.36M Projector 1.48M Total 203.72M
Quickstart
1 import torch
2 from huggingface_hub import snapshot_download
3 import sys , os
4
5 path = snapshot_download ( "Team-LEMON/lemon" )
6 sys . path . insert ( 0 , path )
7
8 from modeling_lemon import LemonEncoder
9 from tokenization_zest import ZESTTokenizer
10
11 tok = ZESTTokenizer . from_pretrained ( path )
12 model = LemonEncoder . from_pretrained (
13 os . path . join ( path , "model.safetensors" ) ,
14 os . path . join ( path , "config.json" ) ,
15 )
16 model . eval ( )
17
18 seqs = [ "MKTAYIAKQRQISFVKSHFSRQ" , "ACDEFGHIKLMNPQRSTVWY" ]
19 enc = tok . batch_encode_plus ( seqs , max_length = 512 , padding = True )
20 with torch . no_grad ( ) :
21 emb = model . embed ( enc [ "input_ids" ] , enc [ "attention_mask" ] ) # [2, 384]
22 print ( emb . shape ) # torch.Size([2, 384])
23
24 sim = model . similarity ( emb [ : 1 ] , emb [ 1 : ] )
25 print ( "cosine-like similarity:" , sim . item ( ) )
Reproducing Table 1
The eval_retrieval.py script and all three benchmark datasets are bundled in this
repo. No external downloads required.
Run all three datasets in one command:
1 from huggingface_hub import snapshot_download
2 path = snapshot_download ( "Team-LEMON/lemon" )
1 cd /path/to/snapshot
2 python eval_retrieval.py # runs SCOPe + SCOP + CATH-S20
3 python eval_retrieval.py --scope # SCOPe only
4 python eval_retrieval.py --cath # CATH-S20 only
5 python eval_retrieval.py --scop # SCOP only
Test-Time Augmentation (TTA) with Trie-Dropout:
TTA improves retrieval by averaging embeddings from multiple stochastic tokenizations.
python eval_retrieval.py --dropout 0.45 --tta 5 # 5 stochastic passes, averagedTTA Gain (SCOPe, seed=42):
Level Metric Baseline TTA (d=0.45, k=5) Gain fold AUROC 0.9025 0.9080 +0.0055 fold mAP 0.3067 0.3197 +0.0130 superfamily AUROC 0.9443 0.9519 +0.0076 superfamily mAP 0.4700 0.4803 +0.0103
To reproduce:
1 # Baseline
2 python eval_retrieval.py --scope --seed 42
3
4 # With TTA
5 python eval_retrieval.py --scope --seed 42 --dropout 0.45 --tta 5
Or from a Jupyter notebook:
1 import sys
2 from huggingface_hub import snapshot_download
3
4 path = snapshot_download ( "Team-LEMON/lemon" )
5 sys . path . insert ( 0 , path )
6
7 from eval_retrieval import run_benchmark , display_results
8
9 results = run_benchmark ( repo = path , seed = 42 ) # deterministic with seed=42
10 display_results ( results )
11
12 # With TTA:
13 # results = run_benchmark(repo=path, seed=42, dropout=0.1, tta_passes=8)
Expected output (seed=42, deterministic):
Dataset Level AUROC mAP SCOPe fold 0.9025 0.3066 SCOPe superfamily 0.9443 0.4700 CATH-S20 architecture 0.8871 0.3128 CATH-S20 topology 0.9580 0.5381 SCOP fold 0.9062 0.2919
Results are deterministic with --seed 42 (default).
CATH uses Architecture/Topology levels; SCOP/SCOPe uses Fold/Superfamily.
Bundled dataset provenance:
Circular Permutation Detection (CIRPIN SCOPe40)
Zero-shot detection of circularly permuted protein pairs using cosine similarity of LEMON embeddings.
Benchmark: CIRPIN SCOPe40 — 18,127 pairs (1,967 positive CP pairs) from ASTRAL SCOPe 2.08 at 40% identity.
Results (seed=42):
Configuration AUROC AUPRC Accuracy Baseline 0.7413 0.3035 0.8990 TTA (d=0.45, k=5) 0.7576 0.3066 0.8987 Gain +0.0163 +0.0031 -
TTA improves CP detection by averaging embeddings over multiple stochastic tokenizations.
To reproduce:
1 # Baseline
2 python eval_circular_permutation.py --fasta data/cirpin/scope40.fa --pairs data/cirpin/pairs.tsv --seed 42
3
4 # With TTA
5 python eval_circular_permutation.py --fasta data/cirpin/scope40.fa --pairs data/cirpin/pairs.tsv --seed 42 --dropout 0.45 --tta 5
Requirements
torch>=2.0
safetensors
huggingface_hubNotes
Input sequences should be standard single-letter amino-acid strings.
The tokenizer handles unknown characters via <MASK> token fallback.
model.embed() returns L2-normalised embeddings; use dot product for
cosine similarity.model.similarity() applies a learned temperature scalar. 
