Retrieval-Aware Contrastive Decoding Rule

• The paper introduces a novel retrieval-aware contrastive decoding rule that integrates multiple document-specific experts to enhance cross-document reasoning for RAG systems.
• It employs parallel decoding with individualized KV caches, calibrating expert outputs against a model prior and leveraging retrieval relevance for token selection.
• Empirical evaluations reveal significant accuracy improvements and up to 180× faster time-to-first-token compared to traditional long-prompt attention methods.

Parallel Context-of-Experts Decoding (Pced) is a training-free, decode-time framework for Retrieval-Augmented Generation (RAG) that enables efficient multi-document reasoning by aggregating evidence at the decoding step rather than through the standard attention mechanism. Unlike traditional long-prompt approaches, which concatenate retrieved documents into the input and incur quadratic costs, Pced isolates each retrieved document as an “expert” with its own context, synchronizing expert predictions via a novel retrieval-aware contrastive decoding rule. This design recovers cross-document reasoning capabilities without the need to construct a shared attention context, offering robust performance and significant latency improvements in practical RAG systems (Corallo et al., 13 Jan 2026).

1. Operational Workflow and Framework Overview

Pced replaces long prompt concatenation with a set of document-specific experts, each with its own precomputed key-value (KV) cache from the underlying LLM (LM). The framework operates as follows:

• Offline Preparation:
  • Construct a document datastore $\mathcal{D}$. For each document $d_i$, store (a) its retrieval embedding $e_i$ and (b) its precomputed LM KV cache $K_i$.
  • At query time:
  • Retrieve and rerank top-$N$ documents, mapping their retrieval and reranker scores into normalized relevance $r_k \in (0,1)$ using the harmonic mean.
  • Instantiate $N$ “contextual experts” with caches $\{K_1, ..., K_N\}$ and one “amateur” expert representing the model prior with $K_0 = \emptyset$.
• Parallel Encoding:
  • All $N+1$ experts process the query prefix $q$ and shared generation history in parallel, yielding per-expert logits $s_k \in \mathbb{R}^{|V|}$ at each step.
• Retrieval-Aware Contrastive Decoding:
  • Calibrate each expert $k$ against the amateur expert and inject the retrieval prior to produce adjusted scores $\hat{s}_k$.
  • Select the next token $y_t$ by maximizing $\max_k \hat{s}_k(v)$ over vocabulary entries $v$.
  • Append $y_t$ to the shared generation history for all experts and update each expert’s KV cache.
• Evidence Stitching:
  • By always feeding the latest prefix and chosen token to all caches, Pced enables amortized crossing between expert contexts without constructing a joint attention mechanism.

2. Mathematical Formulation

The retrieval-aware contrastive decoding in Pced is mathematically defined as follows:

• Notation:
  • Vocabulary $V$
  • Generation step $t$
  • Expert $k \in \{0, ..., N\}$
  • Raw logits: $s_k(v) \in \mathbb{R}$, where $v \in V$
  • Amateur (prior) logits: $s_0(v)$
  • Contextual expert logits: $s_k(v)$ for $k=1 \dots N$
  • Relevance prior: $r_k \in (0,1)$
  • Contrastive strength: $\beta_0 \geq 0$
  • Retrieval-gating weight: $\gamma > 0$ (empirical $\gamma=2.5$)
• Calibrated scores:

$$\hat{s}_k(v) = (1+\beta_0)\,s_k(v) - \beta_0\,s_0(v) + \gamma\,\log r_k$$

Here $(1+\beta_0)\,s_k(v) - \beta_0\,s_0(v)$ represents the contrastive adjustment, and $\gamma\,\log r_k$ introduces a bias toward higher-relevance documents.

• Token Selection Rule:

$$y_t = \arg\max_{v \in V}\;\max_{k \in \{1,\dots,N\}} \hat{s}_k(v)$$

At each decoding step, the rule selects the token and expert combination that maximizes the calibrated score.

3. Algorithmic Details

The main generation loop for Pced is summarized in the following pseudocode:

