CAFE: Retrieval Head-based Coarse-to-Fine Information Seeking to Enhance Multi-Document QA Capability

TL;DR

This paper addresses the challenge of retrieval and reasoning in long-context multi-document QA by LLMs, where adding more documents degrades performance. It proposes CAFE, a training-free, two-stage coarse-to-fine framework that leverages internal retrieval heads to first filter background documents and then steer attention to the most relevant content. By pre-detecting suitable retrieval heads and applying coarse-grained filtering followed by fine-grained attention steering, CAFE reduces noise and improves reliance on gold evidence; it demonstrates substantial gains across five benchmarks and three models, including up to 22.1% SubEM improvement over SFT and 13.7% over RAG on Mistral. The approach also shows robustness to longer inputs, ablation-supported necessity of both stages, and manageable but higher inference latency, suggesting practical applicability to broad long-document reasoning tasks.

Abstract

Advancements in Large Language Models (LLMs) have extended their input context length, yet they still struggle with retrieval and reasoning in long-context inputs. Existing methods propose to utilize the prompt strategy and retrieval head to alleviate this limitation. However, they still face challenges in balancing retrieval precision and recall, impacting their efficacy in answering questions. To address this, we introduce $\textbf{CAFE}$, a two-stage coarse-to-fine method to enhance multi-document question-answering capacities. By gradually eliminating the negative impacts of background and distracting documents, CAFE makes the responses more reliant on the evidence documents. Initially, a coarse-grained filtering method leverages retrieval heads to identify and rank relevant documents. Then, a fine-grained steering method guides attention to the most relevant content. Experiments across benchmarks show CAFE outperforms baselines, achieving up to 22.1% and 13.7% SubEM improvement over SFT and RAG methods on the Mistral model, respectively.

Figures (5)

• Figure 1: LLMs' performance on HotpotQA varies with the number of input documents. Solid lines represent performance with the gold document, while dashed lines show performance as more documents are added.
• Figure 2: Overall framework of our proposed CAFE approach. The red, blue, and yellow bar charts represent the gold, distracting, and background documents, respectively.
• Figure 3: The impact of hyperparameters $M$ (documents per retrieval head) and $K$ (retrieval heads) on Llama-3.1-8B-Instruct. The top row shows recall and performance for coarse-grained filtering, while the bottom row illustrates changes for fine-grained steering.
• Figure 4: Attention distribution of retrieval heads used in the coarse-grained filtering stage.
• Figure 5: Attention distribution of retrieval heads used in the fine-grained steering stage.