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

arXiv: https://arxiv.org/abs/2602.03306 · HTML: https://mlsanity.com/paper/2602.03306/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 11 of its 19 references that are in the map.
- First posted on arXiv 3 Feb 2026.

## Builds on (most cited first, 11 of 11)

- [Dense Passage Retrieval for Open-Domain Question Answering](https://mlsanity.com/paper/2004.04906/) (2020)
- [Billion-scale similarity search with GPUs](https://mlsanity.com/paper/1702.08734/) (2017)
- [Approximate Nearest Neighbor Negative Contrastive Learning for Dense Text Retrieval](https://mlsanity.com/paper/2007.00808/) (2020)
- [LLM2Vec: Large Language Models Are Secretly Powerful Text Encoders](https://mlsanity.com/paper/2404.05961/) (2024)
- [SPLADE: Sparse Lexical and Expansion Model for First Stage Ranking](https://mlsanity.com/paper/2107.05720/) (2021)
- [Generative Representational Instruction Tuning](https://mlsanity.com/paper/2402.09906/) (2024)
- [Matryoshka-Adaptor: Unsupervised and Supervised Tuning for Smaller Embedding Dimensions](https://mlsanity.com/paper/2407.20243/) (2024)
- [SMEC: Rethinking Matryoshka Representation Learning for Retrieval Embedding Compression](https://mlsanity.com/paper/2510.12474/) (2025)
- [Search-Adaptor: Embedding Customization for Information Retrieval](https://mlsanity.com/paper/2310.08750/) (2023)
- [ECLIPSE: Contrastive Dimension Importance Estimation with Pseudo-Irrelevance Feedback for Dense Retrieval](https://mlsanity.com/paper/2412.14967/) (2024)
- [Statistical Foundations of DIME: Risk Estimation for Practical Index Selection](https://mlsanity.com/paper/2601.05649/) (2026)

## Abstract (arXiv)

Dense retrieval represents queries and documents as high-dimensional embeddings, but these representations can be redundant at the query level: for a given information need, only a subset of dimensions is consistently helpful for ranking. Prior work addresses this via pseudo-relevance feedback (PRF) based dimension importance estimation, which can produce query-aware masks without labeled data but often relies on noisy pseudo signals and heuristic test-time procedures. In contrast, supervised adapter methods leverage relevance labels to improve embedding quality, yet they learn global transformations shared across queries and do not explicitly model query-aware dimension importance. We propose a Query-Aware Adaptive Dimension Selection framework that learns to predict per-dimension importance directly from query embedding. We first construct oracle dimension importance distributions over embedding dimensions using supervised relevance labels, and then train a predictor to map a query embedding to these label-distilled importance scores. At inference, the predictor selects a query-aware subset of dimensions for similarity computation based solely on the query embedding, without pseudo-relevance feedback. Experiments across multiple dense retrievers and benchmarks show that our learned dimension selector improves retrieval effectiveness over the full-dimensional baseline as well as PRF-based masking and supervised adapter baselines.