Raghav06123/SecondsubDA3410

Leak CURBER enzyme classification — homology-graph neural network

A graph attention network predicts one EC number per protein from the protein's homology neighbourhood. Retrieval supplies the graph; the neural network makes every prediction.

Pipeline

1. Retrieval (builds the graph): TF-IDF k-mer search over the 186k unique training enzymes (top 1,000 by 5-mers + top 300 by 4-mers), then Smith-Waterman local alignment (BLOSUM62, gap 11/1, pyopal). The query and its 20 best hits form a graph.

2. Graph neural network (gnn.py): each neighbour is a node with 11 features:

   • 7 edge features from the alignment: score / query self-score, score / neighbour self-score, log(score / best score), rank, log length ratio, best score / query self-score, log(1 + score);
   • 4 agreement features: how much of the neighbourhood shares this neighbour's EC at levels 1, 2, 3, 4.

   Features -> MLP (11 -> 128 -> 128) -> 2 transformer layers (4 heads, pre-norm) that pass messages between all neighbours -> a weight per neighbour (output head plus a residual alpha * log(score/best), alpha learned). The EC with the largest total softmax weight is the prediction.

Training

• Graph for every training enzyme: its top-60 Smith-Waterman neighbours among the other training enzymes.
• To mimic test conditions (validation enzymes are distant from train: median best alignment = 0.38 of self-score, vs 0.92 inside train), each training example hides neighbours closer than a threshold drawn from the validation distribution, then keeps the 20 best remaining.
• Loss: negative log of the softmax weight on neighbours carrying a true EC. AdamW, lr 1e-3, 5 epochs, PyTorch; weights exported to NumPy (gnn.npz, 1.1 MB, ~280k parameters). The residual starts the network at a hand-tuned similarity vote, so training learns corrections to it.

Local validation (4,915 distant training enzymes, weighted to match validation)

Metric	Hand-tuned vote (v2)	Graph network
L1	0.947	0.951
L2	0.917	0.920
L3	0.903	0.906
L4	0.749	0.749

Leaderboard (public validation): v2 vote L4 0.773.

Alternatives tested: ESM2-8M embeddings (too slow on 2 CPUs: ~76 min for validation); a k-mer embedding-bag classifier alone (L4 0.53 locally).

Runtime

11,367 validation inputs in 357 s (graph network step: 9 s), 4.1 GB RAM, 2 vCPU. Queries not aligned by 480 s keep a 5-mer nearest-neighbour prediction, so the run stays under 10 minutes.

Files

• predict.py — inference entry point
• gnn.py, gnn.npz — graph neural network (NumPy forward pass, weights)
• kmer.py — sparse k-mer index
• db.parquet — unique training sequences with their ECs
• train.py — builds db.parquet from train.parquet
• inputs.py, task.json, ec_vocabulary.json — from the starter