ADD vector search plugin based on TurboQuant - #8994
Open
pubkey wants to merge 5 commits into
Open
Conversation
Adds the `vector-turboquant` plugin that runs approximate nearest neighbor search over the documents of an RxCollection. It implements the TurboQuant quantizer (https://arxiv.org/abs/2504.19874), so an embedding is stored with 1, 2 or 4 bits per dimension instead of 32 bits as float32: - Normalization: the length of a vector is stored separately. - Rotation: a randomized block Walsh-Hadamard transform turns the coordinates of any dataset into the same normal distribution in d * log(64) time and preserves dot products exactly. - Quantization: a Lloyd-Max codebook for the normal distribution, computed with the Lloyd iteration, so there is no training step. - Length renormalization: one dot product per vector during encoding removes the downward bias of the estimated scores. - Optional per coordinate calibration (TQ+). `collection.addVectorIndex()` fills the index from the collection and keeps it in sync with the change stream. `search()` scores the packed codes with a per byte lookup table and supports `cosine`, `dotProduct` and `euclidean` plus an allowlist to combine a vector search with a normal query. `serializeTurboQuantIndex()` writes the whole index into a binary blob so it does not have to be rebuilt on the next start. Measured on 10k random vectors with 384 dimensions: recall@4 is 100% at 4 bit and the index needs 8 times less memory than float32. Adds testcases for the codebook, the rotation, the index, the serialization and the collection integration, plus the https://rxdb.info/vector-search.html docs page. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_013XoBX9R7x2RZvKHUvavxXS
Contributor
✅ Verify Test Reproduction: Tests FAILED without the fix (expected)This confirms the changed tests correctly reproduce the bug that the source changes fix. This workflow runs the changed tests without the source fix to verify they reproduce the bug. Show output |
The normal CDF of the vector-turboquant plugin uses the rational approximation of Hart (1968). codespell reads the name as a misspelling of 'heart', so it is added to the ignore list. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_013XoBX9R7x2RZvKHUvavxXS
Contributor
|
This looks like a great idea. I wonder if it might benefit from some sort of existing wasm implementation that uses SIMD etc? There's various JS and TS repos in github that implement things like this. A few selected at random: https://github.com/zlaabsi/turboquant-wasm |
Owner
Author
|
Thanks, that is worth weighing. I checked all three: Where the time actually goes in this implementation, measured on Node.js 22 on a single core:
SIMD would make the scan several times faster. But in a local-first app the embedding model dominates: transformers.js needs a few hundred milliseconds per document, so at the corpus sizes that fit on a user device the scan is rarely the bottleneck. The The arguments against bundling WASM into this plugin:
A middle path fits the existing conventions better: keep the dependency-free TypeScript scorer as the default and make the scorer pluggable, the same way the SQLite RxStorage takes the engine from the user through I have not changed anything for this, because taking on a third party dependency for a core plugin is the maintainer's call rather than mine. Generated by Claude Code |
Contributor
|
Owner
Author
|
Decision on this: the plugin stays pure TypeScript. The deciding argument is one I had underweighted, and it is stronger than the portability and dependency points I listed above: crossing the JavaScript to WASM boundary costs latency on every call. A search has to copy the query vector in and the results back out, and RxDB optimizes for latency. At the corpus sizes that fit on a user device, that per-call overhead eats into what SIMD wins back inside the kernel. So no dependency and no change here. Thanks for raising it, the links were worth checking. Generated by Claude Code |
Answers why the scoring loop is plain TypeScript: a WASM kernel with SIMD runs the inner loop faster, but every search pays the cost of copying the query into the WASM memory and the results back out, which competes with what SIMD wins inside the kernel. Asked in #8994 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_013XoBX9R7x2RZvKHUvavxXS
Contributor
|
So no transformers.js then on non-chromium browsers (which don't yet support webgpu)? How will the embedding work without wasm? |
The entry implied that WASM support is uneven on Deno, which is wrong, and it leaned on portability as the main argument. The real argument is the fixed cost of crossing the JavaScript to WASM boundary on every search. transformers.js is the opposite case and does run on WASM by default, because one crossing per document is amortized over a few hundred milliseconds of compute, so the entry now says that instead. Raised in #8994 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_013XoBX9R7x2RZvKHUvavxXS
Owner
Author
|
You are right, and my portability argument was the weak part of that answer. transformers.js runs on WASM by default and only uses WebGPU when you pass What actually decides it is narrower than "WASM is not available", and the two cases are not symmetric:
So the plugin keeps the scorer on the JavaScript side, while the model next to it stays on WASM. Those are consistent positions, not contradictory ones. React Native is the remaining case where WASM itself is a problem, and it is a smaller part of the argument than I made it sound. I have corrected the FAQ entry accordingly in 3204e7c: it now leads with the per-call boundary cost, states that transformers.js runs on WASM by default, and drops the Deno claim. On your earlier point that the three libraries were picked at random and better or more mature ones may exist, including something in Rust or Zig: agreed, and I only evaluated those three. If a mature implementation does turn out to beat the TypeScript scorer end to end, including the marshalling, that is worth revisiting with a benchmark rather than on priors. Generated by Claude Code |
…h-turboquant-plugin-s403e9
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
3 participants
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.
This PR contains:
vector-turboquantplugin)Describe the problem you have without this PR
RxDB has the
vectorplugin with exact distance functions, but no index. To run a similarity search you have to fetch every embedding from the storage and compare it with the query. The vector database article measures that at around 700ms for 10k documents, and it grows linearly. Keeping the rawfloat32embeddings in memory instead is expensive: 100k documents with 768 dimensions need 293 MB.This PR adds a vector index that stores the embeddings in a compressed form, so a search does not read the raw vectors at all.
What it does
The plugin implements TurboQuant, a data-oblivious quantizer from Google Research, in the same spirit as turbovec. An embedding is stored with 1, 2 or 4 bits per dimension instead of 32 bits:
d * log(64)time instead of thed * dof a full random orthogonal matrix and preserves dot products exactly. Blocks of 64 keep the padding at 63 dimensions at most, while padding to the next power of two would waste 33% for the common sizes like 768 or 1536.A search rotates the query into the same space and scores it against the packed codes with a per byte lookup table. The vectors are never decompressed.
API
addVectorIndex()reads the collection once and then follows the change stream, so inserts, updates and deletes are reflected in the next search.search()supportscosine,dotProductandeuclidean, plus anallowlistof ids so a vector search can be combined with a normal Mango query.TurboQuantIndexalso works standalone, andserializeTurboQuantIndex()writes the whole index into a binary blob so it does not have to be rebuilt on the next start.error$emits the errors that a user providedembeddingfunction threw while a change event was processed.Measured results
Recall on 10k random vectors with 384 dimensions and 200 queries.
recall@kis the share of queries where the exact nearest neighbor was inside the firstkresults:Random vectors are the hardest case because in high dimensions all of them have nearly the same distance to each other. The relative error of the estimated scores is about 10% at 4 bit, which matches the distortion of the Lloyd-Max quantizer at that bit width. Scanning 100k vectors of 768 dimensions at 4 bit takes about 100ms on a single core.
The Lloyd-Max solver reproduces the published values from J. Max, "Quantizing for minimum distortion" (1960) to four decimals, and the rotation turns a one-hot vector, the worst case for the quantizer, into coordinates with a variance of 1.00 and a kurtosis of 3.02.
Todos
41 testcases in
test/unit/vector-turboquant.test.tscover the codebook, the rotation, the index, the calibration, the serialization and the collection integration. They pass against both the memory and the dexie storage.npm run lintandnpm run check-typesare clean andnpm run docs:buildsucceeds.New docs page at
docs-src/docs/vector-search.md, linked in the sidebar and from the vector database article, the fulltext search page and the RxPipeline page. The outdated sentence in the vector database article saying RxDB has no dedicated vector plugin was updated.🤖 Generated with Claude Code
https://claude.ai/code/session_013XoBX9R7x2RZvKHUvavxXS
Generated by Claude Code