PaperlayerView exact revision on arXiv

Source-linked AI summary

Parallel Context-of-Experts Decoding for Retrieval Augmented Generation

Giulio Corallo, Paolo Papotti

arXiv:2601.08670v1cs.AIcs.CL

TL;DR

RAG must reconcile the reasoning benefits of concatenated documents with their prefill costs and the interaction loss of independently cached documents. PCED shifts aggregation to decoding by synchronizing per-document expert logits with a retrieval-aware contrastive rule. It recovers cross-document reasoning, often matches or exceeds full-context baselines, and substantially reduces time-to-first-token.

  • Problem

    RAG concatenation creates prefill bottlenecks and long-context reasoning failures, while independent document KV caches can degrade cross-document reasoning.

  • Method

    PCED is a training-free framework that treats each retrieved document as an expert and fuses its logits with an amateur prior using retrieval-aware contrastive decoding.

  • Results

    PCED consistently outperforms APE on multi-document reasoning, often matches or exceeds full-context concatenation, and achieves over 180× faster TTFT with 0.14s versus 25.50s.

  • Takeaways & Limitations

    PCED preserves modular per-document KV caches while recovering cross-document reasoning and providing more distractor-robust, scalable generation.

  • Takeaways & Limitations

    PCED requires full access to per-expert token-level logits, restricting direct application to closed-source or API-only models exposing limited outputs.

Abstract

from arXiv · show

Retrieval Augmented Generation faces a trade-off: concatenating documents in a long prompt enables multi-document reasoning but creates prefill bottlenecks, while encoding document KV caches separately offers speed but breaks cross-document interaction. We propose Parallel Context-of-Experts Decoding (Pced), a training-free framework that shifts evidence aggregation from the attention mechanism to the decoding. Pced treats retrieved documents as isolated "experts", synchronizing their predictions via a novel retrieval-aware contrastive decoding rule that weighs expert logits against the model prior. This approach recovers cross-document reasoning capabilities without constructing a shared attention across documents.

1 Introduction

RAG must balance long-context multi-document reasoning against prefill cost and the reasoning failures caused by scattered evidence. PCED addresses this trade-off by aggregating isolated document experts during decoding rather than constructing joint attention.

  • Long concatenated RAG prompts make inference dominated by prefill latency and can hinder integration of evidence spread across documents.
  • Independent KV-cache encoding reduces prefill cost but can substantially degrade multi-hop and reasoning-intensive performance by removing cross-document attention.
  • PCED is a training-free framework that shifts document aggregation from attention to decoding.
  • Each retrieved document acts as an expert, while a no-context amateur prior helps select tokens according to retrieval support.
  • PCED uses token-level expert switching and retrieval-integrated priors to aggregate evidence without shared attention, outperforming prior parallel methods by up to 70 points and delivering over 180× faster time-to-first-token.

2 Related Work

Related work addresses parallel KV caching, recovery of cross-document interaction, and context-aware decoding. PCED preserves per-document modularity while targeting cross-document reasoning during inference.

  • Prior parallel-encoding methods reduce prefill cost with independently cached documents, but differ in whether they use masking, fine-tuning, or decoupled encoding-generation interfaces.
  • PCED studies evidence aggregation across multiple cached documents while preserving their independent modularity.
  • Cache-merging methods restore cross-document attention through selective recomputation, bridging tokens, or training-free alignment after independent encoding.
  • Context-aware decoding shifts probability toward context-supported tokens, but most formulations assume one supportive context.
  • DvD extends context-aware decoding to multiple documents but collapses them into a single input sequence.

3 Methodology

PCED prepares independently retrievable document KV caches, processes document experts and an amateur prior in parallel, and fuses their logits at each decoding step. Retrieval and reranking scores become priors that gate expert influence.

  • Offline KV cache preparation: The datastore stores each document’s embedding for retrieval alongside its precomputed KV cache.
  • Retrieval and relevance scoring: For each query, PCED retrieves top-N documents, reranks them, and fuses recall-oriented retrieval scores with precision-oriented reranker scores into per-document relevance scores.
  • Parallel Context-of-Experts: PCED runs one amateur prior and one contextual expert per retrieved document in a single batched forward pass.
  • Retrieval-aware contrastive decoding: Retrieval-aware contrastive decoding calibrates each contextual expert against the amateur logits and applies retrieval gating.
  • Decoding: At every step, PCED selects the highest-scoring candidate across experts and appends the chosen token to all experts’ shared generation histories.
  • Experimental variants: The methodology evaluates PCED variants using Sparse, Dense, or ColBERT relevance signals against concatenation, KV-merging, and MapReduce baselines.

