Parallel Key-Value Cache Fusion for Position Invariant RAG

TL;DR

This work tackles the Lost in the Middle problem in retrieval-augmented generation by introducing KV-Fusion, a two-decoder framework that preloads shared key-value caches from retrieved passages and trains a decoder to generate answers conditioned on these caches. By injecting uniform positional information and grounding generation in reshaped KV caches, KV-Fusion achieves position-invariant outputs and strong robustness to irrelevant passages across Natural Questions, TriviaQA, and POPQA. The approach reduces reliance on reranking and demonstrates near-perfect token-level consistency under context-order perturbations, marking a notable advance for decoder-only RAG systems. Practically, KV-Fusion enables reliable, order-agnostic QA in open domains with large context sets and enhances robustness against context noise in retrieval pipelines.

Abstract

Recent advancements in Large Language Models (LLMs) underscore the necessity of Retrieval Augmented Generation (RAG) to leverage external information. However, LLMs are sensitive to the position of relevant information within contexts and tend to generate incorrect responses when such information is placed in the middle, known as `Lost in the Middle' phenomenon. In this paper, we introduce a framework that generates consistent outputs for decoder-only models, irrespective of the input context order. Experimental results for three open domain question answering tasks demonstrate position invariance, where the model is not sensitive to input context order, and superior robustness to irrelevent passages compared to prevailing approaches for RAG pipelines.

Figures (6)

• Figure 1: Illustration of the KV-Fusion model: Generated tokens remain consistent even when the retrieved passages are shuffled.
• Figure 2: Overview of KV-Fusion Architecture. $\mathcal{D}_\textrm{p}$ denotes Prefill decoder and $\mathcal{D}_\textrm{t}$ represents Trainable decoder. We employ the off-the-shelf LLM to extract the key and value states of the retrieved contexts independently. Then reshaping these caches to train the LLM with task instructions along with questions to generate answers.
• Figure 3: Comparison of EM Accuracy between KV-Llama3 and Llama3 across different gold context positions. KV-Llama3 maintains its accuracy, while Llama3 shows a tendency for the 'lost in the middle' problem.
• Figure 4: Accuracies of baseline and KV models in two scenarios: 1) POS1, where the gold context is positioned first, and 2) Shuffled, where contexts are randomly ordered. KV models maintain their accuracy on both cases, while baseline models struggle in shuffled setting, leading in a wider accuracy gap between the baseline and KV-models.
• Figure 5: Comparison of EM Accuracy between KV-Llama3.1 and Llama3.1 across different gold context positions. With varying gold context positions, KV-Llama3.1 illustrates consistent accuracies across datasets. However, Llama3.1 suffers from the 'Lost in the middle' problem, which can be resolved by KV-Fusion models.