From ce91c56175f7119a5e38bb5ea582f572b39685c6 Mon Sep 17 00:00:00 2001 From: Sharma Date: Fri, 3 Jul 2026 13:10:21 +0530 Subject: [PATCH 1/2] Implement auto-tuning for IVF's nprobe with recall estimation and deterministic query selection --- src/index/ivfTune.ts | 59 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 59 insertions(+) create mode 100644 src/index/ivfTune.ts diff --git a/src/index/ivfTune.ts b/src/index/ivfTune.ts new file mode 100644 index 0000000..07f335c --- /dev/null +++ b/src/index/ivfTune.ts @@ -0,0 +1,59 @@ +// Auto-tuning for IVF's recall/latency knob (nprobe). +// +// The insight that makes this cheap: an IVF query misses a true neighbour only +// when the neighbour's cluster isn't probed, and clusters are probed in +// query→centroid score order. So for one sample query, ONE exact scan (ground +// truth) plus one centroid ranking tells us the recall at EVERY possible nprobe +// at once: a ground-truth row is found at nprobe p iff the rank of its cluster +// is < p. Tuning is then a histogram over those ranks — no per-nprobe +// re-querying, no binary search. +// +// This holds exactly for fp32 IVF (probed scores are exact, so any probed +// ground-truth row outranks every non-ground-truth candidate). For IVF×quant +// the same argument applies with quantized scores, so the tuner measures the +// IVF-induced recall loss in the space the index actually scores in — which is +// precisely the loss nprobe controls (the quantization loss is handled +// separately by the exact fp32 re-rank). + +export interface TunedNprobe { + /** Smallest nprobe whose estimated recall meets the target. */ + nprobe: number; + /** Estimated recall@k at that nprobe (fraction of ground-truth hits probed). */ + recall: number; +} + +/** + * Pick the smallest nprobe meeting `targetRecall` from the observed cluster + * ranks of ground-truth neighbours. `neededRanks[i]` is, for one ground-truth + * hit, the probe rank (0 = nearest cluster to the query) of the cluster that + * hit lives in; hits from all sample queries are pooled. + */ +export function chooseNprobe(neededRanks: ArrayLike, nlist: number, targetRecall: number): TunedNprobe { + const total = neededRanks.length; + if (total === 0 || nlist <= 1) return { nprobe: Math.max(1, nlist), recall: 1 }; + const hist = new Uint32Array(nlist); + for (let i = 0; i < total; i++) hist[neededRanks[i]!]!++; + let cum = 0; + for (let p = 1; p <= nlist; p++) { + cum += hist[p - 1]!; + if (cum / total >= targetRecall) return { nprobe: p, recall: cum / total }; + } + return { nprobe: nlist, recall: 1 }; // probing everything is exact by construction +} + +/** + * Deterministic evenly-spaced pick of `want` sample indices out of `filled` + * (the reservoir is already a uniform sample of the corpus, so a stride keeps + * tuning reproducible without another RNG). + */ +export function pickTuneQueries(filled: number, want: number): number[] { + const n = Math.min(filled, want); + const out: number[] = []; + for (let i = 0; i < n; i++) out.push(Math.floor((i * filled) / n)); + return out; +} + +/** Sample queries drawn per tuning pass. */ +export const TUNE_QUERIES = 32; +/** Recall is measured @k=10 (clamped to the corpus size). */ +export const TUNE_K = 10; From 00769fc2ac73270a9925c493120e5243a0c5f23d Mon Sep 17 00:00:00 2001 From: Sharma Date: Fri, 3 Jul 2026 16:34:05 +0530 Subject: [PATCH 2/2] Add IVF auto-tuning feature with target recall support - Implemented auto-tuning for nprobe based on target recall in IVF index. - Updated IVFConfig to include targetRecall parameter for auto-tuning. - Enhanced documentation to reflect changes in IVF auto-tuning and usage examples. - Added tests to validate auto-tuning functionality and its impact on recall. - Created a new example demonstrating the IVF auto-tuning feature. --- CHANGELOG.md | 14 +++++ README.md | 2 +- docs/api-reference.md | 4 +- docs/features.md | 26 +++++++++ examples/20-ivf-autotune.html | 104 ++++++++++++++++++++++++++++++++++ examples/index.html | 1 + src/index.ts | 8 +++ src/index/ivf.ts | 68 +++++++++++++++++++++- src/index/ivfquant.ts | 63 +++++++++++++++++++- src/types.ts | 19 ++++++- tests/browser/gpu.test.ts | 73 ++++++++++++++++++++++++ tests/unit/ivfTune.test.ts | 64 +++++++++++++++++++++ 12 files changed, 436 insertions(+), 10 deletions(-) create mode 100644 examples/20-ivf-autotune.html create mode 100644 tests/unit/ivfTune.test.ts diff --git a/CHANGELOG.md b/CHANGELOG.md index f59212d..0a35031 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,17 @@ the public API may still shift between minor versions). ## [Unreleased] ### Added +- IVF auto-tuning: `ann: { targetRecall: 0.95 }` replaces hand-picking + `nprobe`. After each k-means build the index estimates recall@10 on ≤ 32 + sample queries drawn from the corpus (one exact scan each) and adopts the + smallest `nprobe` whose estimated recall meets the target — one scan per + query prices every candidate `nprobe` at once via the rank of each true + neighbour's cluster (`src/index/ivfTune.ts`). Works on fp32 IVF and the + IVF × quant combo (tuned in the quantized space the index scans; the exact + fp32 re-rank still covers quantization loss). An explicit `nprobe` disables + tuning; per-query `{ nprobe }` overrides still apply. `stats()` now reports + `nprobe` (the effective default) and `tunedRecall` (the tuner's estimate). + New example: `examples/20-ivf-autotune.html`. - Metadata filtering (FR-7): `query()`/`queryText()`/`queryBatch()` accept a Mongo-ish `filter` in `QueryOptions` — bare values for equality plus `$eq`, `$ne`, `$in`, `$gt`/`$gte`/`$lt`/`$lte` (AND across fields; unknown operators @@ -49,6 +60,9 @@ the public API may still shift between minor versions). benchmarks can't drift apart on what "GPU time" or "recall" means. ### Fixed +- A per-query `{ nprobe }` override is now truly consumed by that query alone; + previously it silently stuck as the default for all subsequent queries on + the same IVF store. - `transformersEmbedder()` now sets `env.allowLocalModels = false` before loading a model. transformers.js defaults this to `true` and probes a local `/models/...` path first; since this library never ships local model files, diff --git a/README.md b/README.md index 1f279bb..e5bfae8 100644 --- a/README.md +++ b/README.md @@ -93,7 +93,7 @@ Each links to a short guide with runnable code: | [Persistence](./docs/features.md#persistence-m2) | Versioned binary snapshots to OPFS (or IndexedDB), auto-load on `create()`, export/import as a `Blob`. | | [Encryption at rest](./docs/features.md#encryption-at-rest-m6) | AES-256-GCM + PBKDF2 passphrase envelope for persisted/exported snapshots. | | [Quantization (TurboQuant)](./docs/features.md#quantization--turboquant-int8-m3a) | int8/int4/1-bit codes via randomized Hadamard rotation + exact fp32 re-rank — ~4×/8×/32× less memory. | -| [Approximate search (IVF)](./docs/features.md#approximate-search--ivf-m4) | GPU-assisted k-means clustering; queries scan only the nearest `nprobe` clusters. Combines with quantization for the ~1M-row path. | +| [Approximate search (IVF)](./docs/features.md#approximate-search--ivf-m4) | GPU-assisted k-means clustering; queries scan only the nearest `nprobe` clusters. Combines with quantization for the ~1M-row path. Set `targetRecall: 0.95` instead of `nprobe` and the index auto-tunes itself against recall measured on your own data. | | [Graph search (HNSW)](./docs/features.md#graph-search--hnsw-m7) | Layered proximity graph with O(log N) beam search — incremental inserts (no rebuild), all metrics, works without WebGPU. Graph persists in snapshots, so loads skip the rebuild (~180×). | | [GPU graph search](./docs/features.md#gpu-graph-search-m7b) | CAGRA-style WGSL kernel: the whole beam search in one dispatch, one workgroup per query — `queryBatch()` searches many queries concurrently (~4× the CPU walk at 60k×768×128). | | [Text retrieval / embedder](./docs/features.md#text-retrieval--on-device-embedder-m5) | `addText`/`queryText` via a zero-dep hashing embedder or an optional real semantic model (transformers.js). | diff --git a/docs/api-reference.md b/docs/api-reference.md index 38faf54..2db9a75 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -80,7 +80,7 @@ linked from each entry below. A discriminated union on `type`. Omitting `type` (or `'ivf'`) selects IVF, so pre-M7 configs are unchanged. -**`IVFConfig`** — `type?: 'ivf'`, `nlist` (default ≈ `sqrt(count)`, clamped `[16, 4096]`), `nprobe` (default ≈ 5% of `nlist`), `sampleSize` (default `50_000`, reservoir sample for k-means training), `iters` (default `8` Lloyd iterations), `seed`. +**`IVFConfig`** — `type?: 'ivf'`, `nlist` (default ≈ `sqrt(count)`, clamped `[16, 4096]`), `nprobe` (default: the auto-tuned value when `targetRecall` is set, else ≈ 5% of `nlist`; setting it explicitly disables auto-tuning), `targetRecall` (0–1, e.g. `0.95` — auto-tune `nprobe`: after each build, recall@10 is estimated on ≤ 32 sample queries drawn from the corpus, one exact scan each, and the smallest `nprobe` meeting the target becomes the default; the result surfaces as `stats().nprobe` / `stats().tunedRecall`), `sampleSize` (default `50_000`, reservoir sample for k-means training), `iters` (default `8` Lloyd iterations), `seed`. **`HNSWConfig`** — `type: 'hnsw'`, `M` (graph out-degree per layer, layer 0 keeps 2·M; default `16`), `efConstruction` (build beam width; default `200`), `efSearch` (query beam width, clamped ≥ k; default `64`), `seed` (level RNG, reproducible builds), `search` (`'cpu' | 'gpu'`, default `'cpu'` — `'gpu'` runs the single-dispatch beam-search kernel; needs WebGPU, `M ≤ 32`, `efSearch ≤ 256`, corpus within one storage buffer, and shines on `queryBatch`). @@ -135,7 +135,7 @@ Execution picks a strategy per query: ### `Stats` (from `stats()`) -`count`, `deleted?`, `dimension`, `metric`, `device: 'webgpu' | 'wasm'`, `lastQueryMs?`, `lastQueryGpuMs?` (GPU kernel portion), `lastQueryCpuMs?` (CPU re-rank/filtering portion), `persist?: 'opfs' | 'indexeddb'`, `quantBits`, `nlist?` (IVF cluster count once built), `maxLevel?` (HNSW top graph layer once non-empty), `graphSearch?: 'gpu' | 'cpu'` (which engine answers HNSW queries), `chunks?` (>1 once the corpus spans multiple GPU buffers), `ingest?: 'worker' | 'main-thread'` (where quantized rotate+quantize ran), `train?: 'worker' | 'main-thread'` (where the ANN index build ran — IVF k-means updates, or HNSW graph construction). +`count`, `deleted?`, `dimension`, `metric`, `device: 'webgpu' | 'wasm'`, `lastQueryMs?`, `lastQueryGpuMs?` (GPU kernel portion), `lastQueryCpuMs?` (CPU re-rank/filtering portion), `persist?: 'opfs' | 'indexeddb'`, `quantBits`, `nlist?` (IVF cluster count once built), `nprobe?` (IVF default clusters-per-query — explicit, auto-tuned, or the 5% heuristic — once built), `tunedRecall?` (estimated recall@10 at the auto-tuned `nprobe`; only present when `targetRecall` tuning ran), `maxLevel?` (HNSW top graph layer once non-empty), `graphSearch?: 'gpu' | 'cpu'` (which engine answers HNSW queries), `chunks?` (>1 once the corpus spans multiple GPU buffers), `ingest?: 'worker' | 'main-thread'` (where quantized rotate+quantize ran), `train?: 'worker' | 'main-thread'` (where the ANN index build ran — IVF k-means updates, or HNSW graph construction). ### `SupportInfo` (from `isSupported()`) diff --git a/docs/features.md b/docs/features.md index eef3e16..5d5e73a 100644 --- a/docs/features.md +++ b/docs/features.md @@ -193,6 +193,32 @@ Real embeddings cluster well, so recall@10 ≥ 0.95 is reachable at a small > **Try it:** [`examples/03-ivf-index.html`](../examples/03-ivf-index.html) +**Auto-tuning (`targetRecall`).** Rather than hand-picking `nprobe`, state the +recall you want: + +```ts +const db = await BrowserVec.create({ + dimension: 768, + metric: 'cosine', + ann: { targetRecall: 0.95 }, // auto-tune nprobe to ≈95% recall@10 +}); +``` + +After each k-means build the index estimates recall on ≤ 32 sample queries drawn +from the corpus itself and adopts the smallest `nprobe` whose estimated recall +meets the target. One exact scan per sample query prices *every* candidate +`nprobe` at once: a true neighbour is found at `nprobe = p` exactly when its +cluster ranks among the query's `p` nearest centroids, so recall-vs-nprobe falls +out of a histogram of those ranks — no per-`nprobe` re-querying. Easy corpora +get a faster index, hard ones a more thorough one. The choice surfaces as +`stats().nprobe` and `stats().tunedRecall`; an explicit `nprobe` disables +tuning, and per-query `{ nprobe }` overrides still apply. Works on fp32 IVF and +the IVF × quant combo (there the tuner measures the clustering loss in the +quantized space the index scans; quantization loss itself is still covered by +the exact fp32 re-rank). Tuning cost: ≤ 32 exact GPU scans per rebuild. + +> **Try it:** [`examples/20-ivf-autotune.html`](../examples/20-ivf-autotune.html) + **IVF × int8 (the 1M path).** Add `quantBits: 8` to an `ann` store and the corpus is both clustered *and* int8-quantized — so ~1M×768 fits in a single ~1 GB buffer *and* each query scans only the probed lists: diff --git a/examples/20-ivf-autotune.html b/examples/20-ivf-autotune.html new file mode 100644 index 0000000..c16ca68 --- /dev/null +++ b/examples/20-ivf-autotune.html @@ -0,0 +1,104 @@ + + + + + +BrowserVec — 20. IVF auto-tuning + + + +

