[fix] Avoid materializing the full similarity matrix in mine_hard_negatives without FAISS#3816
Merged
tomaarsen merged 4 commits intoJun 12, 2026
Conversation
Contributor
There was a problem hiding this comment.
Pull request overview
This PR fixes excessive RAM usage in mine_hard_negatives when use_faiss=False by avoiding construction of the full (n_queries, n_corpus) similarity matrix. It implements query-axis batching for the non-FAISS similarity search path (mirroring the FAISS batching approach), adds parameter validation, and introduces a regression test to ensure batched behavior is equivalent to a single-batch run.
Changes:
- Batch non-FAISS similarity computation +
topkover query chunks controlled byfaiss_batch_size. - Add
faiss_batch_size <= 0validation and make search progress bars respectverbose. - Add a regression test asserting chunking behavior and output equivalence vs single-batch execution.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
sentence_transformers/util/hard_negatives.py |
Implements batched non-FAISS similarity search, improves docstring/validation, and gates progress bars by verbose. |
tests/util/test_hard_negatives.py |
Adds a regression test validating non-FAISS chunking behavior and mined output equivalence. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Co-authored-by: Copilot Autofix powered by AI <[email protected]>
tomaarsen
reviewed
Jun 12, 2026
tomaarsen
left a comment
Member
There was a problem hiding this comment.
Thanks for this. I've made some modifications re. verbose, which I'll merge.
Co-authored-by: Tom Aarsen <[email protected]>
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
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.
Summary
Fixes #3164.
With
use_faiss=False,mine_hard_negativescurrently computes the full(n_queries, n_corpus)similarity matrix before runningtopk:For the scale mentioned in #3164 (Natural Questions, roughly 100k queries x 100k
corpus entries), that float32 matrix is about 40GB before overhead, so it can
easily exhaust 16-32GB machines. The FAISS branch already avoids this by searching
in
faiss_batch_sizechunks over the query axis.This PR applies the same query-axis batching to the non-FAISS branch: compute
similarity +
topkper chunk, then concatenate the per-chunk results. Since eachrow is ranked independently, this preserves the mined negatives while bounding
the intermediate similarity matrix to
faiss_batch_size * len(corpus).I also generalized the
faiss_batch_sizedocstring, added a clearValueErrorfor
faiss_batch_size <= 0, and made both search progress bars respectverbose.Local peak RSS delta with
all-MiniLM-L6-v2on CPU and a synthetic corpus:In local equivalence checks, the mined outputs matched exactly before/after for
n-tuple and triplet formats, including non-divisible chunk sizes.
Tests
Added a regression test that checks both the query-axis chunking behavior
(
[3, 3, 2]forfaiss_batch_size=3) and output equivalence with a single-batchrun. The same test fails on the previous implementation, which calls
model.similarityonce for all 8 queries.