AdaWidth: Query-Adaptive Embedding Width for Dense Retrieval
Shubing Yang, Dongfang Zhao
cs.IR
2026-08-25
A shared orthogonal rotation packs signal into prefixes and a router picks width per query, matching SOTA NDCG@10 with 55%–84% fewer dimensions.
Dense retrievers score 768- to 4096-dimensional vectors, and cost scales with width. Compression has two standard answers. Nested representations (Matryoshka and post-hoc adapters) give every query the same prefix length, which must cover the hardest queries. Query-specific coordinate pickers (DIME, Learning-to-Select) keep a different, usually scattered, subset per query, so the index still stores full width and scoring gathers from a full record.
Queries in one task differ widely in how many dimensions they need before the ranking settles. Some stabilize on a short prefix; others need much more. AdaWidth keeps one shared coordinate order and varies only how far along that order each query is scored.
The encoder stays frozen. Two modules sit between it and the index.
An orthogonal prefix adapter is a product of 64 Householder reflections, the same rotation R for queries and documents. Orthogonality leaves every full-width inner product unchanged, so full-width behavior matches the original encoder exactly. Training can only move discriminative mass from the tail into the prefix. The loss sums multi-positive contrastive terms at widths 64, 128, and 256, plus score matching and geometry matching against the frozen full-width teacher, with extra weight on contestable candidates. At F=64 and D=2048 the adapter has 131k parameters. At deploy time it collapses to an explicit matrix; documents are rotated offline once.
A prefix router first ranks the shared 500-candidate shortlist at 64 dimensions (the shortlist is mined once from the frozen encoder at 64, for every method). It reads 18 order statistics of that ranking: top scores, the gap across the cutoff, entropy of the top-10, and similar, then predicts whether widening to 128 or 256 would change NDCG@10. Only a fraction ρ of queries pay for the wider prefix. Partial sums from stage one are reused, so mean cost is shortlist size times mean width, not two full passes.
A prefix-sufficiency argument says a relevant document leaves the top-k exactly when at least k non-relevant ones overtake it at that width. Required width is an order statistic of competition at the cutoff: it grows with log corpus size, falls with log retrieval depth, and is heavy-tailed across queries.
Six tasks (FiQA, ArguAna, Quora, a 2M MS MARCO slice, OK-VQA, A-OKVQA) times five frozen encoders (E5-Mistral-7B down to Nomic 137M, plus Qwen3-VL-2B).
On text, AdaWidth leads three retrained baselines at nearly every operating point. Mean NDCG@10 is +5.88 versus Matryoshka-Adaptor, +9.94 versus SMEC, +15.77 versus Learning-to-Select. At 32 retained coordinates Matryoshka needs 2.36× the width to match; at 48, 3.04×. AdaWidth at 64 already exceeds Matryoshka at 256. Absolute averages on aligned task–encoder pairs: 0.431 at 32, 0.534 at 64, 0.597 at 256, against Matryoshka's 0.377, 0.472, 0.510. Matching the others' NDCG@10 uses 55% to 84% fewer dimensions per query.
Ablations show where the money is. Mean NDCG across widths: prefix truncation 0.4027, router only 0.4158, adapter only 0.6121, both 0.6134. The adapter takes almost all of the gain; the router adds 0.13 overall and 0.45 on corpora larger than 100k documents. Sweeping Householder count from 8 to 128 moves NDCG by 1.03; F=8 (13,824 parameters) already leads the strongest baseline by more than 7 points. Fitted width versus N and k has R² 0.926–0.973 on text, with the predicted signs.
A production encoder can stay frozen. Install an inner-product-preserving rotation in front of the index, then decide per query how far to score. Prefixes are contiguous, so the index stores N×h values and a query reads exactly the coordinates it multiplies, cleaner than a scattered per-query mask. Gains show up on both Matryoshka-pretrained and non-Matryoshka encoders, so they do not require a pretrained prefix order.
This saves arithmetic at rerank time. It is not a claim about scanning fewer dimensions inside a full-corpus ANN. The shortlist is mined at 64 dimensions; missed relevant documents are gone.
The protocol reranks a 500-candidate shortlist, not end-to-end nearest neighbor search. How prefix width should be cut on an ANN graph or inverted file is untested. The router needs a labeled train split to fit NDCG gain; unlabeled settings keep only the adapter. Adapter-plus-router is nearly tied with adapter alone, so the adaptive-width story is mostly the rotation. An orthogonal map cannot raise full-width quality; the ceiling is the frozen encoder. MS MARCO is a 2M subsample. Hyperparameters were swept on SciFact and NFCorpus, then frozen for the reported grid.