README.md

---
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

```python
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:**

```python
from huggingface_hub import snapshot_download
path = snapshot_download("Team-LEMON/lemon")
```

```bash
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.

```bash
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:
```bash
# 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:**

```python
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](https://scop.berkeley.edu/downloads/scopeseq-2.08/) |
| `data/cath_s20.fa` | 15 043 | CATH v4.4.0 S20 — [cathdb.info](https://www.cathdb.info/wiki/doku/?id=data:index) |
| `data/cath_s20_labels.tsv` | 15 043 | CATH domain list v4.4.0 (S20 subset) — [cathdb.info](https://www.cathdb.info/wiki/doku/?id=data:index) |
| `data/scop175.fa` | 31 073 | SCOP 1.75 — [plm-zero-shot-remote-homology-evaluation](https://github.com/amoldwin/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:
```bash
# 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.