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.