FAISS ANN search backend for proteogram similarity - #10
Open
swapnilashtekar wants to merge 4 commits into
Open
swapnilashtekar wants to merge 4 commits into
swapnilashtekar wants to merge 4 commits into
Conversation
Adds an optional approximate-nearest-neighbour search path as an alternative to brute-force cosine similarity, for corpora where the O(N^2) scan dominates runtime. - proteogram/v2/faiss_search.py: FaissIndex wrapper (IVFFlat, and IVF-PQ for very large corpora). Embeddings are L2-normalised before indexing so inner-product search is true cosine similarity. - Img2Vec gains four thin delegators (build/save/load/similarities_faiss) so scripts keep a single entry point. - measure_similarity_v2.py: --faiss, --faiss_pq, --faiss_index_file. Brute-force remains the default; the index is cached to disk and reused unless --overwrite is passed. - Adds faiss-cpu to dependencies and relocks. Both search paths share the same preprocessing function so ANN results are directly comparable to the brute-force baseline.
swapnilashtekar
marked this pull request as draft
September 12, 2026 05:58
Review fixes on top of the initial FAISS backend:
nlist defaulted to N instead of sqrt(N). The expression
`N ** _NLIST_SQRT_FACTOR ** 0.5` parses as `N ** (1.0 ** 0.5)` because ** is
right-associative, so every vector got its own Voronoi cell. FAISS warned on
every build and training cost scaled quadratically (3.7s vs 0.5s at 20k x 512).
search_all() silently truncated deep rankings. measure_similarity_v2.py asks
for the full corpus ordering so Recall@K works at any K, but an IVF search
only returns vectors in the cells it probes: 200 of 2008 requested results
came back per query. The blank trailing CSV cells read back as NaN and
crashed evaluate_methods_v2.py in target.split(','). search_all() now widens
nprobe to reach the requested depth, restoring it afterwards so one deep
search does not leave the index scanning exhaustively, and the CSV is sized
to the shortest ranking actually produced.
Added --faiss_top_k, because a full-corpus ranking forces an exhaustive scan
and is 15x slower than brute force at the current corpus size (1.67s vs
0.11s at N=2008). Capping the depth is what makes the ANN path worth using.
faiss-cpu moves to an optional "search" extra, matching the lazy import and
install hint already in faiss_search.py. Adds tests/test_faiss_search.py
covering the defaults, truncation, nprobe state, persistence and PQ
fallbacks.
Drops docs/improvements_v2.md: it plans three improvements, only one of
which is implemented here, and its snippets no longer match the shipped API.
swapnilashtekar
marked this pull request as ready for review
September 12, 2026 06:16
Adds docs/faiss_search.md and wires it into the README's Step 4 and the scripts reference table. The thing the doc leads with is that --faiss on its own is slower than the brute-force path it replaces. measure_similarity_v2.py ranks the whole corpus per query so Recall@K works at any K, and an exhaustive ranking gives an IVF index nothing to skip, so it does brute-force work plus indexing overhead. Measured on random 512-d vectors: 1.46s vs 0.11s at N=2008, and 77.37s vs 4.28s at N=13503, the released demo corpus size. Capping the depth with --faiss_top_k 100 turns that into 1.60s vs 4.28s. Also records what has not been established: recall on real proteogram embeddings. On random vectors, which have no cluster structure for the coarse quantiser to exploit, Recall@20 is 0.26 at the default nprobe and 0.71 at nlist//2. Real embeddings are clustered by fold and superfamily so this probably understates it, but it should be measured against the brute-force path before any --faiss numbers are used in a GTalign/USalign/Foldseek comparison.
swapnilashtekar
marked this pull request as draft
September 12, 2026 06:36
The performance numbers in the previous commit were measured on random gaussian vectors and were wrong in both directions. Replaced with a sweep over the released 13,503-proteogram corpus embeddings. Two errors. The claimed 2.7x speedup at top-100 compared FAISS against brute force doing a full argsort of the entire NxN matrix, when only the top K was needed; against a chunked argpartition top-K the same comparison is 1.1x. And the recall figures came from random vectors, which understate real embeddings by about 3x at the same scanned fraction (Recall@10 of 0.31 vs 0.87 at 8.7% of the corpus) because the model clusters structures by fold and superfamily, giving the coarse quantiser real structure to exploit. On the real corpus the feature looks considerably better than the random-vector numbers suggested. Brute-force top-10 is 1.82s; the index reaches 4.8x at Recall@10 0.962 (nprobe=4, 3.4% of the corpus) and 12.6x at 0.841 (nprobe=1). The shipped default nprobe = nlist//10 sits at the conservative end, 1.9x at 0.991. What has not changed is that a full-corpus ranking is the wrong way to use an ANN index: at nprobe = nlist it scans everything and is 5x slower than brute force, so --faiss_top_k is required rather than optional. Documents the random-vs-real comparison so the mistake is not repeated, and notes that the sweep should be re-run on substantially larger corpora since nlist = sqrt(N) shifts the useful nprobe range.
swapnilashtekar
marked this pull request as ready for review
September 12, 2026 06:37
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds an optional approximate-nearest-neighbour search path as an alternative to brute-force cosine similarity, for corpora where the O(N^2) scan dominates runtime.
Both search paths share the same preprocessing function so ANN results are directly comparable to the brute-force baseline.