license: mit
tags:
- protein
- biology
- sequence-encoder
- contrastive-learning
- lemon
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
import torch
from huggingface_hub import snapshot_download
import sys, os
path = snapshot_download("Team-LEMON/lemon")
sys.path.insert(0, path)
from modeling_lemon import LemonEncoder
from tokenization_zest import ZESTTokenizer
tok = ZESTTokenizer.from_pretrained(path)
model = LemonEncoder.from_pretrained(
os.path.join(path, "model.safetensors"),
os.path.join(path, "config.json"),
)
model.eval()
seqs = ["MKTAYIAKQRQISFVKSHFSRQ", "ACDEFGHIKLMNPQRSTVWY"]
enc = tok.batch_encode_plus(seqs, max_length=512, padding=True)
with torch.no_grad():
emb = model.embed(enc["input_ids"], enc["attention_mask"]) # [2, 384]
print(emb.shape) # torch.Size([2, 384])
sim = model.similarity(emb[:1], emb[1:])
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:
from huggingface_hub import snapshot_download
path = snapshot_download("Team-LEMON/lemon")
cd /path/to/snapshot
python eval_retrieval.py # runs SCOPe + SCOP + CATH-S20
python eval_retrieval.py --scope # SCOPe only
python eval_retrieval.py --cath # CATH-S20 only
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, averaged
TTA 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:
# Baseline
python eval_retrieval.py --scope --seed 42
# With TTA
python eval_retrieval.py --scope --seed 42 --dropout 0.45 --tta 5
Or from a Jupyter notebook:
import sys
from huggingface_hub import snapshot_download
path = snapshot_download("Team-LEMON/lemon")
sys.path.insert(0, path)
from eval_retrieval import run_benchmark, display_results
results = run_benchmark(repo=path, seed=42) # deterministic with seed=42
display_results(results)
# With TTA:
# 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:
| File | Sequences | Original source |
|---|---|---|
data/scope_10_2.08.fa |
7 117 | SCOPe 2.08, 10% seq-id — scop.berkeley.edu |
data/cath_s20.fa |
15 043 | CATH v4.4.0 S20 — cathdb.info |
data/cath_s20_labels.tsv |
15 043 | CATH domain list v4.4.0 (S20 subset) — cathdb.info |
data/scop175.fa |
31 073 | SCOP 1.75 — plm-zero-shot-remote-homology-evaluation |
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:
# Baseline
python eval_circular_permutation.py --fasta data/cirpin/scope40.fa --pairs data/cirpin/pairs.tsv --seed 42
# With TTA
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_hub
Notes
- 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.