| --- |
| license: unlicense |
| tags: |
| - protein |
| - gene-ontology |
| - protein-function-prediction |
| - sparse-encoder |
| - cafa |
| - lafa |
| library_name: pytorch |
| --- |
| |
| # protea-sparse-knn |
|
|
| Protein function prediction by nearest-neighbour Gene Ontology transfer, |
| performed in the space of a learned sparse encoder rather than in the raw |
| embedding space. |
|
|
| This repository holds the frozen artefacts a container needs to run the |
| method: the trained encoder and the reference bank. It is a LAFA |
| submission bundle, released so the evaluation can be reproduced. |
|
|
| ## The idea |
|
|
| A protein language model places proteins near each other when their |
| sequences look alike. Function does not follow sequence that closely, and |
| the gap is what limits nearest-neighbour transfer. |
|
|
| This encoder is a single linear map from the 768-dimensional Ankh-base |
| embedding to a 2048-dimensional dictionary, keeping the 128 largest |
| components. It is trained so that the cosine between two proteins' codes |
| matches the information-content-weighted Lin semantic similarity of their |
| propagated GO closures, with hard negatives mined from embedding-near |
| neighbours. In other words, it learns a geometry in which distance means |
| functional similarity. |
|
|
| Held out on 4,000 proteins absent from the training pool, over 120,000 |
| pairs, correlation against Lin semantic similarity: |
|
|
| | representation | Spearman | |
| | --- | --- | |
| | raw Ankh-base embedding | +0.168 | |
| | encoder, dense | +0.530 | |
| | encoder, top-128 (what ships here) | +0.467 | |
|
|
| The gain of the shipped code over the raw embedding is **+0.300**. Fitting |
| inside the earlier release instead moves it to +0.293, so the result does |
| not depend on which release it was fitted in. The figures are differences |
| rather than ratios on purpose: the baseline is small enough that the same |
| two runs read as +179 and +162 percent, a spread that is an artefact of |
| the denominator. |
|
|
| The encoder also removes most of the anisotropy of the raw space: mean |
| cosine between random pairs falls from 0.551 to 0.086. |
|
|
| ## Contents |
|
|
| | file | what it is | |
| | --- | --- | |
| | `encoder.pt` | the trained encoder, 6.3 MB | |
| | `MANIFEST.json` | what it was fitted on | |
| | `VALIDATION.json` | the held-out measurement above | |
| | `codes_idx.npy` | dictionary positions of the kept components, `uint16` | |
| | `codes_val.npy` | their values, `float16` | |
| | `accessions.npy` | row order, mapping a code back to a protein | |
| | `donors.tsv.gz` | the annotations transferred by the neighbours | |
| | `BANK.json` | what the bank was built from | |
| | `EMBEDDING_RECIPE.json` | the backbone recipe the bank was built with | |
|
|
| `EMBEDDING_RECIPE.json` is not documentation. A consumer must embed its |
| queries with that exact recipe, or the queries and the bank end up in |
| different geometries and accuracy degrades without anything failing. |
|
|
| ## Training and cutoff |
|
|
| Everything is fitted once and frozen. Nothing is refitted at inference. |
|
|
| | item | value | |
| | --- | --- | |
| | backbone | `ElnaggarLab/ankh-base`, last layer, mean pooled, L2-normalised, truncated at 2048 | |
| | annotations | GOA release 227, published 2025-09-04 | |
| | ontology | `releases/2025-07-22` | |
| | bank rows | 575,503 canonical proteins, over 487,237 distinct sequences | |
| | donor annotations | 5,880,402 over 557,071 proteins | |
| | encoder training pool | 60,000 proteins | |
| | dictionary | 2048, top 128 kept | |
| | seed | 42 | |
|
|
| The ontology is the last snapshot at or before the annotation release. |
| The training driver refuses to run if that ordering is violated, because |
| a later ontology would put terms that did not exist at the cutoff on the |
| label side of the objective. |
|
|
| The earlier release, GOA 226 of 2025-05-03, is reserved as the window for |
| choosing hyperparameters and is not used to fit what ships here. |
|
|
| ## Running it |
|
|
| The container that consumes this bundle lives in the PROTEA repository |
| under `apps/lafa_sparse_knn`, with its bind-mount contract in |
| `METHOD_CARD.md`. Retrieval is exact brute force on the processor, and |
| the bank is held sparse, so the whole method fits in about 1.5 GB plus |
| the backbone. |
|
|
| ## Two things a consumer should know |
|
|
| **The bank is keyed by protein, not by sequence.** 575,503 canonical |
| accessions resolve to 487,237 distinct sequences, so proteins sharing a |
| sequence appear as separate rows carrying their own annotations. They are |
| distinct donors, which is the intent, but they are also exact ties in |
| retrieval: a query matching a sequence held by five accessions fills five |
| neighbour slots with the same vector. |
|
|
| **Self-retrieval is the parity check.** Embedding a protein that is in |
| the bank returns that protein first, at cosine 0.99 or above (measured on |
| 20 proteins: 19 return themselves and the twentieth returns an accession |
| sharing its exact sequence). If a consumer sees lower, its embedding |
| recipe has drifted from `EMBEDDING_RECIPE.json` and the results should |
| not be trusted. |
|
|
| ## Scope |
|
|
| This is the retrieval core on its own. It carries no reranker, no |
| alignment features, no taxonomic prior and no per-aspect model. It is an |
| intermediate release of a system under active development, published so |
| the containerised inference path can be exercised end to end. |
|
|
| ## Citation |
|
|
| Francisco Miguel Perez Canales. PROTEA: protein functional |
| embedding-based annotation. |
|
|
