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

TL;DR

This paper introduces a query-aware adaptive dimension selection method for dense retrieval, which learns to identify relevant embedding dimensions from queries to improve retrieval accuracy without relying on pseudo-relevance feedback.

Contribution

It proposes a novel framework that predicts query-specific important dimensions for dense retrieval, outperforming existing methods that use heuristics or global transformations.

Findings

Improves retrieval effectiveness over full-dimensional baselines.

Outperforms pseudo-relevance feedback-based masking methods.

Effective across multiple dense retrievers and benchmarks.

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 \emph{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.