Learning to Select: Query-Aware Adaptive Dimension Selection for Dense Retrieval
Paper · 2026 · Zhanyu Wu, Richong Zhang and Zhijie Nie
Summary
Dense retrievers encode queries and documents as high-dimensional vectors, but for any single query only part of those dimensions actually helps ranking. Earlier DIME-style methods, including Eclipse, estimate per-query dimension importance from pseudo-relevance feedback or LLM-written pseudo-documents. That signal is noisy and adds heuristic steps at inference. Supervised search adapters, by contrast, learn one global transformation shared by all queries.
The authors combine the strengths of both approaches. Using relevance labels, they build an oracle importance distribution for each training query. Each coordinate is scored by contrasting a gain-weighted centroid of relevant documents with the mean of sampled hard negatives, and the scores are then softened with a temperature softmax. A single fully connected layer on frozen query embeddings is trained with KL divergence to reproduce this distribution. At test time it keeps the top-k query dimensions and zeros out the rest, so document embeddings and ANN indexes such as FAISS stay untouched.
On encoders including Qwen-8B and GritLM and on benchmarks like MS MARCO, SciFact and NFCorpus, the selector outperforms full-dimensional retrieval, PRF masking and adapters. This makes it a cheap way to improve effectiveness without reindexing.
Key points
- Results are reported as peak NDCG@10, chosen from a sweep of the retained-dimension ratio in 2% steps. On MS MARCO with Qwen-8B, scores rise from 0.646 (baseline) to 0.714 with the selector, compared with 0.698 for the Adapter and 0.668 for Eclipse. With GritLM they rise from 0.536 to 0.632.
- Stacking the selector on adapter outputs gives the best MS MARCO scores: 0.731 for Qwen-8B and 0.637 for GritLM. On SciFact the combination falls back to the adapter level of 0.883, while the selector alone reaches 0.906 with GritLM.
- A Jaccard analysis of top-2048 dimension sets over 20,000 query pairs explains this pattern. On SciFact, adapters make the selected dimensions much more uniform across queries (overlap rises from about 0.34 to roughly 0.48–0.51). MS MARCO changes only modestly (from about 0.36–0.37 to 0.39–0.40).
- Training on LLM-generated pseudo-queries instead of labels recovers much of the gain on MS MARCO. On SciFact and NFCorpus, however, it helps only marginally.
- The authors note three limitations: the method depends on good labels, it is capped by the quality of the frozen encoder, and it transfers weakly across domains. Code is available at github.com/wzy2001wzy/learning2select.
Institutions
Fields, topics and concepts
📚 Shelves 6The field of science this work belongs to (OpenAlex fields).Computer Science · 18,614 worksArtificial Intelligence · 17,965 worksArts and Humanities · 17,450 worksLanguage and Linguistics · 17,296 worksComputer Vision and Pattern Recognition · 3,486 worksInformation Systems · 1,730 works
🧭 Topics 1Research subjects and languages the work is about.Information retrieval & search · 552 works
Cite
Cite as Zhanyu Wu, Richong Zhang and Zhijie Nie. 2026. Learning to Select: Query-Aware Adaptive Dimension Selection for Dense Retrieval. Proceedings of the 64th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers).
@inproceedings{wu2026learning,
title = {{Learning to Select: Query-Aware Adaptive Dimension Selection for Dense Retrieval}},
author = {Zhanyu Wu, Richong Zhang and Zhijie Nie},
booktitle = {Proceedings of the 64th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)},
year = {2026},
doi = {10.18653/v1/2026.acl-long.849},
url = {https://doi.org/10.18653/v1/2026.acl-long.849}
}References
- Parishad BehnamGhader, Vaibhav Adlakha, Marius Mosbach, Dzmitry Bahdanau, Nicolas Chapados, and Siva Reddy. 2024. Llm2vec: Large language models are secretly powerful text encoders. arXiv preprint arXiv:2404.05961.
- Cesare Campagnano, Antonio Mallia, and Fabrizio Silvestri. 2025. Unveiling dime: Reproducibility, generalizability, and formal analysis of dimension importance estimation for dense retrieval. In Proceedings of the 48th International ACM SIGIR Conference on Research and Development in Information Ret… find
- Giulio D’Erasmo, Cesare Campagnano, Antonio Mallia, Pierpaolo Brutti, Nicola Tonellotto, and Fabrizio Silvestri. 2026. Statistical foundations of dime: Risk estimation for practical index selection. arXiv preprint arXiv:2601.05649. find
- Giulio D’Erasmo, Giovanni Trappolini, Fabrizio Silvestri, and Nicola Tonellotto. 2025. Eclipse: Contrastive dimension importance estimation with pseudo-irrelevance feedback for dense retrieval. In Proceedings of the 2025 International ACM SIGIR Conference on Innovative Concepts and Theories in Infor… find
- Guglielmo Faggioli, Nicola Ferro, Raffaele Perego, and Nicola Tonellotto. 2024. Dimension importance estimation for dense information retrieval. In Proceedings of the 47th International ACM SIGIR Conference on Research and Development in Information Retrieval, pages 1318–1328. find
- Guglielmo Faggioli, Nicola Ferro, Raffaele Perego, and Nicola Tonellotto. 2025a. Codime: a counterfactual approach for dimension importance estimation through click logs. In Proceedings of the 48th International ACM SIGIR Conference on Research and Development in Information Retrieval, pages 2330– 2… find
- Guglielmo Faggioli, Nicola Ferro, Raffaele Perego, and Nicola Tonellotto. 2025b. Getting off the dime: Dimension pruning via dimension importance estimation for dense information retrieval. ACM Transactions on Information Systems, 44(1):1–34. find
- Thibault Formal, Carlos Lassance, Benjamin Piwowarski, and Stéphane Clinchant. 2021a. Splade v2: Sparse lexical and expansion model for information retrieval. arXiv preprint arXiv:2109.10086.
- Thibault Formal, Benjamin Piwowarski, and Stéphane Clinchant. 2021b. Splade: Sparse lexical and expansion model for first stage ranking. In Proceedings of the 44th International ACM SIGIR Conference on Research and Development in Information Retrieval, pages 2288–2292.
- Jeff Johnson, Matthijs Douze, and Hervé Jégou. 2019. Billion-scale similarity search with gpus. IEEE Transactions on Big Data, 7(3):535–547.
- Vladimir Karpukhin, Barlas Oguz, Sewon Min, Patrick SH Lewis, Ledell Wu, Sergey Edunov, Danqi Chen, and Wen-tau Yih. 2020. Dense passage retrieval for open-domain question answering. In EMNLP (1), pages 6769–6781.
- Aditya Kusupati, Gantavya Bhatt, Aniket Rege, Matthew Wallingford, Aditya Sinha, Vivek Ramanujan, William Howard-Snyder, Kaifeng Chen, Sham Kakade, Prateek Jain, and 1 others. 2022. Matryoshka representation learning. Advances in Neural Information Processing Systems, 35:30233–30249. find
- Niklas Muennighoff, SU Hongjin, Liang Wang, Nan Yang, Furu Wei, Tao Yu, Amanpreet Singh, and Douwe Kiela. 2024. Generative representational instruction tuning. In ICLR 2024 Workshop: How Far Are We From AGI.
- OpenAI. 2024. New embedding models and api updates. https://openai.com/blog/ new-embedding-models-and-api-updates. Accessed: 2024-01-25. ↗
- Stephen Robertson, Hugo Zaragoza, and 1 others. 2009. The probabilistic relevance framework: Bm25 and beyond. Foundations and Trends® in Information Retrieval, 3(4):333–389. find
- Joseph John Rocchio Jr. 1971. Relevance feedback in information retrieval. The SMART retrieval system: experiments in automatic document processing. find
- Lee Xiong, Chenyan Xiong, Ye Li, Kwok-Fung Tang, Jialin Liu, Paul Bennett, Junaid Ahmed, and Arnold Overwijk. 2020. Approximate nearest neighbor negative contrastive learning for dense text retrieval. arXiv preprint arXiv:2007.00808.
- Jinsung Yoon, Yanfei Chen, Sercan Arik, and Tomas Pfister. 2024a. Search-adaptor: Embedding customization for information retrieval. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 12230–12247. find
- Jinsung Yoon, Rajarishi Sinha, Sercan O Arik, and Tomas Pfister. 2024b. Matryoshka-adaptor: Unsupervised and supervised tuning for smaller embedding dimensions. In Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing, pages 10318–10336. find
- Biao Zhang, Lixin Chen, Tong Liu, and Bo Zheng. 2025a. Smec: Rethinking matryoshka representation learning for retrieval embedding compression. In Proceedings of the 2025 Conference on Empirical Methods in Natural Language Processing, pages 26220–26233. find
Updated
This page in Markdown → https://md.research.elamur.ai/free/en/papers/wu-2026-learning-select-query-aware-adaptive/