Added support for adaptive beam search [skip ci]

This commit is contained in:
Andrew Kane
2025-12-05 15:26:49 -08:00
parent b7c0085afe
commit f6edf9245a
2 changed files with 7 additions and 3 deletions

View File

@@ -1282,6 +1282,7 @@ Thanks to:
- [k-means++: The Advantage of Careful Seeding](https://theory.stanford.edu/~sergei/papers/kMeansPP-soda.pdf)
- [Concept Decompositions for Large Sparse Text Data using Clustering](https://www.cs.utexas.edu/users/inderjit/public_papers/concept_mlj.pdf)
- [Efficient and Robust Approximate Nearest Neighbor Search using Hierarchical Navigable Small World Graphs](https://arxiv.org/ftp/arxiv/papers/1603/1603.09320.pdf)
- [Distance Adaptive Beam Search for Provably Accurate Graph-Based Nearest Neighbor Search](https://arxiv.org/pdf/2505.15636)
## History

View File

@@ -829,6 +829,7 @@ HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation in
HnswUnvisited *unvisited = palloc(lm * sizeof(HnswUnvisited));
int unvisitedLength;
bool inMemory = index == NULL;
double lambda = 0.12;
if (v == NULL)
{
@@ -884,7 +885,7 @@ HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation in
HnswSearchCandidate *f = HnswGetSearchCandidate(w_node, pairingheap_first(W));
HnswElement cElement;
if (c->distance > f->distance)
if (c->distance > f->distance * (1 + lambda))
break;
cElement = HnswPtrAccess(base, c->element);
@@ -904,8 +905,10 @@ HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation in
HnswSearchCandidate *e;
double eDistance;
bool alwaysAdd = wlen < ef;
double fDistance;
f = HnswGetSearchCandidate(w_node, pairingheap_first(W));
fDistance = f->distance * (1 + lambda);
if (inMemory)
{
@@ -920,13 +923,13 @@ HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation in
/* Avoid any allocations if not adding */
eElement = NULL;
HnswLoadElementImpl(blkno, offno, &eDistance, q, index, support, inserting, alwaysAdd || discarded != NULL ? NULL : &f->distance, &eElement);
HnswLoadElementImpl(blkno, offno, &eDistance, q, index, support, inserting, alwaysAdd || discarded != NULL ? NULL : &fDistance, &eElement);
if (eElement == NULL)
continue;
}
if (eElement == NULL || !(eDistance < f->distance || alwaysAdd))
if (eElement == NULL || !(eDistance < fDistance || alwaysAdd))
{
if (discarded != NULL)
{