Improving Neural Retrieval with Attribution-Guided Query Rewriting

TL;DR

This work tackles the brittleness of neural information retrieval when queries are underspecified by introducing an attribution-guided query rewriting framework. It computes token-level contributions from the retriever using Integrated Gradients over the top-$k$ retrieved documents and uses these scores as soft guidance in an LLM prompt to rewrite the query while preserving intent. The rewritten query is then submitted to the same retriever in a closed loop, enabling retriever-aware improvements without retraining. Empirical results on BEIR with SPLADE and TCT-ColBERT show consistent gains in $\text{nDCG}$, $\text{MAP}$, and precision, especially for implicit or domain-specific needs, demonstrating that explanations can be effectively leveraged to improve neural retrieval in practice.

Abstract

Neural retrievers are effective but brittle: underspecified or ambiguous queries can misdirect ranking even when relevant documents exist. Existing approaches address this brittleness only partially: LLMs rewrite queries without retriever feedback, and explainability methods identify misleading tokens but are used for post-hoc analysis. We close this loop and propose an attribution-guided query rewriting method that uses token-level explanations to guide query rewriting. For each query, we compute gradient-based token attributions from the retriever and then use these scores as soft guidance in a structured prompt to an LLM that clarifies weak or misleading query components while preserving intent. Evaluated on BEIR collections, the resulting rewrites consistently improve retrieval effectiveness over strong baselines, with larger gains for implicit or ambiguous information needs.