질의 인지적 적응형 차원 선택을 위한 밀집 검색: 선택 학습

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

밀집 검색은 질의와 문서를 고차원 임베딩으로 표현하지만, 이러한 표현은 질의 수준에서 중복될 수 있습니다. 즉, 특정 정보 요구 사항에 대해, 순위 결정에 도움이 되는 것은 전체 차원 중 일부에 불과합니다. 기존 연구에서는 가짜 관련성 피드백(PRF) 기반의 차원 중요도 추정 방법을 사용하여 레이블이 없는 데이터로 질의 인지적 마스크를 생성하지만, 종종 노이즈가 많은 가짜 신호와 휴리스틱한 런타임 절차에 의존합니다. 반면, 지도 학습 기반 어댑터 방법은 관련성 레이블을 활용하여 임베딩 품질을 향상시키지만, 모든 질의에 대해 공유되는 전역 변환을 학습하며 질의 인지적 차원 중요도를 명시적으로 모델링하지 않습니다. 본 논문에서는 질의 임베딩으로부터 각 차원의 중요도를 직접 예측하는 질의 인지적 적응형 차원 선택 프레임워크를 제안합니다. 먼저, 지도 학습을 통해 얻은 관련성 레이블을 사용하여 임베딩 차원에 대한 이상적인 중요도 분포를 구축하고, 이 분포를 활용하여 질의 임베딩을 해당 레이블로 정제된 중요도 점수로 매핑하는 예측 모델을 학습합니다. 추론 시, 예측 모델은 가짜 관련성 피드백 없이 질의 임베딩만을 사용하여 질의 인지적인 차원 부분 집합을 선택하여 유사도 계산을 수행합니다. 여러 밀집 검색 모델과 벤치마크에 대한 실험 결과, 제안하는 학습된 차원 선택기는 전체 차원을 사용하는 기준 모델뿐만 아니라 PRF 기반 마스킹 및 지도 학습 기반 기준 모델보다 검색 성능을 향상시킵니다.

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.