4 Experimental Setup

The experiments compare how context is incorporated while fixing the underlying models, prompts, and retrieved candidates. They cover RAG, ICL, long-context QA with distractors, and multiple relevance-scoring variants.

  • Tasks: PCED is tested on RAG, in-context learning, and long-context question answering with distractors.
  • Controlled comparison: All methods use the same LLMs, prompts, and retrieved candidates, varying only the mechanism of context incorporation.
  • Datasets and metrics: LOFT experiments retrieve a fixed top-90 document pool, use Subspan Exact Match for RAG, Exact Match for ICL, and official metrics for LongBench.
  • Datasets and metrics: Long-context evaluation uses query-focused LongBench subsets and tests distractor robustness by adding two uniformly sampled distractors to the gold document.
  • PCED variants: The three PCED variants differ only in the relevance signal used to weight experts, while the retrieved document set remains identical.
  • Baselines: Baselines include single- or all-document concatenation, APE KV-cache merging, and MapReduce per-document summarization followed by aggregation.

5 Results and Discussion

PCED recovers cross-document reasoning and improves robustness to distractors while reducing latency through parallel KV-cache conditioning and decode-time expert aggregation. Its ablations show that both retrieval priors and contrastive calibration contribute to performance.

  • Cross-Document Reasoning Emerges at Decode Time: PCED outperforms APE on multi-document QA and ICL, improving QAMPARI from 7 to 77 and exceeding MAPREDUCE by up to 23 points.PCED-DENSE also outperforms Corpus in Ctx (All) in 11/16 settings.
  • Cross-Document Reasoning Emerges at Decode Time: Token-level expert switching stitches multi-hop evidence by moving from a bridging-entity document to a second document while sharing generation history.MAPREDUCE remains superior in some settings but requires multiple LLM calls, whereas PCED aggregates evidence in one decoding procedure.
  • Less Noise, More Accuracy: PCED improves single-document retrieval tasks under distractors by isolating documents and downweighting irrelevant experts through retrieval-aware contrastive decoding.On NQ, PCED-DENSE reaches 85 with Llama and 81 with Mistral, above their listed concatenation baselines.
  • Efficiency at Scale: PCED achieves over 180× faster TTFT, reducing latency from 25.50s to 0.14s, and cuts end-to-end latency by approximately 1.7× at 65k context.The comparison uses continuous batching and PagedAttention for both methods.
  • Ablations: Removing either the retrieval prior or contrastive calibration causes large accuracy drops, while Max aggregation best supports multi-hop expert switching.Soft mixtures can help in single-document settings.

6 Conclusion

PCED is a training-free, cache-native decoding framework for efficient multi-document reasoning. It replaces long-context attention with retrieval-aware expert logit fusion, preserving modular KV caches while recovering cross-document reasoning and improving distractor robustness.

  • 6 Conclusion: PCED enables efficient multi-document reasoning under parallel, cache-native conditioning.It is presented as a training-free decoding framework.
  • 6 Conclusion: PCED preserves KV-cache modularity while recovering cross-document reasoning through retrieval-aware expert logit fusion at decode time.The approach replaces long-context attention rather than constructing a shared attention context.
  • 6 Conclusion: PCED matches or surpasses full-context baselines and is more robust to distractors.The conclusion describes document scaling as flexible with batch size rather than limited by the training context window.

Limitations

PCED’s limitations concern model access, retrieval quality, and storage requirements. These constraints bound applicability to open or self-hosted models, retrieved evidence quality, and static, read-heavy corpora.

  • Model access: PCED requires full access to per-expert token-level logits, restricting direct use with closed-source or API-only language models.Such models may expose only sampled tokens or limited candidate log-probabilities.
  • Retrieval quality: PCED cannot recover evidence that retrieval omits, and low relevance scores can underweight or eliminate otherwise relevant experts.Retrieval-aware decoding mitigates noise from weak documents but does not restore absent evidence.
  • Storage-computation trade-offs: Precomputed KV caches impose storage costs that scale linearly with corpus size and hidden-state dimensionality.FP16 caches for the LOFT HOTPOTQA corpus require approximately 11.04 GB with LLAMA-3.1-8B.
  • Storage-computation trade-offs: PCED is best suited to read-heavy, write-rare workloads with static corpora, where storage costs are justified by lower query-time latency.The stated deployment example is an enterprise knowledge base.

