Work

ADR: Attention Head Detection and Reweighting Enhance RAG Performance in a Positional-Encoding-Free Paradigm

Year: 2025

Type: article

Abstract: Retrieval-augmented generation (RAG) has established a new search paradigm, in which large language models integrate external resources to compensate for their inherent knowledge limitations. However, limited context awareness reduces the performance of large language models in RAG tasks. Existing solutions incur additional time and memory overhead and depend on specific positional encodings. In this paper, we propose Attention Head Detection and Reweighting (ADR), a lightweight and general framework. Specifically, we employ a recognition task to identify RAG-suppressing heads that limit the model’s context awareness. We then reweight their outputs with learned coefficients to mitigate the influence of these RAG-suppressing heads. After training, the weights are fixed during inference, introducing no additional time overhead and remaining agnostic to the choice of positional embedding. Experiments on PetroAI further demonstrate that ADR enhances the context awareness of fine-tuned models.

Source: Information

DOI: 10.3390/info16100900

Authors: Mingwei Wang, Xiaobo Li, Qian Zeng, Xingbang Liu, Minghao Yang, Zhichen Jia

Institutions: Research Institute of Petroleum Exploration and Development, China National Petroleum Corporation (China)

Language: English

Topic: Topic Modeling

Subfield: Artificial Intelligence

Field: Computer Science

Domain: Physical Sciences

Open Access status: gold

Estimated APC paid: $1,930

Funders: China National Petroleum Corporation, National Key Research and Development Program of China