BERAG: Bayesian Ensemble Retrieval-Augmented Generation for Knowledge-based Visual Question Answering

arXiv cs.CL / 4/27/2026

📰 NewsDeveloper Stack & InfrastructureModels & Research

Key Points

  • The paper argues that standard RAG for question answering—by concatenating all retrieved documents into one context—makes individual-document contributions hard to trace and worsens the “lost-in-the-middle” problem, especially with long contexts and visual data.
  • It introduces BERAG, a Bayesian Ensemble RAG framework that conditions a language model on each retrieved document separately and uses Bayesian posterior probabilities as ensemble weights, updated token-by-token during generation.
  • The method supports probabilistic re-ranking, parallel memory usage, and clearer attribution of how each document influenced the final answer, which is advantageous for large document collections.
  • Experiments on knowledge-based visual question answering show substantial improvements over standard RAG, including gains on Document VQA and multimodal “needle-in-a-haystack” benchmarks, and BERAG is shown to mitigate lost-in-the-middle.
  • The authors also report practical mechanisms: using document posteriors to detect insufficient grounding and trigger “deflection,” and using document pruning to speed up decoding versus standard RAG.

Abstract

A common approach to question answering with retrieval-augmented generation (RAG) is to concatenate documents into a single context and pass it to a language model to generate an answer. While simple, this strategy can obscure the contribution of individual documents, making attribution difficult and contributing to the ``lost-in-the-middle'' effect, where relevant information in long contexts is overlooked. Concatenation also scales poorly: computational cost grows quadratically with context length, a problem that becomes especially severe when the context includes visual data, as in visual question answering. Attempts to mitigate these issues by limiting context length can further restrict performance by preventing models from benefiting from the improved recall offered by deeper retrieval. We propose Bayesian Ensemble Retrieval-Augmented Generation (BERAG), along with Bayesian Ensemble Fine-Tuning (BEFT), as a RAG framework in which language models are conditioned on individual retrieved documents rather than a single combined context. BERAG treats document posterior probabilities as ensemble weights and updates them token by token using Bayes' rule during generation. This approach enables probabilistic re-ranking, parallel memory usage, and clear attribution of document contribution, making it well-suited for large document collections. We evaluate BERAG and BEFT primarily on knowledge-based visual question answering tasks, where models must reason over long, imperfect retrieval lists. The results show substantial improvements over standard RAG, including strong gains on Document Visual Question Answering and multimodal needle-in-a-haystack benchmarks. We also demonstrate that BERAG mitigates the ``lost-in-the-middle'' effect. The document posterior can be used to detect insufficient grounding and trigger deflection, while document pruning enables faster decoding than standard RAG.