Beyond RAG vs. Long-Context: Learning Distraction-Aware Retrieval for Efficient Knowledge Grounding
LDAR introduces a distraction-aware retrieval method using similarity distribution bands, improving knowledge grounding with 5%+ accuracy gains and reduced token usage.
Key Findings
Methodology
LDAR employs a similarity distribution-based band retrieval strategy, integrating a Transformer encoder to predict optimal similarity intervals (qL, qU). This adaptive mechanism dynamically balances information coverage and distraction, guided by maximizing the probability of correct answers. The retriever is lightweight, relying solely on cosine similarity scores, and is trained via reinforcement learning to minimize interference from distractors. It does not require fine-tuning the large language model, making it scalable and cost-effective. Extensive experiments across multiple architectures and six knowledge-intensive benchmarks demonstrate that LDAR surpasses fixed retrieval strategies, especially under limited model capacity, by effectively reducing irrelevant information while maintaining high coverage.
Key Results
- Across six datasets, LDAR improves accuracy by an average of 5.2%, with performance exceeding traditional long-context methods by over 10% at 128K token length scenarios. It reduces retrieval tokens by 30%, showing efficiency gains. In models like GPT-4o and Qwen-2.5-7B, LDAR maintains robustness even with high distractor presence, outperforming baseline retrieval strategies.
- Ablation studies confirm that band-based retrieval better manages the trade-off between coverage and distraction compared to random or full retrieval. The method adapts to model capacity, retrieving fewer passages for open-source models and more for closed-source, aligning retrieval scope with model ability.
- In multi-task evaluations, LDAR consistently enhances answer correctness, demonstrating its potential for knowledge-intensive applications. Its ability to dynamically tune retrieval scope based on similarity distribution makes it versatile for real-world deployment.
Significance
This work advances the state-of-the-art in knowledge grounding by addressing the limitations of fixed long-context or retrieval-based methods. LDAR’s adaptive, learning-based approach effectively mitigates distraction, reduces computational costs, and enhances model robustness. It offers a scalable solution for integrating external knowledge into large models, facilitating more accurate and efficient AI systems in tasks like question answering, reasoning, and multi-modal understanding. Its capacity to adapt to different model capabilities paves the way for broader deployment in industry and academia, enabling smarter, more reliable AI assistants.
Technical Contribution
LDAR introduces a novel similarity distribution band retrieval mechanism, leveraging a Transformer encoder to predict adaptive similarity intervals. This approach transforms the complex subset selection problem into a low-dimensional, smooth control space, enabling efficient learning and generalization. The training maximizes the likelihood of correct answers conditioned on retrieved passages, using reinforcement learning to optimize the retrieval policy. Unlike fixed or heuristic methods, LDAR’s capacity-aware, self-adaptive retrieval significantly reduces irrelevant information and enhances answer accuracy, all without fine-tuning the large language model itself.
Novelty
This is the first work to formulate a similarity distribution-based band retrieval strategy that dynamically adjusts retrieval scope according to model capacity. It departs from prior heuristic or static retrieval methods, offering a principled, learnable mechanism that balances coverage and distraction. The integration of Transformer-based prediction of similarity intervals and reinforcement learning for policy optimization represents a significant innovation, enabling models to adaptively manage information flow in long-context scenarios.
Limitations
- The approach relies on the quality of pre-trained embeddings; in noisy or out-of-distribution data, similarity measures may mislead the retrieval policy, reducing effectiveness.
- While reducing distraction, the method may still struggle in extremely dense or ambiguous information environments, requiring further contextual understanding.
- Training involves reinforcement learning, which can be sample-inefficient and computationally intensive, especially on very large datasets or in real-time applications.
Future Work
Future directions include integrating multi-modal signals to enhance similarity estimation, extending the framework to multi-turn dialogue and multi-modal tasks, and exploring unsupervised or semi-supervised training to reduce reliance on labeled data. Additionally, combining LDAR with continual learning techniques could enable models to adapt dynamically to evolving knowledge bases, further improving robustness and scalability in real-world applications.
AI Executive Summary
In recent years, large language models (LLMs) like GPT-4 and Gemini 2.5 have demonstrated remarkable capabilities in processing long texts, pushing the boundaries of natural language understanding. Traditionally, to incorporate external knowledge, the Retrieval-Augmented Generation (RAG) framework has been employed, retrieving relevant passages from external corpora to ground model outputs. However, as models grow in capacity, directly feeding full long contexts becomes feasible and sometimes preferable, but this approach introduces new challenges: high token costs, information redundancy, and the notorious ‘lost in the middle’ phenomenon where crucial information in the middle of long inputs is forgotten or overlooked.
This paper proposes LDAR, a novel, learning-based retrieval strategy that dynamically balances information coverage and distraction. Unlike fixed or heuristic retrieval methods, LDAR employs a Transformer encoder to predict an optimal similarity interval (band) from the distribution of cosine similarity scores between the query and candidate passages. This adaptive band selection enables the model to focus on the most relevant information while minimizing distracting, irrelevant passages. The retriever is trained via reinforcement learning, optimizing for the likelihood of the model producing correct answers, and does not require fine-tuning the large language model itself, making it computationally efficient.
Extensive experiments across six knowledge-intensive tasks and multiple model architectures demonstrate that LDAR consistently outperforms fixed retrieval strategies, improving accuracy by over 5% on average and reducing token usage by 30%. Notably, the method adapts to different model capacities, retrieving fewer passages for open-source models and more for closed-source ones, aligning retrieval scope with model ability. These results highlight the potential of adaptive, distraction-aware retrieval to enhance the efficiency and robustness of knowledge-grounded AI systems.
The broader impact of this work lies in its ability to address fundamental bottlenecks in long-text processing, offering a scalable, cost-effective solution for integrating external knowledge. Future work aims to extend LDAR to multi-modal scenarios, incorporate unsupervised learning, and enable continuous knowledge updates, paving the way for smarter, more reliable AI assistants capable of handling complex, real-world tasks with minimal distraction and maximal efficiency.
Deep Dive
Abstract
Retrieval-Augmented Generation (RAG) is a framework for grounding Large Language Models (LLMs) in external, up-to-date information. However, recent advancements in context window size allow LLMs to process inputs of up to 128K tokens or more, offering an alternative strategy: supplying the full document context directly to the model, rather than relying on RAG to retrieve a subset of contexts. Nevertheless, this emerging alternative strategy has notable limitations: (i) it is token-inefficient to handle large and potentially redundant contexts; (ii) it exacerbates the `lost in the middle' phenomenon; and (iii) under limited model capacity, it amplifies distraction, ultimately degrading LLM output quality. In this paper, we propose LDAR (Learning Distraction-Aware Retrieval), an adaptive retriever that learns to retrieve contexts in a way that mitigates interference from distracting passages, thereby achieving significantly higher performance with reduced token usage compared to long-context approaches. Extensive experiments across diverse LLM architectures and six knowledge-intensive benchmarks demonstrate the effectiveness and robustness of our approach, highlighting the importance of balancing the trade-off between information coverage and distraction.