Skip to content

(diversity) attribute diversity implementation with post process and adaptive L#1250

Draft
narendatha wants to merge 15 commits into
mainfrom
u/narendatha/diverse-postprocess-adaptive
Draft

(diversity) attribute diversity implementation with post process and adaptive L#1250
narendatha wants to merge 15 commits into
mainfrom
u/narendatha/diverse-postprocess-adaptive

Conversation

@narendatha

@narendatha narendatha commented Jul 10, 2026

Copy link
Copy Markdown
Contributor

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-k result set contains at most diverse_results_k vectors 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_search feature 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 k per 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-L pool:

  1. Rerank with exact distances. Full-precision distances computed during traversal are reused from the accessor's distance_cache; only cache-missed candidates trigger a vertex load + evaluate_similarity recompute.
  2. Sort candidates ascending by exact distance (true nearest-first order).
  3. Greedy bucket selection. Walk the sorted candidates keeping a HashMap<bucket, count>; emit a candidate only while its bucket count is < diverse_results_k. The output buffer truncates the stream to k.

Cost is essentially free over a normal search (no extra IO), but recall depends on the caller choosing an L large enough that the top-L pool 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 L per query:

  1. During the greedy walk, sample the attribute bucket of each newly visited node into a HashMap<bucket, count> until sample_count nodes are seen.
  2. Compute the diversity yield as (Σ_b min(count_b, k)) / visited — i.e. the fraction of visited nodes that would survive bucket selection. For k = 1 this is distinct_buckets / visited.
  3. Feed that yield into the existing compute_adaptive_l piecewise scaler (shared with inline-filter search):
    • yield ≥ 50% → 1× (no growth)
    • 10–50% → 2×
    • < 10% → log-scale, 2^(-log10(yield)) (1% → 4×, 0.1% → 8×), clamped to scale_factor
  4. Resize the search-list scratch once to the new L, continue the walk, then hand the enlarged pool to the same Design A post-processor.

When buckets are diffuse the yield is high and L is 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-dataset L tuning required.

Results

Benchmarked against diverse ground truth (k=10, diverse_results_k=1) on three datasets spanning both regimes; full tables and charts in diverse-search-results.md.

  • Concentrated attributes (Enron, YFCC camera): Design B wins recall at every L (+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.
  • Diffuse attribute (Caselaw 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

  • Core: new DiverseAdaptiveSearch strategy; compute_adaptive_l / AdaptiveL accessors made pub(crate); removed experimental_diversity_search feature gate across diskann, diskann-disk, diskann-providers, diskann-bftree.
  • Disk: AttributeBucketDiversity post-processor; SearchMode::DiverseAttribute { provider, diverse_attribute_id, diverse_results_k, adaptive_l } dispatching to plain KNN (A) or adaptive search (B).
  • Benchmark: FileAttributeProvider, disk-search config fields for attributes / diversity / adaptive-L.
  • Docs: benchmark results write-up with recall/QPS tables and Mermaid charts.

Testing

  • Unit tests for diverse_matched_slots and adaptive-L growth behavior.
  • Existing diversity search + queue retain/truncate tests migrated off the removed feature gate.

Note: QPS figures in the doc are best-of-N single-threaded on a contended dev machine — treat them as indicative; recall and IO counts are the authoritative comparison.

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.
Comment thread diverse-search-results.md
@@ -0,0 +1,225 @@
# Attribute-Diversity Search — Benchmark Results

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This doc and the attached plots should be moved to the Wiki.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@narendatha narendatha Jul 17, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 leaves L unchanged.

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.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 denominator num_buckets·k caps 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.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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·k already exceeds the sample, so both paths cap the denominator at sample_visited.
  • Few buckets (YFCC): fallback grows L more 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.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@narendatha narendatha Jul 17, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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_k across more datasets (Enron, YFCC, Caselaw doc_id) to map where it helps vs. is a no-op, and
  • explore auto-deriving yield_k from 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.
Comment thread diskann/src/graph/search/diverse_adaptive_search.rs
@@ -37,7 +37,7 @@ use serde::{Deserialize, Serialize};
use crate::{

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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-commenter

codecov-commenter commented Jul 17, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 28.72629% with 263 lines in your changes missing coverage. Please review.
✅ Project coverage is 90.01%. Comparing base (1552e51) to head (0ae8f2a).

Files with missing lines Patch % Lines
...iskann/src/graph/search/diverse_adaptive_search.rs 42.91% 133 Missing ⚠️
diskann-disk/src/search/provider/disk_provider.rs 0.00% 96 Missing ⚠️
diskann-disk/src/search/search_mode.rs 0.00% 13 Missing ⚠️
diskann-benchmark/src/inputs/disk.rs 0.00% 12 Missing ⚠️
diskann/src/graph/search/inline_filter_search.rs 50.00% 6 Missing ⚠️
diskann/src/neighbor/diverse_priority_queue.rs 0.00% 3 Missing ⚠️
Additional details and impacted files

Impacted file tree graph

@@            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     
Flag Coverage Δ
miri 90.01% <28.72%> (-0.16%) ⬇️
unittests 89.68% <28.72%> (-0.16%) ⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines Coverage Δ
diskann-benchmark/src/utils/mod.rs 85.55% <ø> (ø)
diskann-providers/src/index/diskann_async.rs 96.26% <ø> (+0.04%) ⬆️
diskann/src/graph/misc.rs 90.47% <ø> (+10.47%) ⬆️
diskann/src/neighbor/mod.rs 98.96% <ø> (ø)
diskann/src/neighbor/queue.rs 98.93% <ø> (+0.61%) ⬆️
diskann/src/neighbor/diverse_priority_queue.rs 98.45% <0.00%> (ø)
diskann/src/graph/search/inline_filter_search.rs 95.16% <50.00%> (-2.80%) ⬇️
diskann-benchmark/src/inputs/disk.rs 1.36% <0.00%> (-0.08%) ⬇️
diskann-disk/src/search/search_mode.rs 80.00% <0.00%> (-11.96%) ⬇️
diskann-disk/src/search/provider/disk_provider.rs 87.77% <0.00%> (-6.16%) ⬇️
... and 1 more

... and 2 files with indirect coverage changes

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

///
/// 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(

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this doesn't need to be pub(crate) anymore

inner: Knn,
provider: Arc<P>,
diverse_results_k: usize,
adaptive_l: AdaptiveL,

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@magdalendobson

magdalendobson commented Jul 17, 2026

Copy link
Copy Markdown
Contributor

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants