MMR reranking for diverse KNN results via mmr_lambda hidden column#267
Open
MayCXC wants to merge 4 commits intoasg017:mainfrom
Open
MMR reranking for diverse KNN results via mmr_lambda hidden column#267MayCXC wants to merge 4 commits intoasg017:mainfrom
mmr_lambda hidden column#267MayCXC wants to merge 4 commits intoasg017:mainfrom
Conversation
Add Maximal Marginal Relevance (MMR) support to vec0 virtual table. When mmr_lambda is provided in a KNN query, candidates are over-fetched and then greedily re-selected to balance relevance against diversity. API: WHERE embedding MATCH ? AND k = 10 AND mmr_lambda = 0.7 - mmr_lambda range [0.0, 1.0]: 1.0 = pure relevance, 0.0 = pure diversity - Over-fetch factor: 5x (capped at k_max=4096) - Supports float32, int8, and bit vector types - All distance metrics (L2, cosine, L1, hamming) - Zero impact when mmr_lambda is not provided
9 test functions covering: - Cosine diversity (baseline vs lambda=1.0, 0.5, 0.0) - L2 distance metric compatibility - Int8 vector element type - Cluster monopoly breaking - Composition with distance constraints - Composition with partition keys - Edge cases (k=1, k=0) - Error handling (invalid lambda range) - Insert guard for hidden column
The copy-back loop iterated k_target times, but the greedy selection loop can terminate early via `if (best_idx < 0) break`, leaving the tail of out_rowids/out_distances uninitialized. Add an n_selected counter and out_n_selected output parameter so only actually-selected entries are copied back. The caller now sets k_used = n_selected instead of k_used = k_original. Credit: mceachen (vlasky#6)
The relevance term was already normalized to [0,1] via max_dist, but the diversity term used raw distances (1 - d). For cosine this is fine since values are bounded [0,1], but for L2 and L1 the two terms operated on different scales, making mmr_lambda behave unpredictably. Normalize the diversity term the same way (d / max_dist) so both terms are on a consistent [0,1] scale regardless of distance metric. Ported from vlasky/sqlite-vec@8d4ef9e.
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
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.
refs #266
Adds Maximal Marginal Relevance reranking to
vec0KNN queries. Whenmmr_lambdais provided, candidates are over-fetched and then greedily re-selected to balance relevance against diversity.distancecolumn for pagination/custom thresholds #165 / WIP: Support constaints ondistancecolumn in KNN queries, for pagination and range queries #166) and partition keysTests
9 pytest + syrupy snapshot tests covering:
All 158 existing tests pass (+ 8 skipped). Verified in both aarch64 and x86_64 VMs.