A Evaluation Setup

The evaluation fixes prompts, dataset fields, and random seeds across methods while varying context incorporation. It covers LOFT and LongBench tasks, prompt normalization, score processing, and a controlled latency benchmark.

  • Evaluation protocol: All compared methods use fixed dataset fields, prompts, and answer prefixes, varying only the mechanism of context incorporation.Experiments use random seed 42, and reported numbers are generally from one deterministic run per method.
  • LOFT setup: LOFT evaluation uses standardized system, context, question, and answer-format fields for RAG and ICL tasks.RAG uses the Figure 4 configuration, while ICL enforces a strict output format defined in Figure 5.
  • PCED setup: PCED instantiates separate contextual experts from dataset prompts and documents, plus an amateur expert using only the system prompt.All experts share the same question template while conditioning on separate KV caches.
  • Latency benchmark: Latency evaluation uses 64 fixed-length documents, one containing a secret code, with exact 2048-token document budgets.The query requires verbatim secret-code generation under controlled context length.
  • Score processing: Retrieval and reranker relevance signals are normalized to a common range before combining them with a harmonic mean.Values are clipped below one and away from zero before entering log r_k.

C.1 Impact of Contrastive Strength (β)

Contrastive strength controls how strongly expert predictions are sharpened against the amateur prior, while retrieval-prior weight controls the influence of relevance scores. The ablations favor dynamic contrastive strength and a moderate retrieval-prior weight over fixed or extreme settings.

  • Contrastive strength: β = 0 generally degrades performance relative to the best contrastive settings, supporting a positive contrastive component.The ablation interprets amateur-logit subtraction as isolating knowledge supplied by retrieved documents.
  • Contrastive strength: Fixed β values are inconsistent across tasks: β = 0.25 performs well on Llama-NQ, whereas β = 0.75 performs well on Llama-HotpotQA but drops on Mistral-HotpotQA.Task-specific optima do not transfer reliably across datasets and models.
  • Contrastive strength: Dynamic β remains competitive across all evaluated models and datasets without per-task tuning.The authors select Dynamic as the default for stability across diverse retrieval scenarios.
  • Retrieval-prior weight: γ controls retrieval-score influence through γ log r_k, with evaluated values spanning 0.5 to 4.0 around the default γ = 2.5.The sweep compares under-weighting, over-weighting, and moderate gating regimes.
  • Retrieval-prior weight: γ = 0.5 lowers Llama-NQ to 75, while γ = 4.0 reaches 87 on Llama-NQ but reduces Llama-HotpotQA to 64 versus 66.These results illustrate degradation from under-weighting and inconsistency from excessive retrieval gating.

C.3 Contrastive Signal vs. Retrieval Score Only

The ablations separate the retrieval prior from the contrastive signal and compare token-wise Max aggregation with probabilistic alternatives. Retrieval guidance is essential, while the contrastive signal benefits models differently and Max aggregation is strongest for multi-hop reasoning.

  • Component Analysis: Retrieval guidance is foundational: removing it causes the Only Contrastive setting to fail across all benchmarks, with Llama NQ dropping to 52.
  • Component Analysis: For LLAMA-3.1, removing the contrastive signal lowers NQ from 85 to 70, whereas MISTRAL remains more robust with retrieval scores alone.
  • Component Analysis: The full PCED framework achieves the highest absolute scores across the reported model cases.
  • Aggregation Rule Ablation: On HOTPOTQA, Max aggregation beats MoE by 8 points, 64 versus 56, supporting sharper token-level expert switching for multi-hop reasoning.

C.5 Robustness to Candidate Pool Size (k)

PCED remains stable as the retrieved candidate pool grows from k = 8 to k = 128. This tolerance allows higher recall for difficult queries without sacrificing generation quality.

  • Performance remains nearly constant despite a 16× increase in experts from k = 8 to k = 128.
  • NQ stays near 85, while HOTPOTQA varies only marginally between 63 and 65 across candidate-pool sizes.
  • The retrieval prior γ log r_k gates low-relevance experts and prevents distractor accumulation as the pool expands.
  • At k = 128, PCED preserves generation quality while allowing users to maximize recall for difficult queries.
Loading 2601.08670v1…