PEAR: Position-Embedding-Agnostic Attention Re-weighting Enhances Retrieval-Augmented Generation with Zero Inference Overhead

TL;DR

This paper proposes Position-Embedding-Agnostic attention Re-weighting (PEAR), which enhances the context awareness of LLMs with zero inference overhead, while outperforming competitive baselines in accuracy and efficiency across various RAG tasks.

Abstract

Large language models (LLMs) enhanced with retrieval-augmented generation (RAG) have introduced a new paradigm for web search. However, the limited context awareness of LLMs degrades their performance on RAG tasks. Existing methods to enhance context awareness are often inefficient, incurring time or memory overhead during inference, and many are tailored to specific position embeddings. In this paper, we propose Position-Embedding-Agnostic attention Re-weighting (PEAR), which enhances the context awareness of LLMs with zero inference overhead. Specifically, on a proxy task focused on context copying, we first detect heads which suppress the models' context awareness thereby diminishing RAG performance. To weaken the impact of these heads, we re-weight their outputs with learnable coefficients. The LLM (with frozen parameters) is optimized by adjusting these coefficients to minimize loss on the proxy task. As a result, the coefficients are optimized to values less than one, thereby reducing their tendency to suppress RAG performance. During inference, the optimized coefficients are fixed to re-weight these heads, regardless of the specific task at hand. Our proposed PEAR offers two major advantages over previous approaches: (1) It introduces zero additional inference overhead in terms of memory usage or inference time, while outperforming competitive baselines in accuracy and efficiency across various RAG tasks. (2) It is independent of position embedding algorithms, ensuring broader applicability.

Figures (7)

• Figure 1: An illustration of causal mediation methods for circuit discovery.
• Figure 2: Suppose in layer $l$, $A^{(l,h)}$ is discovered as a RAG-suppression head. PEAR re-weights its output with a learnable scalar $\tau^{(l,h)}$.
• Figure 3: Heatmaps of $\Delta\pi$ scores for each head across three LLMs ($n=10$).
• Figure 4: The learned coefficients of PEAR (on Llama2-7B-chat and $K=40$).
• Figure 5: Heatmaps of $\Delta\pi$ scores for each head of llama2-7B-chat ($n=15, n=25, n=50$).