Skip to content

Use ScoreTracker to avoid wasteful searching for very large k #384

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
marianotepper merged 2 commits into from
Jan 9, 2025
Merged

Use ScoreTracker to avoid wasteful searching for very large k #384

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
978a420
clarify
Jan 9, 2025
829aa92
use scoreTracker to short circuit new edge evaluation once we hit a l…
Jan 9, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 8 additions & 12 deletions jvector-base/src/main/java/io/github/jbellis/jvector/graph/GraphSearcher.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -264,11 +264,8 @@ private SearchResult resume(int initialVisited, int topK, int rerankK, float thr
rerankedResults.setMaxSize(topK);

int numVisited = initialVisited;
// A bound that holds the minimum similarity to the query vector that a candidate vector must
// have to be considered -- will be set to the lowest score in the results queue once the queue is full.
var minAcceptedSimilarity = Float.NEGATIVE_INFINITY;
// track scores to predict when we are done with threshold queries
var scoreTracker = threshold > 0 ? new ScoreTracker.TwoPhaseTracker(threshold) : ScoreTracker.NO_OP;
var scoreTracker = threshold > 0 ? new ScoreTracker.TwoPhaseTracker(threshold) : new ScoreTracker.TwoPhaseTracker(1.0);
VectorFloat<?> similarities = null;

// add evicted results from the last call back to the candidates
Expand All @@ -283,37 +280,36 @@ private SearchResult resume(int initialVisited, int topK, int rerankK, float thr
while (candidates.size() > 0) {
// we're done when we have K results and the best candidate is worse than the worst result so far
float topCandidateScore = candidates.topScore();
if (topCandidateScore < minAcceptedSimilarity) {
if (approximateResults.size() >= rerankK && topCandidateScore < approximateResults.topScore()) {
break;
}
// when querying by threshold, also stop when we are probabilistically unlikely to find more qualifying results
if (scoreTracker.shouldStop()) {
if (threshold > 0 && scoreTracker.shouldStop()) {
break;
}

// process the top candidate
int topCandidateNode = candidates.pop();
if (acceptOrds.get(topCandidateNode) && topCandidateScore >= threshold) {
addTopCandidate(topCandidateNode, topCandidateScore, rerankK);

// update minAcceptedSimilarity if we've found K results
if (approximateResults.size() >= rerankK) {
minAcceptedSimilarity = approximateResults.topScore();
}
}

// if this candidate came from evictedResults, we don't need to evaluate its neighbors again
if (previouslyEvicted.get(topCandidateNode)) {
continue;
}

// skip edge loading if we've found a local maximum and we have enough results
if (scoreTracker.shouldStop() && candidates.size() >= rerankK - approximateResults.size()) {
continue;
}

// score the neighbors of the top candidate and add them to the queue
var scoreFunction = scoreProvider.scoreFunction();
var useEdgeLoading = scoreFunction.supportsEdgeLoadingSimilarity();
if (useEdgeLoading) {
similarities = scoreFunction.edgeLoadingSimilarityTo(topCandidateNode);
}

var it = view.getNeighborsIterator(topCandidateNode);
for (int i = 0; i < it.size(); i++) {
var friendOrd = it.nextInt();
Expand Down
8 changes: 5 additions & 3 deletions jvector-base/src/main/java/io/github/jbellis/jvector/graph/ScoreTracker.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -55,7 +55,7 @@ class TwoPhaseTracker implements ScoreTracker {
private int recentEntryIndex;

// Heap of the best scores seen so far
AbstractLongHeap bestScores;
BoundedLongHeap bestScores;

// observation count
private int observationCount;
Expand Down Expand Up @@ -87,8 +87,10 @@ public boolean shouldStop() {
return false;
}

// we're in phase 2 if the 99th percentile of the recent scores is worse than the best score
// (paper suggests median, but experimentally that is too prone to false positives.
// We're in phase 2 if the 99th percentile of the recent scores evaluated is lower
// than the worst of the best scores seen.
//
// (paper suggests using the median of recent scores, but experimentally that is too prone to false positives.
// 90th does seem to be enough, but 99th doesn't result in much extra work, so we'll be conservative)
double windowMedian = StatUtils.percentile(recentScores, 99);
double worstBest = sortableIntToFloat((int) bestScores.top());
Expand Down