BERAG: Bayesian Ensemble Retrieval-Augmented Generation for Knowledge-based Visual Question Answering
2026-04-24 • Computation and Language
Computation and Language
AI summaryⓘ
The authors highlight a problem with how question-answering models handle multiple documents by just sticking them together, which can hide important info and slow things down. They introduce BERAG, a new method that looks at each document separately and updates how much it trusts each one as it generates answers, making it easier to know which documents helped. This method also runs faster and works better when there are lots of documents or pictures involved. Their experiments show that BERAG improves performance on tasks needing detailed reasoning across many documents, especially in visual question answering.
Retrieval-Augmented GenerationBayes' RuleVisual Question AnsweringDocument AttributionEnsemble MethodsProbabilistic Re-RankingLanguage ModelsContext LengthFine-TuningMultimodal Reasoning
Authors
Jinghong Chen, Jingbiao Mei, Guangyu Yang, Bill Byrne
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.