← All examples · Home

+

20. IVF auto-tuning (targetRecall)

+

Instead of hand-picking nprobe, pass targetRecall: after each k-means build the index estimates recall@10 on sample queries drawn from your own data (one exact scan each, ≤32 queries) and adopts the smallest nprobe whose estimated recall meets the target. The choice surfaces as stats().nprobe / stats().tunedRecall. GPU-only, like all IVF.

+ +
ready.
+ +

+ + + + +

+ + + + diff --git a/examples/index.html b/examples/index.html index 92bfa95..b8d7f54 100644 --- a/examples/index.html +++ b/examples/index.html @@ -47,6 +47,7 @@

Core concepts

  • M3 TurboQuant quantization int8 / int4 quantized storage with asymmetric query + exact re-rank
  • M4 IVF approximate index cluster-based ANN, nprobe sweep, recall vs speed trade-off
  • M4 IVF × quant combo clustered + quantized — the 1M-vector path with int8 or int4
    • +
  • M4 IVF auto-tuning targetRecall picks nprobe for you — recall estimated on your own data at build time
  • M7 HNSW graph index graph-based ANN, efSearch sweep, incremental inserts, works without WebGPU
  • M7b GPU graph search CAGRA-style beam search in one dispatch — queryBatch, CPU vs GPU crossover
  • React React hooks useVectorStore, useSimilaritySearch, useEmbedding, useRetriever — reusable React patterns
    • diff --git a/src/index.ts b/src/index.ts index e6d494e..ef3da88 100644 --- a/src/index.ts +++ b/src/index.ts @@ -719,6 +719,14 @@ export class BrowserVec { const nlist = (this.index as { nlist: number }).nlist; if (nlist > 0) out.nlist = nlist; } + if ('nprobe' in this.index) { + const nprobe = (this.index as { nprobe: number }).nprobe; + if (nprobe > 0) out.nprobe = nprobe; + } + if ('tunedRecall' in this.index) { + const recall = (this.index as { tunedRecall: number | undefined }).tunedRecall; + if (recall !== undefined) out.tunedRecall = recall; + } if ('maxLevel' in this.index) { const maxLevel = (this.index as { maxLevel: number }).maxLevel; if (maxLevel >= 0) out.maxLevel = maxLevel; diff --git a/src/index/ivf.ts b/src/index/ivf.ts index 6b679c3..bc851e0 100644 --- a/src/index/ivf.ts +++ b/src/index/ivf.ts @@ -22,6 +22,7 @@ import { tracedGpuWait } from '../engine/profile.js'; import { topK, type FlatHit, type VectorIndex } from './flat.js'; import { GpuTopK } from './gpuTopk.js'; import { mulberry32 } from '../quant/prng.js'; +import { chooseNprobe, pickTuneQueries, TUNE_K, TUNE_QUERIES } from './ivfTune.js'; const WORKGROUP_SIZE = 64; @@ -30,6 +31,13 @@ export interface IVFParams { nlist?: number; /** Clusters scanned per query. Default ≈ 5% of nlist (min 1). */ nprobe?: number; + /** + * Auto-tune nprobe to hit this recall@10 (0–1, e.g. 0.95). After each build, + * recall is estimated on sample queries drawn from the corpus and the + * smallest nprobe meeting the target becomes the default. Ignored when an + * explicit `nprobe` is set. + */ + targetRecall?: number; /** Reservoir sample size kept for training. Default 50_000. */ sampleSize?: number; /** Max points actually fed to k-means per iteration (GPU-assigned). Default 16_384. */ @@ -50,6 +58,8 @@ export class IVFIndex implements VectorIndex { private rows = 0; private builtRows = -1; // rows count at last build (-1 = never built) private nprobeOverride: number | undefined; // per-query override (see setNprobe) + private tunedNprobe: number | undefined; // auto-tuned default (see tune()) + private tunedRecallEst: number | undefined; // estimated recall@k at tunedNprobe // Reservoir sample (CPU) for k-means; bounded so ingest stays memory-light. private readonly sampleCap: number; @@ -129,6 +139,18 @@ export class IVFIndex implements VectorIndex { this.nprobeOverride = n; } + /** Effective default nprobe (explicit > auto-tuned > 5% heuristic). 0 until built. */ + get nprobe(): number { + if (this.nlistActual === 0) return 0; + const n = this.params.nprobe ?? this.tunedNprobe ?? Math.max(1, Math.round(this.nlistActual * 0.05)); + return Math.max(1, Math.min(n, this.nlistActual)); + } + + /** Estimated recall@k at the auto-tuned nprobe (undefined unless targetRecall tuning ran). */ + get tunedRecall(): number | undefined { + return this.tunedRecallEst; + } + private makePipeline(code: string, label: string): GPUComputePipeline { const module = this.ctx.device.createShaderModule({ label: `browservec:${label}`, code }); return this.ctx.device.createComputePipeline({ @@ -229,6 +251,43 @@ export class IVFIndex implements VectorIndex { this.listOffset = offset; this.listRows = listRows; this.builtRows = this.rows; + + // Auto-tune nprobe against the recall target (an explicit nprobe wins). + if (this.params.targetRecall !== undefined && this.params.nprobe === undefined) { + await this.tune(clusters); + } + } + + /** + * Estimate recall as a function of nprobe and pick the smallest one meeting + * `targetRecall` (see ivfTune.ts for why one exact scan per sample query + * covers every nprobe at once). Queries are drawn from the reservoir sample — + * real corpus rows — and the trivial self-hit is dropped from the ground + * truth so it doesn't inflate recall. Cost: ≤ TUNE_QUERIES exact scans per + * build, on the GPU path the index already uses. + */ + private async tune(rowCluster: Uint32Array): Promise { + const target = Math.min(1, Math.max(0, this.params.targetRecall!)); + const k = Math.min(TUNE_K, this.rows - 1); + if (k < 1 || this.nlistActual <= 1) { + this.tunedNprobe = 1; + this.tunedRecallEst = 1; + return; + } + const dim = this.dim; + const rank = new Int32Array(this.nlistActual); + const needed: number[] = []; + for (const si of pickTuneQueries(this.sampleFilled, TUNE_QUERIES)) { + const q = this.sample.subarray(si * dim, si * dim + dim); + const gt = await this.search(q, this.nlistActual, k + 1); // probe everything = exact + const order = this.pickProbes(q, this.nlistActual); + for (let r = 0; r < order.length; r++) rank[order[r]!] = r; + // gt[0] is the query row itself (a corpus row) — skip it. + for (let h = 1; h < gt.length; h++) needed.push(rank[rowCluster[gt[h]!.row]!]!); + } + const { nprobe, recall } = chooseNprobe(needed, this.nlistActual, target); + this.tunedNprobe = nprobe; + this.tunedRecallEst = recall; } /** @@ -297,8 +356,11 @@ export class IVFIndex implements VectorIndex { if (this.ctx.lost) throw new Error('GPU device lost; re-create the store'); if (this.rows === 0) return []; if (this.builtRows !== this.rows) await this.build(); + return this.search(queryVec, this.resolveNprobe(), k); + } - const nprobe = this.resolveNprobe(); + /** Probe → gather → GPU scan → top-k, at an explicit nprobe (index must be built). */ + private async search(queryVec: Float32Array, nprobe: number, k: number): Promise { const probes = this.pickProbes(queryVec, nprobe); const candidates = this.gatherCandidates(probes); const n = candidates.length; @@ -320,8 +382,8 @@ export class IVFIndex implements VectorIndex { } private resolveNprobe(): number { - const def = Math.max(1, Math.round(this.nlistActual * 0.05)); - const n = this.nprobeOverride ?? this.params.nprobe ?? def; + const n = this.nprobeOverride ?? this.nprobe; + this.nprobeOverride = undefined; // per-query override, consumed once return Math.max(1, Math.min(n, this.nlistActual)); } diff --git a/src/index/ivfquant.ts b/src/index/ivfquant.ts index 940bc96..63594d4 100644 --- a/src/index/ivfquant.ts +++ b/src/index/ivfquant.ts @@ -26,6 +26,7 @@ import { tracedGpuWait } from '../engine/profile.js'; import { topK, type FlatHit, type VectorIndex } from './flat.js'; import { GpuTopK } from './gpuTopk.js'; import type { IVFParams } from './ivf.js'; +import { chooseNprobe, pickTuneQueries, TUNE_K, TUNE_QUERIES } from './ivfTune.js'; const WORKGROUP_SIZE = 64; @@ -63,6 +64,8 @@ export class IVFQuantIndex implements VectorIndex { private listOffset: Int32Array | null = null; private builtRows = -1; private nprobeOverride: number | undefined; + private tunedNprobe: number | undefined; // auto-tuned default (see tune()) + private tunedRecallEst: number | undefined; // estimated recall@k at tunedNprobe // Per-query scratch. private candidatesBuf: GPUBuffer | null = null; @@ -152,6 +155,18 @@ export class IVFQuantIndex implements VectorIndex { this.nprobeOverride = n; } + /** Effective default nprobe (explicit > auto-tuned > 5% heuristic). 0 until built. */ + get nprobe(): number { + if (this.nlistActual === 0) return 0; + const n = this.params.nprobe ?? this.tunedNprobe ?? Math.max(1, Math.round(this.nlistActual * 0.05)); + return Math.max(1, Math.min(n, this.nlistActual)); + } + + /** Estimated recall@k at the auto-tuned nprobe (undefined unless targetRecall tuning ran). */ + get tunedRecall(): number | undefined { + return this.tunedRecallEst; + } + private makePipeline(code: string, label: string): GPUComputePipeline { const module = this.ctx.device.createShaderModule({ label: `browservec:${label}`, code }); return this.ctx.device.createComputePipeline({ @@ -289,6 +304,43 @@ export class IVFQuantIndex implements VectorIndex { this.listOffset = offset; this.listRows = listRows; this.builtRows = this.rows; + + // Auto-tune nprobe against the recall target (an explicit nprobe wins). + if (this.params.targetRecall !== undefined && this.params.nprobe === undefined) { + await this.tune(clusters); + } + } + + /** + * Estimate recall as a function of nprobe and pick the smallest one meeting + * `targetRecall` (see ivfTune.ts). Sample queries come from the reservoir — + * already rotated, so they feed the rotated search path directly. Ground + * truth is the full quantized scan (probe everything), which isolates the + * IVF-induced loss nprobe controls; quantization loss is handled by the + * caller's exact fp32 re-rank as usual. + */ + private async tune(rowCluster: Uint32Array): Promise { + const target = Math.min(1, Math.max(0, this.params.targetRecall!)); + const k = Math.min(TUNE_K, this.rows - 1); + if (k < 1 || this.nlistActual <= 1) { + this.tunedNprobe = 1; + this.tunedRecallEst = 1; + return; + } + const pd = this.paddedDim; + const rank = new Int32Array(this.nlistActual); + const needed: number[] = []; + for (const si of pickTuneQueries(this.sampleFilled, TUNE_QUERIES)) { + const rq = this.sample.subarray(si * pd, si * pd + pd); + const gt = await this.search(rq, this.nlistActual, k + 1); // probe everything = exact + const order = this.pickProbes(rq, this.nlistActual); + for (let r = 0; r < order.length; r++) rank[order[r]!] = r; + // gt[0] is (almost always) the query row itself — skip it either way. + for (let h = 1; h < gt.length; h++) needed.push(rank[rowCluster[gt[h]!.row]!]!); + } + const { nprobe, recall } = chooseNprobe(needed, this.nlistActual, target); + this.tunedNprobe = nprobe; + this.tunedRecallEst = recall; } /** fp32 assignment over a rotated training buffer. */ @@ -410,7 +462,12 @@ export class IVFQuantIndex implements VectorIndex { if (this.builtRows !== this.rows) await this.build(); const rq = this.encoder.rotateQuery(queryVec); // rotated fp32 query (asymmetric) - const probes = this.pickProbes(rq, this.resolveNprobe()); + return this.search(rq, this.resolveNprobe(), k); + } + + /** Probe → gather → GPU scan → top-k over a *rotated* query, at an explicit nprobe. */ + private async search(rq: Float32Array, nprobe: number, k: number): Promise { + const probes = this.pickProbes(rq, nprobe); const candidates = this.gatherCandidates(probes); const total = candidates.length; if (total === 0) return []; @@ -430,8 +487,8 @@ export class IVFQuantIndex implements VectorIndex { } private resolveNprobe(): number { - const def = Math.max(1, Math.round(this.nlistActual * 0.05)); - const n = this.nprobeOverride ?? this.params.nprobe ?? def; + const n = this.nprobeOverride ?? this.nprobe; + this.nprobeOverride = undefined; // per-query override, consumed once return Math.max(1, Math.min(n, this.nlistActual)); } diff --git a/src/types.ts b/src/types.ts index 847371d..a551e74 100644 --- a/src/types.ts +++ b/src/types.ts @@ -67,8 +67,21 @@ export interface IVFConfig { type?: 'ivf'; /** Number of clusters. Default ≈ sqrt(count), clamped to [16, 4096]. */ nlist?: number; - /** Clusters scanned per query. Default ≈ 5% of nlist. Higher = better recall, slower. */ + /** + * Clusters scanned per query. Higher = better recall, slower. Default: the + * auto-tuned value when `targetRecall` is set, else ≈ 5% of nlist. Setting + * this explicitly disables auto-tuning. + */ nprobe?: number; + /** + * Auto-tune nprobe instead of picking it by hand: after each index build the + * store estimates recall@10 on sample queries drawn from your own data (one + * exact scan per query, ≤ 32 queries) and adopts the smallest nprobe whose + * estimated recall meets this target (0–1, e.g. 0.95). The chosen value and + * its estimated recall surface as stats().nprobe / stats().tunedRecall. + * Ignored when `nprobe` is set explicitly. + */ + targetRecall?: number; /** Reservoir sample size used to train k-means. Default 50_000. */ sampleSize?: number; /** Lloyd iterations during build. Default 12. */ @@ -261,6 +274,10 @@ export interface Stats { quantBits: 0 | 1 | 4 | 8; /** IVF cluster count, if an IVF index is in use and built. */ nlist?: number; + /** IVF default clusters-per-query (explicit, auto-tuned, or the 5% heuristic), once built. */ + nprobe?: number; + /** Estimated recall@10 at the auto-tuned nprobe (only when `targetRecall` tuning ran). */ + tunedRecall?: number; /** HNSW top graph layer, if an HNSW index is in use and non-empty. */ maxLevel?: number; /** HNSW query engine: 'gpu' (beam-search kernel) or 'cpu' (graph walk). */ diff --git a/tests/browser/gpu.test.ts b/tests/browser/gpu.test.ts index d4920bc..aa378aa 100644 --- a/tests/browser/gpu.test.ts +++ b/tests/browser/gpu.test.ts @@ -233,6 +233,79 @@ describe.skipIf(!hasGpu)('IVF index', () => { // Scanning every cluster must be exact; nprobe=1 strictly can't beat it. expect(high / queries.length).toBeGreaterThanOrEqual(0.99); expect(low / queries.length).toBeLessThanOrEqual(high / queries.length); + // The override is per-query: the configured default (1) must survive it. + expect(db.stats().nprobe).toBe(1); + } finally { + db.destroy(); + } + }); + + it('targetRecall auto-tunes nprobe and delivers the target on real queries', async () => { + const dim = 32; + const target = 0.95; + const vectors = randomVectors(3000, dim, 44); + const db = await BrowserVec.create({ + dimension: dim, + ann: { type: 'ivf', nlist: 48, targetRecall: target, seed: 3 }, + }); + try { + await db.addBatch(records(vectors)); + let total = 0; + const queries = randomVectors(20, dim, 45); + for (const q of queries) { + total += recall(await db.query(q, { k: 10 }), bruteForce(vectors, q, 10, 'cosine', true)); + } + const stats = db.stats(); + // Tuning ran: a concrete nprobe was chosen and its estimate meets the target. + expect(stats.nprobe).toBeGreaterThanOrEqual(1); + expect(stats.nprobe).toBeLessThanOrEqual(48); + expect(stats.tunedRecall).toBeGreaterThanOrEqual(target); + // Estimated on sampled corpus queries; allow slack on independent ones. + expect(total / queries.length).toBeGreaterThanOrEqual(target - 0.1); + } finally { + db.destroy(); + } + }); + + it('auto-tuning also drives the IVF×quant combo', async () => { + const dim = 32; + const target = 0.9; + const vectors = randomVectors(3000, dim, 46); + const db = await BrowserVec.create({ + dimension: dim, + quantBits: 8, + ann: { type: 'ivf', nlist: 48, targetRecall: target, seed: 4 }, + }); + try { + await db.addBatch(records(vectors)); + let total = 0; + const queries = randomVectors(20, dim, 47); + for (const q of queries) { + total += recall(await db.query(q, { k: 10 }), bruteForce(vectors, q, 10, 'cosine', true)); + } + const stats = db.stats(); + expect(stats.nprobe).toBeGreaterThanOrEqual(1); + expect(stats.tunedRecall).toBeGreaterThanOrEqual(target); + // int8 adds its own (re-ranked) loss on top of the tuned IVF loss. + expect(total / queries.length).toBeGreaterThanOrEqual(target - 0.15); + } finally { + db.destroy(); + } + }); + + it('an explicit nprobe disables auto-tuning', async () => { + const dim = 16; + const vectors = randomVectors(500, dim, 48); + const db = await BrowserVec.create({ + dimension: dim, + ann: { type: 'ivf', nlist: 16, nprobe: 4, targetRecall: 0.99, seed: 5 }, + }); + try { + await db.addBatch(records(vectors)); + await db.query(randomVectors(1, dim, 49)[0]!, { k: 5 }); + const stats = db.stats(); + expect(stats.nprobe).toBe(4); + expect(stats.tunedRecall).toBeUndefined(); } finally { db.destroy(); } diff --git a/tests/unit/ivfTune.test.ts b/tests/unit/ivfTune.test.ts new file mode 100644 index 0000000..637f8da --- /dev/null +++ b/tests/unit/ivfTune.test.ts @@ -0,0 +1,64 @@ +// The nprobe auto-tuner's pure core: given the probe rank of each ground-truth +// hit's cluster, pick the smallest nprobe whose cumulative recall meets the +// target (see src/index/ivfTune.ts for why one scan covers every nprobe). + +import { describe, expect, it } from 'vitest'; +import { chooseNprobe, pickTuneQueries } from '../../src/index/ivfTune.js'; + +describe('chooseNprobe', () => { + it('picks the smallest nprobe whose cumulative recall meets the target', () => { + // 10 hits: 6 in the nearest cluster (rank 0), 3 at rank 1, 1 at rank 7. + const ranks = [0, 0, 0, 0, 0, 0, 1, 1, 1, 7]; + expect(chooseNprobe(ranks, 16, 0.5)).toEqual({ nprobe: 1, recall: 0.6 }); + expect(chooseNprobe(ranks, 16, 0.9)).toEqual({ nprobe: 2, recall: 0.9 }); + expect(chooseNprobe(ranks, 16, 0.95)).toEqual({ nprobe: 8, recall: 1 }); + }); + + it('recall is monotone in nprobe: a higher target never yields a smaller nprobe', () => { + const ranks = [0, 2, 2, 3, 5, 5, 5, 9]; + let prev = 0; + for (const target of [0.1, 0.3, 0.5, 0.7, 0.9, 1]) { + const { nprobe, recall } = chooseNprobe(ranks, 10, target); + expect(recall).toBeGreaterThanOrEqual(target); + expect(nprobe).toBeGreaterThanOrEqual(prev); + prev = nprobe; + } + }); + + it('target 1 requires probing up to the worst hit, never past nlist', () => { + expect(chooseNprobe([0, 4], 5, 1)).toEqual({ nprobe: 5, recall: 1 }); + expect(chooseNprobe([0, 0], 5, 1)).toEqual({ nprobe: 1, recall: 1 }); + }); + + it('degenerate inputs fall back safely', () => { + expect(chooseNprobe([], 16, 0.95)).toEqual({ nprobe: 16, recall: 1 }); + expect(chooseNprobe([0], 1, 0.95)).toEqual({ nprobe: 1, recall: 1 }); + expect(chooseNprobe([], 0, 0.95)).toEqual({ nprobe: 1, recall: 1 }); + }); + + it('accepts typed arrays', () => { + expect(chooseNprobe(Uint32Array.from([0, 1, 1, 3]), 8, 0.75)).toEqual({ nprobe: 2, recall: 0.75 }); + }); +}); + +describe('pickTuneQueries', () => { + it('returns evenly spaced distinct indices within range', () => { + const picks = pickTuneQueries(1000, 32); + expect(picks).toHaveLength(32); + expect(new Set(picks).size).toBe(32); + for (const p of picks) { + expect(p).toBeGreaterThanOrEqual(0); + expect(p).toBeLessThan(1000); + } + expect(picks).toEqual([...picks].sort((a, b) => a - b)); + }); + + it('clamps to the filled count when the reservoir is small', () => { + expect(pickTuneQueries(5, 32)).toEqual([0, 1, 2, 3, 4]); + expect(pickTuneQueries(0, 32)).toEqual([]); + }); + + it('is deterministic', () => { + expect(pickTuneQueries(777, 32)).toEqual(pickTuneQueries(777, 32)); + }); +});