BEE-RAG: Balanced Entropy Engineering for Retrieval-Augmented Generation

TL;DR

This work tackles the instability of retrieval-augmented generation under long contexts by introducing Balanced Entropy Engineering for RAG (BEE-RAG), which enforces entropy invariance across context lengths via a document-specific balancing entropy factor $\beta$. It combines balanced context entropy (BCE), intrinsic multi-importance inference (IMI) for zero-shot importance estimation, and a parameter-efficient adaptive balancing entropy factor learning mechanism to enable domain adaptation with minimal updates. The theoretical analysis links the balancing factors to entropy stability and demonstrates practical efficacy through zero-shot and lightweight fine-tuning across four QA benchmarks, including multi-hop tasks, with improvements in both accuracy and robustness to retrieval quality. Overall, BEE-RAG provides a principled, scalable solution for leveraging long-context information in RAG while maintaining efficiency and generalization across model scales and datasets.

Abstract

With the rapid advancement of large language models (LLMs), retrieval-augmented generation (RAG) has emerged as a critical approach to supplement the inherent knowledge limitations of LLMs. However, due to the typically large volume of retrieved information, RAG tends to operate with long context lengths. From the perspective of entropy engineering, we identify unconstrained entropy growth and attention dilution due to long retrieval context as significant factors affecting RAG performance. In this paper, we propose the balanced entropy-engineered RAG (BEE-RAG) framework, which improves the adaptability of RAG systems to varying context lengths through the principle of entropy invariance. By leveraging balanced context entropy to reformulate attention dynamics, BEE-RAG separates attention sensitivity from context length, ensuring a stable entropy level. Building upon this, we introduce a zero-shot inference strategy for multi-importance estimation and a parameter-efficient adaptive fine-tuning mechanism to obtain the optimal balancing factor for different settings. Extensive experiments across multiple RAG tasks demonstrate the effectiveness of BEE-RAG.

Figures (4)

• Figure 1: In vanilla RAG, increasing input context length raises the information entropy of attention scores (left), potentially harming performance, while the LLM is less focused on important segments (right). This paper proposes a balanced entropy engineering strategy, which maintains entropy stability for longer contexts and guide the LLM to focus on critical segments.
• Figure 2: The overview of the proposed BEE-RAG. The left side shows the overall architecture design, while the right side illustrates the zero-shot strategy intrinsic multi-importance inference and the parametric-efficient fine-tuning strategy adaptive balancing factor learning.
• Figure 3: Effect of LLM capabilities with various Scales. Mean metrics are reported over four datasets.
• Figure 4: Effect of context length with various document numbers. Mean metrics are reported over four datasets.