(diversity) attribute diversity implementation with post process and adaptive L#1250
(diversity) attribute diversity implementation with post process and adaptive L#1250narendatha wants to merge 15 commits into
Conversation
Remove the experimental_diversity_search feature gate (graduate the diverse-search primitives to always-on) and add a post-processing-only attribute-bucket diversity path: a plain greedy search over the top-L pool followed by AttributeBucketDiversity, which keeps at most diverse_results_k results per distinct attribute value. Exposed as SearchMode::DiverseAttribute and wired into the disk benchmark via new attributes/diverse_attribute_id/diverse_results_k inputs.
Add DiverseAdaptiveSearch, a greedy graph search that samples attribute-bucket concentration during traversal and grows the search list L when few distinct buckets are seen (yield = sum min(count_b,k)/visited). The enlarged pool feeds the existing AttributeBucketDiversity post-processor. Wired through SearchMode::DiverseAttribute's optional adaptive_l and the disk benchmark. Reuses inline-filter compute_adaptive_l. Includes unit tests for the bucket-yield sizing.
AttributeBucketDiversity::post_process previously recomputed exact distances for every candidate in the L-pool, fetching each full-precision vector and re-evaluating the distance. This duplicated work already done during traversal and made Design A slower than the queue method at higher L. Mirror RerankAndFilter: consult scratch.distance_cache first and only fetch/recompute for cache misses, then sort and run the greedy per-bucket selection. Recall and IO counts are unchanged; redundant vector fetches and distance recomputation are eliminated. Add diverse-search-results.md documenting the Enron/Caselaw 3-way comparison (standard, queue, Design A, Design B) and this optimization.
| @@ -0,0 +1,225 @@ | |||
| # Attribute-Diversity Search — Benchmark Results | |||
There was a problem hiding this comment.
This doc and the attached plots should be moved to the Wiki.
There was a problem hiding this comment.
Sure. will do this once we integrate attribute getting code done.
| //! the search list `L` accordingly. The actual bucket-diverse top-`k` selection | ||
| //! is still performed by a downstream post-processor over the enlarged pool. | ||
| //! | ||
| //! # Motivation |
There was a problem hiding this comment.
Maybe I'm misunderstanding something about this calculation, but it seems like the space of possible values might be constrained in a bad way when the number of attributes is small. If diverse_results_k=1, and the number of buckets is low (say, 60-70 like in some of your experiments), it seems like if the number of points visited is set to the default 1000, the yield could only ever be as high as .06 even with a maximally diverse set, and L might be increased even if the yield is already perfect. It seems like there needs to be some coupling of diverse_results_k with num_buckets and the number of samples in AdaptiveL to make sure that yield is actually allowed to be anywhere from 0 to 1.0.
There was a problem hiding this comment.
Thinking about this more, maybe what we should be doing is waiting for a certain number of samples to accumulate, and then computing yield = (Σ_b min(count_b, k)) / num_buckets * k. Then, we would probably want to scale the adaptive L a bit differently from how filtered search does it. For filtered search, a specificity of .5 over 1000 samples is really good and an indication that we probably don't need to search more. On the other hand, with my proposed definition, a yield of .5 would be pretty bad and would indicate that we have only reached <= 50% recall so far and probably need to search quite a lot more.
There was a problem hiding this comment.
Thanks — both points were spot on. Fixed in the latest push:
1. Yield is now coupled to num_buckets, k, and the sample size. The metric is:
yield = (Σ_b min(count_b, k)) / min(sample_visited, num_buckets·k)
The denominator is the tight upper bound on the numerator, so yield always spans the full [0, 1]:
- few buckets / large sample → denominator is
num_buckets·k, so covering every bucket reaches 1.0 (your 60–70 bucket case); - many buckets / small sample → denominator is
sample_visited, so a fully distinct sample reaches 1.0 and diffuse data leavesLunchanged.
num_buckets comes from a new AttributeValueProvider::num_buckets() (defaults to None, which disables growth); the file-backed provider reports its distinct-value count.
2. Diversity yield is scaled differently from filter specificity. Design B no longer reuses the inline-filter piecewise curve. L now scales linearly from the deficit 1 − yield, so a yield of 0.5 grows L substantially (≈ halfway to the max multiplier) rather than treating 0.5 as "good enough" — matching your point that 0.5 yield ≈ ≤50% recall and needs much more search.
There was a problem hiding this comment.
Followed up by rerunning Design B across all four datasets after the yield-normalization change, and regenerated the charts. Recall is deterministic (authoritative); QPS is kept as best-of-several single-threaded runs since it varies 10–30× from disk contention.
Design B recall @ L = 20 / 40 / 80 / 100:
| Dataset (attribute) | Regime | Before | After |
|---|---|---|---|
| Enron | concentrated, uncorrelated | 73.4 / 83.6 / 89.8 / 91.0 | 79.9 / 87.4 / 92.1 / 93.0 |
Caselaw doc_id |
diffuse | 92.0 / 97.5 / 99.1 / 99.3 | 92.0 / 97.5 / 99.1 / 99.3 (unchanged) |
Caselaw court_jurisdiction |
concentrated + distance-correlated | 59.2 / 69.6 / 77.6 / 79.9 | 57.3 / 68.2 / 76.6 / 79.0 |
YFCC camera |
concentrated, ~uncorrelated | 83.6 / 92.5 / 95.4 / 96.3 | 88.4 / 94.2 / 96.5 / 96.9 |
Takeaways matching the intent of the review:
- Concentrated + uncorrelated (Enron, YFCC): clear recall gains — normalizing the yield by achievable diversity drives more appropriate over-fetch. +6.5 pts on Enron and +4.9 pts on YFCC at L=20.
- Diffuse (Caselaw
doc_id): unchanged — yield ≈ 1.0, so no over-fetch; still matches Design A on recall and IO exactly. Correct behavior preserved. - Concentrated + distance-correlated (Caselaw
court_jurisdiction): a small, honest regression (~1–2 pts). With only 60 buckets the denominatornum_buckets·kcaps at 60, so seeing many courts in the sample reads as good coverage and over-fetches less (IO 193 vs old 210). Design B now sits just below the queue at every L (79.0 vs 79.1 at L=100) rather than edging past it — which actually reinforces the doc's conclusion that in-traversal diversity (the queue) is the right tool for distance-correlated attributes, and post-processing shouldn't be expected to win there.
I also checked in the chart tooling so this is reproducible going forward: assets/diverse-search/chart_data.json (the plotted recall/QPS values) + assets/diverse-search/plot_qps_vs_recall.py. All four PNGs and the results tables/prose are updated.
I will remove these chart tooling scripts once we are close to shipping level with right way to get attributes becomes available. We will also need to discuss how to get the num_buckets attribute information.
There was a problem hiding this comment.
Follow-up: num_buckets-unavailable fallback
Rather than disabling adaptive growth when a provider can't report num_buckets, diverse_yield now falls back to normalizing by sample_visited alone — the original pre-normalization heuristic. So on that path L still grows on concentrated samples (never worse than fixed-L Design A) instead of collapsing to Design A.
Verified by temporarily forcing num_buckets() -> None and re-running Design B:
| Dataset (buckets) | With num_buckets (L=20/40/80) |
None fallback (L=20/40/80) |
|---|---|---|
| Enron (~12,849) | 79.9 / 87.4 / 92.1 | 79.9 / 87.4 / 92.1 (identical) |
| YFCC (67) | 88.4 / 94.2 / 96.5 | 91.5 / 95.4 / 96.9 (over-fetches more) |
- Many buckets (Enron): fallback is a no-op —
num_buckets·kalready exceeds the sample, so both paths cap the denominator atsample_visited. - Few buckets (YFCC): fallback grows
Lmore than necessary (the sample looks concentrated without the bucket total), yielding higher recall at higher IO. Still bounded below by Design A.
Recall stays deterministic; QPS omitted as usual due to run-to-run disk-contention noise. Added a yield_unknown_buckets_uses_sample_size unit test covering the None path.
There was a problem hiding this comment.
Thanks, it's good to see that recall mostly improved after better normalization. I was wondering what sample_visited is set to in the experiments? In particular with L=20, if it's set to the default 1000, are we always seeing 1000 samples and giving the algorithm a chance to adapt?
It's interesting that normalizing didn't help for Caselaw with the court_jurisdiction attribute. I guess one interpretation of what happened there is that since Caselaw has much more correlation between distance and attribute, the buckets were saturated earlier on in the search by points that were far away from the query, but all of the points that are close to the query were from a smaller set of attributes. Then it would look to the Adaptive L calculation as if diversity had been achieved, but actually the attribute-satisfying points closest to the query hadn't been found. Previously with num_samples normalization, L would have just been coincidentally increased because the cap on the ratio was artificially small. Normalizing correctly is still the right thing to do, but we could consider playing with increasing k a bit for a small number of buckets?
There was a problem hiding this comment.
Great questions.
On sample_count: it's a required config field (no default), and in every experiment here it's set to 100, not 1000 — so the sampler estimates yield after the first 100 visited nodes, then resizes L and continues. Even at L=20 the search visits far more than 100 nodes before it would naturally converge (fixed-L=20 Design A already does ~1,300 comparisons / ~40 IOs per query on Enron), so the 100-sample threshold is always reached and adaptation always fires. You can see it in the L=20 Enron numbers: Design B does ~1,880 comps / ~81 IOs vs Design A's ~1,300 / ~40 — i.e. L really was grown mid-search, and recall went 60.0 → 79.9 as a result.
On the court_jurisdiction case: your interpretation matches the data exactly. With 60 buckets and sample_count=100, the correct denominator is min(100, 60·k) = 60, so once the sample has touched all 60 courts (which happens early, dominated by far-away points) the yield saturates at ~1.0 and L stops growing. The old sample_visited normalization divided by 100 while matched could never exceed 60 — so the ratio was capped at 0.6, guaranteeing a ≥40% "deficit" and thus coincidental over-fetch every time. That's precisely why the old numbers were a bit higher here (57.3 vs 59.2 at L=20, with IO 193 vs 210): the previous version was over-fetching for the wrong reason. Correct normalization is still the right call.
On bumping k for few-bucket datasets: I think this is a promising direction. The clean version is to decouple the yield-k from the output diverse_results_k — i.e. require the sample to see, say, 3–4 points per bucket before the yield saturates, even though the final result still caps at 1 per bucket. On a correlated attribute that would keep L growing past the point where the far-away points have merely "covered" the buckets, giving the search a chance to reach the closer attribute-satisfying points. It's a small change (an extra yield_k knob feeding diverse_matched_slots/diverse_yield), and it degrades gracefully on diffuse data where each bucket only has one point anyway. Does it match with your thinking? Will prototype this now.
There was a problem hiding this comment.
Follow-up on the yield_k idea → moved to a separate experiment draft PR #1262
Following the discussion about the jurisdiction regression on distance-correlated attributes, I prototyped the decoupled yield_k knob (sampler per-bucket cap independent of the output diverse_results_k). It works — on jurisdiction, yield_k ≥ 2 recovers the regression and beats the plain-queue baseline, and yield_k = 4 reaches 81.55 recall@10 @ L=100 (vs. 78.97 today) for ~28% more IOs.
However, it adds a knob whose optimal value depends on the (unobservable) attribute/distance correlation, and it interacts with the existing scale_factor / sample_count over-fetch controls. Rather than expand this PR's surface area with a parameter that's hard to specify a default for, I've split it onto its own branch/draft PR so we can:
- sweep
yield_kacross more datasets (Enron, YFCC, Caselaw doc_id) to map where it helps vs. is a no-op, and - explore auto-deriving
yield_kfrom observed bucket saturation so it needn't be a knob at all.
This PR is unaffected — it stays as the clean Design B (yield_k = 1 behavior, no new config surface). Experiment branch: u/narendatha/diverse-yield-k-experiment (draft PR here #1262)
What do you feel about adding another knob that may be hard to tune? Perhaps it depends on scenarios. Let me continue experimenting in that other branch with other datasets meanwhile.
Couple the diversity yield to num_buckets, k, and sample size via the denominator min(sample_visited, num_buckets*k), so yield spans [0, 1] in both few-bucket and diffuse regimes. Add AttributeValueProvider::num_buckets (default None disables growth) and report it from FileAttributeProvider. Scale L linearly from the yield deficit (1 - yield) rather than reusing the inline-filter specificity curve.
Rerun Design B across all four datasets after the yield-normalization fix. Recall improves on the concentrated, uncorrelated datasets (Enron, YFCC) and is unchanged on the diffuse dataset (Caselaw doc_id). On the correlated Caselaw jurisdiction attribute Design B now sits just below the queue at every L. Add reusable chart_data.json + plot_qps_vs_recall.py and regenerate the four PNGs.
…lable When AttributeValueProvider::num_buckets returns None, diverse_yield now normalizes by sample_visited alone instead of skipping adaptive growth. This restores the pre-normalization heuristic for that path: L still grows on concentrated samples (never worse than fixed-L Design A) rather than collapsing to Design A.
| @@ -37,7 +37,7 @@ use serde::{Deserialize, Serialize}; | |||
| use crate::{ | |||
There was a problem hiding this comment.
When this goes into benchmark, can you add generated diverse groundtruth for the two filter datasets in test_data and enable an example benchmark json with them?
There was a problem hiding this comment.
Yes, of course, that is the plan. Will keep this comment open until that happens.
…postprocess-adaptive # Conflicts: # diskann-disk/Cargo.toml # diskann-disk/src/search/provider/disk_provider.rs
…bounded queue Accumulate every visited node with its distance during the adaptive-L traversal and hand the full distance-sorted set to the diversity post-processor, instead of the L-bounded scratch.best navigation queue. This mirrors inline_filter_search and lets diverse-but-distant nodes that were evicted from the frontier by closer non-diverse nodes remain eligible as final answers. Recall@10 improves at every L with identical IO across all four benchmark regimes; the biggest gain is the distance-correlated jurisdiction attribute (L100 78.97 -> 90.19). Also gitignore diskann-benchmark/output/.
Re-ran all four Design B regimes after decoupling the post-processor's candidate pool from the L-bounded frontier. Recall improves at every L with identical IO. Notably the distance-correlated jurisdiction attribute now favours Design B over the in-traversal queue (90.2 vs 79.1 at L=100), inverting the previous conclusion. Updated tables, per-dataset analysis, summary, chart data and regenerated the QPS-vs-recall PNGs (QPS unchanged since IO is identical).
…ature list This branch un-gated diversity search (removed the experimental_diversity_search feature entirely), but the CI and nightly workflows still listed it in DISKANN_FEATURES, causing 'none of the selected packages contains this feature'.
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #1250 +/- ##
==========================================
- Coverage 90.16% 90.01% -0.16%
==========================================
Files 509 512 +3
Lines 97238 98539 +1301
==========================================
+ Hits 87678 88696 +1018
- Misses 9560 9843 +283
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
| /// | ||
| /// Clamped to [1×, max_multiplier] range. | ||
| fn compute_adaptive_l(base_l: usize, visited: usize, matched: usize, max_multiplier: f64) -> usize { | ||
| pub(crate) fn compute_adaptive_l( |
There was a problem hiding this comment.
I think this doesn't need to be pub(crate) anymore
| inner: Knn, | ||
| provider: Arc<P>, | ||
| diverse_results_k: usize, | ||
| adaptive_l: AdaptiveL, |
There was a problem hiding this comment.
It might be worth making AdaptiveL and Option the way it is in inline filtered search. This can be helpful for benchmarking how well regular search can perform without adaptations, scenarios where the user really wants to be able to control how much compute is being used, and perhaps also might be the right thing to do when num_buckets isn't known.
|
Thanks Naren! It was good to see that normalization and using the entire visited set made a big impact. I don't have any more concrete thoughts at the moment on how the algorithm could be improved, although I think it would be interesting to look at estimating distance<->attribute correlation at some point. I can return for another pass once we're ready to move this out of draft. |
Attribute-bucket diversity search (post-processing + adaptive-L)
Summary
Adds two independent, opt-in strategies for attribute-bucket diverse nearest-neighbor search on the on-disk index, so that a top-
kresult set contains at mostdiverse_results_kvectors sharing the same attribute value (e.g. author,doc_id, camera model). Neither strategy modifies the shared neighbor queue or the default search path — both are selected only when a diversity attribute is configured.This also removes the old
experimental_diversity_searchfeature gate (the queue-mutating prototype) and replaces it with these two cleaner designs.Motivation
A plain KNN search over a corpus where many vectors share an attribute value returns highly redundant results (e.g. 30%+ of neighbors from a single bucket). Enforcing "at most
kper bucket" against a diverse ground truth requires either filtering after search or fetching a larger candidate pool. This PR implements and benchmarks both approaches.Algorithms
Design A — post-processing over a fixed pool (
AttributeBucketDiversity,diskann-disk/src/search/provider/disk_provider.rs)Runs an ordinary greedy beam search with the caller's
L, then applies a post-processor over the top-Lpool:distance_cache; only cache-missed candidates trigger a vertex load +evaluate_similarityrecompute.HashMap<bucket, count>; emit a candidate only while its bucket count is< diverse_results_k. The output buffer truncates the stream tok.Cost is essentially free over a normal search (no extra IO), but recall depends on the caller choosing an
Llarge enough that the top-Lpool holds enough distinct buckets.Design B — adaptive-L over-fetch (
DiverseAdaptiveSearch,diskann/src/graph/search/diverse_adaptive_search.rs)Layers a traversal-time sampler on top of Design A to size
Lper query:HashMap<bucket, count>untilsample_countnodes are seen.(Σ_b min(count_b, k)) / visited— i.e. the fraction of visited nodes that would survive bucket selection. Fork = 1this isdistinct_buckets / visited.compute_adaptive_lpiecewise scaler (shared with inline-filter search):2^(-log10(yield))(1% → 4×, 0.1% → 8×), clamped toscale_factorL, continue the walk, then hand the enlarged pool to the same Design A post-processor.When buckets are diffuse the yield is high and
Lis left unchanged, so Design B degrades exactly to Design A with no over-fetch. When buckets are concentrated it over-fetches proportionally to the observed skew — no per-datasetLtuning required.Results
Benchmarked against diverse ground truth (k=10,
diverse_results_k=1) on three datasets spanning both regimes; full tables and charts indiverse-search-results.md.camera): Design B wins recall at everyL(+13 over the queue prototype on Enron @ L=20; on YFCC it climbs to 96% while the queue plateaus near 81%), at a throughput cost from the extra IO.doc_id): Design B performs no over-fetch and matches Design A on recall and IO exactly; both slightly beat the queue prototype via exact-distance reranking.Recommendation: Design B — matches Design A when diversity is free and outperforms both Design A and the queue prototype when the attribute is concentrated, without per-dataset tuning.
Changes
DiverseAdaptiveSearchstrategy;compute_adaptive_l/AdaptiveLaccessors madepub(crate); removedexperimental_diversity_searchfeature gate acrossdiskann,diskann-disk,diskann-providers,diskann-bftree.AttributeBucketDiversitypost-processor;SearchMode::DiverseAttribute { provider, diverse_attribute_id, diverse_results_k, adaptive_l }dispatching to plain KNN (A) or adaptive search (B).FileAttributeProvider, disk-search config fields for attributes / diversity / adaptive-L.Testing
diverse_matched_slotsand adaptive-L growth behavior.