[Experiment do not review - for collab with Magdalen] Decoupled yield_k knob for Design B attribute-diversity search

[narendatha]

Draft — not for merge. This branches off u/narendatha/diverse-postprocess-adaptive
(the Design B attribute-diversity PR) to explore a follow-up idea raised in review.
It is intended as a home for experimentation across datasets, not as shippable
surface area yet. Headline finding: yield_k is a per-dataset tradeoff, not a
universal win — see results below.

Motivation

Design B grows the search list L mid-search from the observed bucket
concentration of the attribute sample. The yield estimator caps each bucket's
contribution at k slots:

matched = Σ_b min(count_b, k)
achievable = min(sample_visited, num_buckets · k)
yield = matched / achievable (∈ [0, 1])

Today k is the same value as diverse_results_k (the per-bucket output cap).
On a distance-correlated attribute — where the near neighbors of a query all
share a few attribute values — the near buckets get "counted as satisfied" after
a single far-away hit. Yield saturates early, adaptive-L stops growing, and the
search terminates before the close attribute-satisfying points are found.
Jurisdiction (60 buckets) regressed slightly below the plain-queue baseline for
exactly this reason.

What this changes

Decouple the sampler's per-bucket cap (yield_k) from the output cap
(diverse_results_k):

• Requiring yield_k points per bucket before that bucket is considered
  "satisfied" changes when yield saturates, and hence how aggressively L grows.
• yield_k defaults to the effective diverse_results_k, so existing configs
  are unchanged (yield_k = 1 reproduces the current Design B behavior
  bit-for-bit).

Plumbing

Layer	Change
diskann/src/graph/search/diverse_adaptive_search.rs	DiverseAdaptiveSearch::new takes yield_k; yield helpers (diverse_matched_slots, diverse_yield) parameterized on it
diskann-disk/src/search/search_mode.rs	SearchMode::DiverseAttribute { …, yield_k } variant + diverse_attribute() constructor param
diskann-disk/src/search/provider/disk_provider.rs	passes *yield_k to the adaptive search; AttributeBucketDiversity processor still uses *diverse_results_k
diskann-benchmark/src/inputs/graph_index.rs	optional yield_k field on the adaptive_l config block (serde default None)
diskann-benchmark/src/disk_index/search.rs	resolves yield_k, defaulting to the effective diverse_results_k

Results — full yield_k sweep

All configs: sample_count=100, scale_factor=4.0, diverse_results_k=1,
recall@10, search_list=[20,40,80,100]. yield_k=1 == current Design B.

Jurisdiction (Caselaw, ~60 buckets, distance-correlated) — helps strongly

yield_k	L=20	L=40	L=80	L=100	mean_ios @ L=100
1	57.30	68.23	76.62	78.97	192.8
2	59.36	69.84	77.84	80.15	213.0
3	60.84	71.08	78.81	80.92	231.2
4	62.18	71.90	79.53	81.55	247.0

YFCC camera (~67 buckets) — helps slightly

yield_k	L=20	L=40	L=80	L=100	mean_ios @ L=100
1	88.44	94.20	96.45	96.92	286.0
2	89.32	94.52	96.43	96.92	307.7
3	89.85	94.69	96.75	97.03	320.5
4	90.19	94.90	96.77	96.95	329.8

Caselaw doc_id (diffuse) — exact no-op (identical for all yield_k)

yield_k	L=20	L=40	L=80	L=100	mean_ios @ L=100
1–4	92.04	97.53	99.06	99.25	~112

Enron (~12.8k buckets, concentrated) — hurts

yield_k	L=20	L=40	L=80	L=100	mean_ios @ L=100
1	79.92	87.38	92.12	93.01	319.8
2	78.37	86.52	91.66	92.57	292.3
3	77.34	85.95	91.33	92.33	273.6
4	76.17	85.34	90.79	92.05	257.9

Why the sign flips — regime analysis

The behavior is fully explained by which term wins in
achievable = min(sample_visited, num_buckets · yield_k):

• num_buckets ≫ sample_count (Enron): denominator is pinned at
  sample_visited. Raising yield_k only loosens the per-bucket caps, which
  raises matched → raises yield → less over-fetch → fewer IOs and lower
  recall. yield_k = 1 is optimal.
• num_buckets ≲ sample_count (jurisdiction, YFCC): denominator is
  num_buckets · yield_k. Raising yield_k lowers yield → more over-fetch
  → higher recall at the cost of more IOs.
• Diffuse attribute (doc_id): buckets rarely exceed one sampled point, so
  min(count_b, yield_k) is insensitive to yield_k — an exact no-op.

Takeaway: there is no safe universal default. The right yield_k depends on
the bucket-count-vs-sample regime and the attribute/distance correlation, which
a client can't easily observe. This is the core argument for keeping it out of the
main Design B PR.

Open questions

1. Auto-derivation instead of a knob. The "helpful" regime is exactly
   num_buckets ≤ sample_count, which is cheaply detectable at runtime. Could
   yield_k be inferred from observed bucket saturation rather than configured?
2. Interaction with scale_factor. Both control over-fetch aggressiveness;
   would a smarter scale_factor heuristic subsume the useful cases?

Planned experiments

• yield_k ∈ {1,2,3,4} sweep on jurisdiction (correlated) — helps
• Sweep on Enron (concentrated, ~12.8k buckets) — hurts
• Sweep on Caselaw doc_id (diffuse) — no-op
• Sweep on YFCC camera (~67 buckets) — helps slightly
• Characterize the recall/IO tradeoff + regime analysis
• Prototype auto-derived yield_k from num_buckets vs sample_count
• Compare against a scale_factor-only heuristic

Testing

• cargo build --release -p diskann-benchmark --features disk-index ✅
• cargo test -p diskann diverse_adaptive_search — 10 tests pass ✅

[magdalendobson]

Thanks for posting these! I think the results make a lot of sense. I agree that experimenting with auto-deriving yield_k might be useful. I don't have strong intuition for whether that could be replaced by tuning scale_factor instead. Since the situation of relatively few, distance-correlated attributes seems pretty realistic (but maybe you know better than me there?), if we could achieve it without an extra parameter that would be ideal.

I feel like this is trending in the direction of some sort of auto-detection of correlation during search, but that might be a pretty big can of worms haha