rishivk144/enzyme-hybrid

Leak CURBER enzyme classification: DIAMOND kNN + ESM-2 neural classifier (hybrid)

Predicts one fully specified EC number per enzyme sequence. It uses only train.parquet for training, plus the pretrained ESM-2 8M protein language model (facebook/esm2_t6_8M_UR50D) as a frozen feature extractor. Runs on CPU with numpy; PyTorch is needed only for training.

Pipeline

1. Quick alignment pass. DIAMOND default mode aligns every unique input against all 186,488 training sequences.
2. Upgrade pass. DIAMOND very-sensitive mode in batches of 3,000, replacing quick-pass hits while time remains.
3. Nearest-neighbour vote. The top 25 hits vote for their recorded ECs with weight (bitscore / best bitscore) ** 16, split 1/n over a hit's n ECs.
4. Rescue pass. DIAMOND ultra-sensitive mode for inputs still without hits (about 2%). It runs only if the time left is enough for it to finish; it has a large fixed start-up cost.
5. Neural pass, for inputs that no DIAMOND pass could match (about 0.5%; about 2% if the rescue pass was skipped):
   • ESM-2 8M embedding. 6 transformer layers, run in numpy and checked against PyTorch to a maximum difference of 2e-6. The output is the mean of the last layer over residues.
   • MLP classifier. 320 → 1024 → 1024 → 4,614 ECs, with GELU and dropout 0.3. It is trained with soft targets (1/n on each of a sequence's n recorded ECs) on the embeddings of all training sequences.
6. Fallback. The most frequent training EC (2.7.7.48), used only if every stage fails.

When time is short, the very-sensitive upgrade has priority over the rescue pass, and the rescue pass over the neural pass.

Local validation

The hidden validation inputs are only about 30–55% identical to any training sequence, so I validated on held-out sequence families:

• MMseqs2 clustering of the training sequences at 50% identity and 80% coverage.
• A holdout of 1,092 whole clusters (11,804 sequences).
• Scores reweighted so each top-hit identity band counts as often as in the real validation inputs, using the official 1/n credit rule.

Model	L4	L3	L2	L1
Most frequent EC (starter baseline)	0.011	0.055	0.143	0.381
ESM-2 8M + MLP alone	0.681	0.796	0.814	0.855
ESM-2 cosine nearest neighbour alone	0.546	0.683	0.706	0.792
DIAMOND kNN, with ultra-sensitive rescue (previous submission)	0.817	0.930	0.942	0.972
This hybrid (DIAMOND with rescue pass, then ESM-2 for inputs still unmatched)	0.818

L4 by the identity of DIAMOND's best hit:

Best-hit identity	Share of inputs	DIAMOND	ESM-2 MLP
20–30%	4%	0.286	0.206
30–40%	22%	0.798	0.437
40–50%	60%	0.867	0.784
50–60%	11%	0.917	0.866
No DIAMOND hit	2%	0.000	0.079

An earlier version of this hybrid replaced the rescue pass with ESM-2. On the local holdout, ESM-2 looked better on those inputs (0.084 vs 0.037 L4 on 107 inputs, a handful of correct answers either way). On the real validation set, that version scored below the DIAMOND-only model. The comparison was too small to trust, so this version keeps the rescue pass and uses ESM-2 only where every DIAMOND pass failed. Those inputs previously got the fixed fallback EC.

What did not help

• Using ESM-2 more widely. A perfect per-input switch between DIAMOND and ESM-2 would reach 0.825. Every confidence-based switching rule tried (ESM-2 probability ≥ 0.5/0.7/0.9 combined with DIAMOND top identity < 30%/40%) scored at or below DIAMOND alone.
• A learned listwise re-ranker over DIAMOND's candidate ECs. It was an MLP on 21 features per candidate (vote shares, bitscore ratios, identity, coverage, support, EC prior, agreement at the 3-digit level), trained on cross-family searches of 4 training folds. It scored 0.8172 against 0.8167 for the hand-tuned vote. With every candidate ranked perfectly, the ceiling is only 0.852; most errors are cases where the right EC is not among DIAMOND's candidates.
• A useful signal, though not a score gain: where DIAMOND and ESM-2 agree (71% of inputs), DIAMOND's L4 is 0.927; where they disagree, it is 0.593.

Runtime and sandbox

• Grader limits. The grader allows 256 open files per process. DIAMOND therefore uses 32 temp bins (about 37 open files, largest temp file about 6 MiB), retrying with 16 or 8 bins on memory or open-file errors.
• Time. On the grader's 2 CPUs: about 40 s for the quick pass, about 240 s for the upgrade pass, about 130 s for the rescue pass, and about 15 s for the neural pass, about 425 s in total. All stages stop 70 s before the 10-minute limit.
• Logging. Sandbox limits, per-stage timings and errors go to stderr.

Files

• predict.py: inference.
• esm_numpy.py: ESM-2 forward pass in numpy.
• train.py: rebuilds the DIAMOND index and labels: python train.py --train ../data/train.parquet.
• embed.py and train_neural.py: rebuild the ESM-2 embeddings and the MLP classifier (requires PyTorch and transformers).
• weights/train.dmnd, labels.json, meta.json: the DIAMOND index and the recorded ECs.
• weights/esm2_8M.npz: ESM-2 8M encoder weights (MIT licence, Meta AI).
• weights/ec_mlp.npz, ec_mlp_classes.json: the trained classifier.
• bin/diamond: DIAMOND v2.1.11 Linux binary (GPL-3.0; source at github.com/bbuchfink/diamond).
• inputs.py, task.json, ec_vocabulary.json: unchanged from the starter.

EC vocabulary: ExPASy ENZYME release 02-Sep-2026, SIB, CC BY 4.0.