Inputs: query q precomputed caches {K₁,…,K_N}, relevance scores {r₁,…,r_N} LLM LM (logit access) contrastive weight schedule β₀ retrieval weight γ Initialize: caches ← {K₀=∅} ∪ {K₁,…,K_N} prefix_tokens ← tokenize(q) β₀ ← undefined generation_history ← prefix_tokens For t = 1 to T_max: # Batch forward pass across experts For k in 0…N: sₖ ← LM.forward(caches[k], generation_history) # Set β₀ dynamically (first step) If t == 1: β₀ ← f_JS-divergence(s₀, s₁…s_N) # Contrastive + retrieval calibration For k = 1…N: For v in V: ĥₖ(v) ← (1+β₀)·sₖ(v) - β₀·s₀(v) + γ·log(rₖ) # Expert selection and token emission (k*, v*) ← argmax over k=1…N, v ∈ V of ĥₖ(v) y_t ← v* # Update history and KV caches Append y_t to generation_history For k in 0…N: caches[k] ← LM.update_cache(caches[k], y_t) If y_t is end-of-sequence: break Output: detokenize(generation_history without the query)

Key operational principles include amortized batch processing across experts, individualized KV cache management, and shared generation history.

4. Computational Complexity and Latency Analysis

Pced offers notable improvements relative to traditional and alternative multi-document fusion schemes:

Scheme	Prefill/TTFT Complexity	Interaction Mechanism
Long-Prompt Attention	$O((N \cdot L)^2)$	Native cross-attention
Separate-KV + Merge	$O(N\cdot L\cdot d)$ (offline), moderate alignment cost (merge)	Requires attention recomputation
Pced	Single batch call for $N+1$ caches, linear per-token cost $O(N\cdot L\cdot d^2)$	Decoding-time expert fusion

Time-to-first-token (TTFT) with Pced demonstrates up to 180× speedup (0.14 s vs. 25.5 s at $K=90$ documents), with end-to-end latency roughly 1.7× faster on 65k-token contexts and 512-token generation. Memory requirements are $O(N\cdot L\cdot d)$ for caches, making Pced well-suited to static, read-heavy corpora.

5. Empirical Evaluation on Multi-Document Reasoning Tasks

Experimental results on LOFT (RAG/ICL) and LongBench benchmarks show substantial accuracy gains over approaches such as APE (Attention Pooling Experts):

• LOFT RAG (Mistral-13B-Instruct):
  • HotpotQA: APE 27 → Pced-Dense 66
  • MuSiQue: APE 11 → Pced-Dense 34
  • NQ: APE 38 → Pced-Dense 81
• LOFT ICL (Mistral-13B-Instruct):
  • Web: APE 58.9 → Pced-Dense 62.2
  • Date: APE 40.0 → Pced-Dense 57.8
• LongBench (Qwen3-8B):
  • Multi-Doc QA (Hotpot): 56.3 → 62.6
  • Few-Shot (TriviaQA): 84.0 → 88.8
  • RepoB-P: 51.1 → 60.1

Ablation studies demonstrate that both the contrastive term ($\beta_0$) and retrieval prior ($\gamma$) are critical; removal causes severe drops in performance (e.g., NQ falls from 85 to 52 without retrieval gating). Max-selection aggregation among experts is essential for multi-hop reasoning (HotpotQA: max 64 vs. mixture 56).

6. Limitations, Extensions, and Unresolved Questions

Pced is subject to several constraints:

• Model Logit Access: Full per-token logits are required; closed-API backends that only return sampled tokens are incompatible.
• Retrieval Quality Sensitivity: The framework is reliant on effective retrieval and reranking; missing or mis-ranked documents render their evidence inaccessible.
• Storage Costs: Linear scaling with the number of documents and hidden-state dimension (e.g., 11 GB for 1 200 passages with Llama-3.1-8B and FP16 caches).

Potential avenues for extension include end-to-end training for expert selection within LLMs (reducing dependency on external retrievers/rerankers), instance-wise dynamic adjustment of $\gamma$ and aggregation rules, and hybrid approaches combining cross-attention bridges with decode-time fusion.

Pced presents a clear trade-off: it preserves the cross-document reasoning power characteristic of long-prompt attention while delivering significant improvements in decoding speed and robustness to noise or large candidate pool sizes (Corallo et al., 13 Jan 2026).