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

Zhanyu Wu, Richong Zhang, Zhijie Nie


Abstract
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.
Anthology ID:
2026.acl-long.849
Volume:
Proceedings of the 64th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)
Month:
July
Year:
2026
Address:
San Diego, California, United States
Editors:
Maria Liakata, Viviane P. Moreira, Jiajun Zhang, David Jurgens
Venue:
ACL
SIG:
Publisher:
Association for Computational Linguistics
Note:
Pages:
18666–18677
Language:
URL:
https://preview.aclanthology.org/ingest-acl/2026.acl-long.849/
DOI:
Bibkey:
Cite (ACL):
Zhanyu Wu, Richong Zhang, and Zhijie Nie. 2026. Learning to Select: Query-Aware Adaptive Dimension Selection for Dense Retrieval. In Proceedings of the 64th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 18666–18677, San Diego, California, United States. Association for Computational Linguistics.
Cite (Informal):
Learning to Select: Query-Aware Adaptive Dimension Selection for Dense Retrieval (Wu et al., ACL 2026)
Copy Citation:
PDF:
https://preview.aclanthology.org/ingest-acl/2026.acl-long.849.pdf
Checklist:
 2026.acl-long.849.checklist.pdf