Skip to content

[fix] Avoid materializing the full similarity matrix in mine_hard_negatives without FAISS#3816

Merged
tomaarsen merged 4 commits into
huggingface:mainfrom
Incheonkirin:fix-mine-hard-negatives-non-faiss-memory
Jun 12, 2026
Merged

[fix] Avoid materializing the full similarity matrix in mine_hard_negatives without FAISS#3816
tomaarsen merged 4 commits into
huggingface:mainfrom
Incheonkirin:fix-mine-hard-negatives-non-faiss-memory

Conversation

@Incheonkirin

Copy link
Copy Markdown
Contributor

Summary

Fixes #3164.

With use_faiss=False, mine_hard_negatives currently computes the full
(n_queries, n_corpus) similarity matrix before running topk:

scores = model.similarity(query_embeddings, corpus_embeddings).to(device)
scores, indices = torch.topk(scores, k=range_max + max_positives, dim=1)

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_size chunks over the query axis.

This PR applies the same query-axis batching to the non-FAISS branch: compute
similarity + topk per chunk, then concatenate the per-chunk results. Since each
row 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_size docstring, added a clear ValueError
for faiss_batch_size <= 0, and made both search progress bars respect verbose.

Local peak RSS delta with all-MiniLM-L6-v2 on CPU and a synthetic corpus:

queries x corpus before after
5k x 50k 1.28 GB 0.55 GB
10k x 100k 4.56 GB 1.09 GB

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] for faiss_batch_size=3) and output equivalence with a single-batch
run. The same test fails on the previous implementation, which calls
model.similarity once for all 8 queries.

pytest tests/util/test_hard_negatives.py -q
ruff check sentence_transformers/util/hard_negatives.py tests/util/test_hard_negatives.py

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 + topk over query chunks controlled by faiss_batch_size.
  • Add faiss_batch_size <= 0 validation and make search progress bars respect verbose.
  • 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.

Comment thread sentence_transformers/util/hard_negatives.py
Comment thread sentence_transformers/util/hard_negatives.py Outdated
Co-authored-by: Copilot Autofix powered by AI <[email protected]>

@tomaarsen tomaarsen left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for this. I've made some modifications re. verbose, which I'll merge.

Comment thread sentence_transformers/util/hard_negatives.py Outdated
Comment thread sentence_transformers/util/hard_negatives.py Outdated
@tomaarsen
tomaarsen merged commit d16e6bf into huggingface:main Jun 12, 2026
17 of 18 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

RAM explosion when using mine hard negatives

3 participants