Elamur Research⌕
☰
home / papers / 2025 / CAFE: Retrieval Head-based Coarse-to-Fine Information Seeking to Enhance Multi-Document QA Capability

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

Paper · 2025 · Han Peng, Jinhao Jiang, Zican Dong et al.

Summary

Long-context LLMs answer multi-document questions poorly once the evidence is buried among many irrelevant or distracting documents. External retrievers and prompt-based selection often miss relevant evidence (low recall), and fine-tuning attention heads needs substantial data and transfers poorly to new domains. The authors propose CAFE, a training-free, two-stage method that uses the model's own retrieval heads, which are attention heads that locate relevant tokens in long inputs. Before inference, two sets of such heads are identified on validation data. In the coarse stage, each head nominates its highest-attention documents. These are reranked by summed head scores and placed at the end of the prompt, and background documents are discarded. In the fine stage, a second head set narrows the candidates further, and attention steering strengthens the question's focus on them. Across five benchmarks and three LLMs, CAFE beats SFT and RAG baselines. Because it requires no training and applies to many tasks, it is a practical way to make long-context QA rely more on the actual evidence.

Key points

  • On Llama-3.1-8B-Instruct, retrieval-head selection reached the highest recall in a pilot study. It scored 0.97 on HotpotQA-8k (in-context selection: 0.64) and 0.93 on HotpotQA-32k (in-context selection: 0.38).
  • In an attention-masking analysis, blocking attention from the question to the gold documents dropped SubEM to roughly the no-document level. Masking attention between the gold documents themselves barely mattered.
  • CAFE improved average SubEM by up to 22.1% over SFT and 13.7% over RAG. The tested models were Llama-3.1-8B-Instruct, Mistral-3-7B-Instruct and Phi-3.5-Mini-Instruct.
  • In the SQuAD lost-in-the-middle test (50 documents), Mistral's baseline SubEM fell from 70 to about 56–58 for middle positions. CAFE stayed at roughly 78–80 across all positions.
  • Steering attention at sentence level scored lower than at document level (SubEM 65.5 vs 70.0), even below no steering (68.5). The coarse stage uses M1 = 4 documents per head and K1 = 4 heads.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

 Institutions

🎓Renmin University of ChinaBeijing · China114 works · 16,118 citations

 Fields, topics and concepts

📚 Shelves 6The field of science this work belongs to (OpenAlex fields).Computer Science · 16,667 worksArtificial Intelligence · 16,194 worksArts and Humanities · 16,039 worksLanguage and Linguistics · 15,891 worksComputer Vision and Pattern Recognition · 2,869 worksInformation Systems · 1,437 works

🧭 Topics 2Research subjects and languages the work is about.Question answering · 441 worksInformation retrieval & search · 473 works

Cite

Cite as Han Peng, Jinhao Jiang, Zican Dong, Wayne Xin Zhao and Lei Fang. 2025. CAFE: Retrieval Head-based Coarse-to-Fine Information Seeking to Enhance Multi-Document QA Capability. Proceedings of the 2025 Conference on Empirical Methods in Natural Language Processing.

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.18653/v1/2025.emnlp-main.655

