Learning to Select: Query-Aware Adaptive Dimension Selection for Dense Retrieval

TL;DR

Dense retrieval often suffers from redundancy in high-dimensional embeddings, where many coordinates contribute little to a given information need. The paper introduces a query-aware dimension selection framework that distills per-query dimension importance from relevance labels into a predictor that masks uninformative coordinates at inference, while leaving document embeddings and indices unchanged. This approach yields strong, consistent improvements across multiple dense retrievers and datasets, with peak gains typically emerging when retaining ~20–40% of dimensions and practical deployment around 30%. By aligning the retained representation with query-specific relevance signals, the method delivers effective, annotation-driven masking that enhances ranking without reindexing or pseudo-labeling noise.

Abstract

Dense retrieval represents queries and docu-002 ments as high-dimensional embeddings, but003 these representations can be redundant at the004 query level: for a given information need, only005 a subset of dimensions is consistently help-006 ful for ranking. Prior work addresses this via007 pseudo-relevance feedback (PRF) based dimen-008 sion importance estimation, which can produce009 query-aware masks without labeled data but010 often relies on noisy pseudo signals and heuris-011 tic test-time procedures. In contrast, super-012 vised adapter methods leverage relevance labels013 to improve embedding quality, yet they learn014 global transformations shared across queries015 and do not explicitly model query-aware di-016 mension importance. We propose a Query-017 Aware Adaptive Dimension Selection frame-018 work that learns to predict per-dimension im-019 portance directly from query embedding. We020 first construct oracle dimension importance dis-021 tributions over embedding dimensions using022 supervised relevance labels, and then train a023 predictor to map a query embedding to these024 label-distilled importance scores. At inference,025 the predictor selects a query-aware subset of026 dimensions for similarity computation based027 solely on the query embedding, without pseudo-028 relevance feedback. Experiments across multi-029 ple dense retrievers and benchmarks show that030 our learned dimension selector improves re-031 trieval effectiveness over the full-dimensional032 baseline as well as PRF-based masking and033 supervised adapter baselines.

Figures (8)

• Figure 1: Our query-aware dimension selection pipeline. We construct a per-query dimension-importance target by contrasting aggregated embeddings of relevant documents against aggregated embeddings of hard negatives. A query-only predictor is trained to match this target distribution via KL divergence. At inference, the predicted importance selects top‑k query dimensions (masking the rest), while document embeddings and the ANN index remain unchanged.
• Figure 2: NDCG@10 as a function of retained dimension ratio $k/D$. Each panel corresponds to one model; our method consistently reaches a peak and plateau around 20--50% of the dimensions.
• Figure 3: Oracle-style dimension selection improves upper bounds, with gains peaking near 20% and further boosted by positive weighting and hard negative mining.
• Figure 4: Effect of temperature.
• Figure 5: Effect of K.