@inproceedings{peng2025cafe,
  title = {{CAFE: Retrieval Head-based Coarse-to-Fine Information Seeking to Enhance Multi-Document QA Capability}},
  author = {Han Peng, Jinhao Jiang, Zican Dong, Wayne Xin Zhao and Lei Fang},
  booktitle = {Proceedings of the 2025 Conference on Empirical Methods in Natural Language Processing},
  year = {2025},
  doi = {10.18653/v1/2025.emnlp-main.655},
  url = {https://doi.org/10.18653/v1/2025.emnlp-main.655}
}

References

39 references, 20 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. Devanshu Agrawal, Shang Gao, and Martin Gajek. 2024. Can’t remember details in long documents? you need some r&r. In Findings of the Association for Computational Linguistics: EMNLP 2024, Miami, Florida, USA, November 12-16, 2024, pages 12692– 12704. Association for Computational Linguistics. find
  2. Shengnan An, Zexiong Ma, Zeqi Lin, Nanning Zheng, Jian-Guang Lou, and Weizhu Chen. 2024. Make your LLM fully utilize the context. In Advances in Neural Information Processing Systems 38: Annual Conference on Neural Information Processing Systems 2024, NeurIPS 2024, Vancouver, BC, Canada, December 10… find
  3. Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, Yuxiao Dong, Jie Tang, and Juanzi Li. 2024a. Longbench: A bilingual, multitask benchmark for long context understanding. In Proceedings of the 62nd Annual Meeting of the Associat…
  4. Yushi Bai, Shangqing Tu, Jiajie Zhang, Hao Peng, Xiaozhi Wang, Xin Lv, Shulin Cao, Jiazheng Xu, Lei Hou, Yuxiao Dong, Jie Tang, and Juanzi Li. 2024b. Longbench v2: Towards deeper understanding and reasoning on realistic long-context multitasks. CoRR, abs/2412.15204.
  5. Yuhan Chen, Ang Lv, Jian Luan, Bin Wang, and Wei Liu. 2024. Hope: A novel positional encoding without long-term decay for enhanced context awareness and extrapolation. CoRR, abs/2410.21216. find
  6. Zican Dong, Junyi Li, Jinhao Jiang, Mingyu Xu, Wayne Xin Zhao, Bingning Wang, and Weipeng Chen. 2025. Longred: Mitigating short-text degradation of long-context large language models via restoration distillation. CoRR, abs/2502.07365. find
  7. Zican Dong, Junyi Li, Xin Men, Xin Zhao, Bingning Wang, Zhen Tian, Weipeng Chen, and Ji-Rong Wen. 2024. Exploring context window of large language models via decomposed positional vectors. In Advances in Neural Information Processing Systems 38: Annual Conference on Neural Information Processing Sys… find
  8. Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, Anirudh Goyal, Anthony Hartshorn, Aobo Yang, Archi Mitra, Archie Sravankumar, Artem Korenev, Arthur Hinsvark, Arun Rao, Aston Zhang, Aurélien Rodriguez,… — The Llama 3 Herd of Models
  9. Tianyu Fu, Haofeng Huang, Xuefei Ning, Genghan Zhang, Boju Chen, Tianqi Wu, Hongyi Wang, Zixiao Huang, Shiyao Li, Shengen Yan, Guohao Dai, Huazhong Yang, and Yu Wang. 2024a. Moa: Mixture of sparse attention for automatic large language model compression. CoRR, abs/2406.14909. find
  10. Yao Fu, Rameswar Panda, Xinyao Niu, Xiang Yue, Hannaneh Hajishirzi, Yoon Kim, and Hao Peng. 2024b. Data engineering for scaling language models to 128k context. In Forty-first International Conference on Machine Learning, ICML 2024, Vienna, Austria, July 21-27, 2024. OpenReview.net.
  11. Aryo Pradipta Gema, Chen Jin, Ahmed Abdulaal, Tom Diethe, Philip Teare, Beatrice Alex, Pasquale Minervini, and Amrutha Saseendran. 2024. Decore: Decoding by contrasting retrieval heads to mitigate hallucinations. CoRR, abs/2410.18860. find
  12. Cheng-Ping Hsieh, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, Yang Zhang, and Boris Ginsburg. 2024. RULER: what’s the real context size of your long-context language models? CoRR, abs/2404.06654.
  13. Jinhao Jiang, Jiayi Chen, Junyi Li, Ruiyang Ren, Shijie Wang, Wayne Xin Zhao, Yang Song, and Tao Zhang. 2024. Rag-star: Enhancing deliberative reasoning with retrieval augmented verification and refinement. CoRR, abs/2412.12881.
  14. Vladimir Karpukhin, Barlas Oguz, Sewon Min, Patrick S. H. Lewis, Ledell Wu, Sergey Edunov, Danqi Chen, and Wen-tau Yih. 2020. Dense passage retrieval for open-domain question answering. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing, EMNLP 2020, Online, Nov…
  15. Jinhyuk Lee, Anthony Chen, Zhuyun Dai, Dheeru Dua, Devendra Singh Sachan, Michael Boratko, Yi Luan, Sébastien M. R. Arnold, Vincent Perot, Siddharth Dalmia, Hexiang Hu, Xudong Lin, Panupong Pasupat, Aida Amini, Jeremy R. Cole, Sebastian Riedel, Iftekhar Naim, Ming-Wei Chang, and Kelvin Guu. 2024. Ca… — Can Long-Context Language Models Subsume Retrieval, RAG, SQL, and More?
  16. Huayang Li, Pat Verga, Priyanka Sen, Bowen Yang, Vijay Viswanathan, Patrick Lewis, Taro Watanabe, and Yixuan Su. 2024a. Alr2 : A retrieve-then-reason framework for long-context question answering. CoRR, abs/2410.03227. — [arXiv:2410.03227](https://arxiv.org/abs/2410.03227) find
  17. Siheng Li, Cheng Yang, Zesen Cheng, Lemao Liu, Mo Yu, Yujiu Yang, and Wai Lam. 2024b. Large language models can self-improve in long-context reasoning. CoRR, abs/2411.08147. — [arXiv:2411.08147](https://arxiv.org/abs/2411.08147) find
  18. Yanyang Li, Shuo Liang, Michael R. Lyu, and Liwei Wang. 2024c. Making long-context language models better multi-hop reasoners. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), ACL 2024, Bangkok, Thailand, August 11-16, 2024, pages 24… find
  19. Nelson F. Liu, Kevin Lin, John Hewitt, Ashwin Paranjape, Michele Bevilacqua, Fabio Petroni, and Percy Liang. 2024. Lost in the middle: How language models use long contexts. Trans. Assoc. Comput. Linguistics, 12.
  20. Tsendsuren Munkhdalai, Manaal Faruqui, and Siddharth Gopal. 2024. Leave no context behind: Efficient infinite context transformers with infini-attention. CoRR, abs/2404.07143. find
  21. Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. 2016. Squad: 100, 000+ questions for machine comprehension of text. In Proceedings of the 2016 Conference on Empirical Methods in Natural Language Processing, EMNLP 2016, Austin, Texas, USA, November 1-4, 2016, pages 2383–2392. The A…
  22. Machel Reid, Nikolay Savinov, Denis Teplyashin, Dmitry Lepikhin, Timothy P. Lillicrap, Jean-Baptiste Alayrac, Radu Soricut, Angeliki Lazaridou, Orhan Firat, Julian Schrittwieser, Ioannis Antonoglou, Rohan Anil, Sebastian Borgeaud, Andrew M. Dai, Katie Millican, Ethan Dyer, Mia Glaese, Thibault Sotti… — Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context
  23. Stephen E. Robertson and Hugo Zaragoza. 2009. The probabilistic relevance framework: BM25 and beyond. Found. Trends Inf. Retr., pages 333–389. find
  24. Jianlin Su, Murtadha H. M. Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. 2024. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, 568:127063.
  25. Hanlin Tang, Yang Lin, Jing Lin, Qingsen Han, Shikuan Hong, Yiwu Yao, and Gongyi Wang. 2024. Razorattention: Efficient KV cache compression through retrieval heads. CoRR, abs/2407.15891. — [arXiv:2407.15891](https://arxiv.org/abs/2407.15891) find
  26. Harsh Trivedi, Niranjan Balasubramanian, Tushar Khot, and Ashish Sabharwal. 2022. 9835 musique: Multi-hop questions via single-hop question composition. Trans. Assoc. Comput. Linguistics, 10. find
  27. Minzheng Wang, Longze Chen, Fu Cheng, Shengyi Liao, Xinghua Zhang, Bingli Wu, Haiyang Yu, Nan Xu, Lei Zhang, Run Luo, Yunshui Li, Min Yang, Fei Huang, and Yongbin Li. 2024. Leave no document behind: Benchmarking long-context llms with extended multi-doc QA. In Proceedings of the 2024
  28. Conference on Empirical Methods in Natural Language Processing, EMNLP 2024, Miami, FL, USA, November 12-16, 2024, pages 5627–5646. Association for Computational Linguistics. find
  29. Yifei Wang, Feng Xiong, Yong Wang, Linjing Li, Xiangxiang Chu, and Daniel Dajun Zeng. 2025. Position bias mitigates position bias: Mitigate position bias through inter-position knowledge distillation. arXiv preprint arXiv:2508.15709.
  30. Wenhao Wu, Yizhong Wang, Guangxuan Xiao, Hao Peng, and Yao Fu. 2024. Retrieval head mechanistically explains long-context factuality. CoRR, abs/2404.15574.
  31. Guangxuan Xiao, Jiaming Tang, Jingwei Zuo, Junxian Guo, Shang Yang, Haotian Tang, Yao Fu, and Song Han. 2024. Duoattention: Efficient long-context LLM inference with retrieval and streaming heads. CoRR, abs/2410.10819.
  32. Wenhan Xiong, Jingyu Liu, Igor Molybog, Hejia Zhang, Prajjwal Bhargava, Rui Hou, Louis Martin, Rashi Rungta, Karthik Abinav Sankararaman, Barlas Oguz, Madian Khabsa, Han Fang, Yashar Mehdad, Sharan Narang, Kshitiz Malik, Angela Fan, Shruti Bhosale, Sergey Edunov, Mike Lewis, Sinong Wang, and Hao Ma.… — Effective Long-Context Scaling of Foundation Models
  33. Qingru Zhang, Chandan Singh, Liyuan Liu, Xiaodong Liu, Bin Yu, Jianfeng Gao, and Tuo Zhao. 2024a. Tell your model where to attend: Post-hoc attention steering for llms. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net. find
  34. Qingru Zhang, Xiaodong Yu, Chandan Singh, Xiaodong Liu, Liyuan Liu, Jianfeng Gao, Tuo Zhao, Dan Roth, and Hao Cheng. 2024b. Model tells itself where to attend: Faithfulness meets automatic attention steering. CoRR, abs/2409.10790. find
  35. Xinrong Zhang, Yingfa Chen, Shengding Hu, Zihang Xu, Junhao Chen, Moo Khai Hao, Xu Han, Zhen Leng Thai, Shuo Wang, Zhiyuan Liu, and Maosong Sun. 2024c. ∞bench: Extending long context evaluation beyond 100k tokens. CoRR, abs/2402.13718. — [arXiv:2402.13718](https://arxiv.org/abs/2402.13718) find
  36. Andrew Zhu, Alyssa Hwang, Liam Dugan, and Chris Callison-Burch. 2024. Fanoutqa: A multi-hop, multi-document question answering benchmark for large language models. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics, ACL 2024 - Short Papers, Bangkok, Thailand,…
  37. An Yang, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chengyuan Li, Dayiheng Liu, Fei Huang, Haoran Wei, Huan Lin, Jian Yang, Jianhong Tu, Jianwei Zhang, Jianxin Yang, Jiaxi Yang, Jingren Zhou, Junyang Lin, Kai Dang, Keming Lu, Keqin Bao, Kexin Yang, Le Yu, Mei Li, Mingfeng Xue, Pei… find
  38. Zhilin Yang, Peng Qi, Saizheng Zhang, Yoshua Bengio, William W. Cohen, Ruslan Salakhutdinov, and Christopher D. Manning. 2018. Hotpotqa: A dataset for diverse, explainable multi-hop question answering. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, Brussel…
  39. Howard Yen, Tianyu Gao, and Danqi Chen. 2024. Long-context language modeling with parallel context encoding. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), ACL 2024, Bangkok, Thailand, August 11-16, 2024, pages 2588–2610. Associati…

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/peng-2025-cafe-retrieval-head-based-coarse/