Parallel Context-of-Experts Decoding for Retrieval Augmented Generation

🍞 Top Bread (Hook): Imagine you’re doing a school project. You collect facts from many books, but copying every page into your notebook would be slow and messy. You wish you could quickly ask each book what it knows and then write your report, one sentence at a time.

🥬 Filling (The Actual Concept): What it is: Retrieval-Augmented Generation (RAG) is when an AI looks up helpful documents first and then uses them to write better, more factual answers. How it works (recipe):

1. Find documents that seem relevant to your question.
2. Put those documents where the AI can read them.
3. Generate the answer using both the documents and the model’s own knowledge. Why it matters: Without RAG, the AI may guess or forget facts, especially about rare details.

🍞 Bottom Bread (Anchor): When you ask, “Who discovered penicillin and in what year?”, RAG retrieves a page about Alexander Fleming and the year 1928 so the AI can answer exactly.

🍞 Top Bread (Hook): You know how reading a super long wall of text makes you slow and tired? AI feels that too when it has to read a giant prompt.

🥬 Filling (The Actual Concept): What it is: The prefill bottleneck is the slowdown that happens when an AI must first read a very long prompt (many documents squished together) before it can start writing. How it works (recipe):

1. Concatenate lots of documents into one huge input.
2. The model attends over every token to build internal states.
3. Only after that can it generate the first word. Why it matters: The longer the prompt, the longer you wait for the first token, making the system feel laggy.

🍞 Bottom Bread (Anchor): If you paste 90 articles into a single prompt, the model spends most of its time just getting ready, like reading all the chapters before writing the first sentence.

🍞 Top Bread (Hook): Imagine keeping quick notes beside each book so you don’t have to reread the entire chapter every time.

🥬 Filling (The Actual Concept): What it is: A KV cache is the model’s saved “memory” (keys and values) from reading a piece of text so the model can reuse it later without re-reading. How it works (recipe):

1. Encode each document once to produce its KV cache.
2. Save these caches in a datastore.
3. At question time, load the caches and skip re-encoding. Why it matters: Without KV caches, the model redoes lots of work and delays answers.

🍞 Bottom Bread (Anchor): It’s like sticky notes summarizing each chapter: you don’t reread the whole book—just peek at the notes to remember key points.

🍞 Top Bread (Hook): When two books contain two halves of a puzzle, you need to connect them to solve the mystery.

🥬 Filling (The Actual Concept): What it is: Cross-document reasoning is when the AI combines clues from different documents to reach the final answer. How it works (recipe):

1. Identify a bridging clue in one document (like a person’s name).
2. Use it to find the missing piece in another document (like the event or date).
3. Put the pieces together into one answer. Why it matters: Without this, the AI answers from only one source and may miss the full story.

🍞 Bottom Bread (Anchor): To answer “Where did the player make his debut?”, the AI might need one page that mentions the player and another that lists the debut event.

🍞 Top Bread (Hook): Picture a teacher asking many students for their opinions; some students are right, some are noisy, and you need a fair way to pick the best answer.

🥬 Filling (The Actual Concept): What it is: Logits are the model’s raw scores for each possible next word before turning them into probabilities. How it works (recipe):

1. For every step, the model assigns a score to each word in the vocabulary.
2. Higher scores mean the model thinks that word is more likely next.
3. These scores guide which word is chosen. Why it matters: If we can adjust these scores smartly, we can make the AI prefer evidence-backed words.

🍞 Bottom Bread (Anchor): If the words “Paris” and “London” both get scores, the one with the higher logit is more likely to be picked as the next word for “Capital of France is …”.

🍞 Top Bread (Hook): Think of comparing a guess with and without seeing the textbook—if reading the book makes a word much more likely, that’s a good sign.

🥬 Filling (The Actual Concept): What it is: Contrastive decoding compares an expert’s suggestion (with context) against a prior suggestion (without that context) to highlight what the document truly adds. How it works (recipe):

1. Compute logits with the document and without it.
2. Subtract the “without” from the “with” to see what the document specifically supports.
3. Prefer words that are boosted by the document. Why it matters: Without this, the model might rely too much on its own habits instead of the evidence.

🍞 Bottom Bread (Anchor): If “1928” becomes way more likely after reading a page about penicillin, contrastive decoding helps pick “1928” over a vague, model-only guess.

The world before this paper: Teams tried two main paths. One, smash many documents into a long prompt to keep cross-document attention; this helped reasoning but caused big prefill delays and “lost in the middle” attention issues. Two, encode each document separately to save time; this was fast but broke cross-document interactions because the model couldn’t easily “see” across caches. People tried merging caches, adding special bridging tokens, or training tricks—but these either needed extra training or still struggled to fully restore clean cross-document reasoning without reintroducing heavy costs. The gap: a way to keep documents modular and fast while still letting the model stitch evidence across them, and to do it without training. The stake for real life: faster, more accurate assistants that can read many files, policies, or reports in seconds and give grounded answers—useful in homework, customer support, coding help, and more.

🍞 Top Bread (Hook): Imagine a panel of librarians. Each librarian read one book and suggests words for your report. You don’t make them all sit in one big noisy meeting; instead, at every sentence you ask, “Which librarian has the strongest evidence for the next word?”

🥬 Filling (The Actual Concept): What it is: Parallel Context-of-Experts Decoding (PCED) treats each retrieved document as its own expert, runs them all in parallel, and picks each next token from the expert that supports it best. How it works (recipe):

1. Precompute a KV cache per document (one expert per doc) and add an “amateur” expert with no document (the model prior).
2. For each next token, get logits from every expert in parallel.
3. Contrast each expert’s logits against the amateur to keep only what the document truly adds.
4. Weigh each expert by how relevant its document is (retrieval scores).
5. Choose the token with the highest adjusted support (a max over experts) and append it to every expert’s history. Why it matters: Without PCED, you either pay long-context costs or you lose cross-document reasoning when using separate caches.

🍞 Bottom Bread (Anchor): If one document says “bridging player → event” and another says “event → year,” the model can hop experts across tokens—first taking a name from expert A, then the event from expert B—to complete the answer.

Three analogies:

1. Debate team: Each debater cites one source. At every word, the judge picks the debater with the strongest evidence. No messy group shout—just pick the best-supported point.
2. Radio channels: You tune to the clearest station for each lyric. If static grows, you switch channels mid-song to keep the music clear.
3. Toolbelt: For each screw or bolt, you grab the right tool. You don’t dump all tools into a single mega-tool; you just pick the best one for the next action.

Before vs. After:

• Before: Long prompts = slow starts; separate caches = fast but weak at stitching facts.
• After: PCED keeps caches separate (fast), yet stitches facts during decoding by switching experts token by token.

🍞 Top Bread (Hook): You know how a coach balances a team’s raw talent with who actually practiced for the game?

🥬 Filling (The Actual Concept): What it is: Retrieval-aware contrastive decoding boosts an expert’s token only if (a) that token is especially supported by its document compared to the model’s prior and (b) the document’s retrieval score says it’s relevant. How it works (recipe):

1. For each expert, compare “with-doc” vs. “no-doc” logits to highlight doc-specific support.
2. Add a bonus based on the document’s relevance score (from the retriever/reranker).
3. Pick the token with the strongest adjusted score across experts. Why it matters: Without the relevance-aware bonus, noisy or off-topic documents can distract the model; without the contrast, the model may ignore the document and rely on its prior.

🍞 Bottom Bread (Anchor): If a reranker says Doc 3 is highly relevant, its helpful tokens get a small lift; if Doc 8 is low relevance, its tokens are downweighted—even if they look tempting.

🍞 Top Bread (Hook): Think of a traffic cop waving cars from different streets through the intersection, one by one, to keep everything moving smoothly.

🥬 Filling (The Actual Concept): What it is: Expert switching at the token level means PCED can choose different experts for different words in the same answer. How it works (recipe):

1. Compute adjusted scores per expert.
2. Take a max across experts per token.
3. Append the chosen token to a shared history so all experts stay in sync for the next step. Why it matters: Without token-level switching, the model might get stuck with one document and miss crucial details from others.

🍞 Bottom Bread (Anchor): In a multi-hop question, the answer might start with a name backed by expert A, then switch to expert B for the event, and finally to expert C for the date—one smooth sentence, several experts.

Why it works (intuition):

• The contrast step cancels the model’s “habit words” and keeps what the document uniquely supports.
• The retrieval prior gates out irrelevant experts so noise doesn’t win.
• The max rule allows sharp, decisive switching across documents, which is especially helpful when different docs are useful at different steps.

Building blocks:

• Experts = per-document KV caches (plus the amateur prior).
• Contrast = highlight doc-added evidence by comparing to the prior.
• Retrieval prior = use fused retriever + reranker scores to trust good docs more.
• Max aggregation = pick the single best-supported token each step.
• Shared history = every chosen token is fed back to all experts to keep them aligned.

Overview (high level): Input (query + top-N retrieved documents and their scores) → Build N contextual experts (one per document) + 1 amateur expert (no context) → Parallel forward pass to get per-expert logits → Retrieval-aware contrastive adjustment per expert → Max across experts to choose next token → Append token to all experts’ histories → Repeat until done → Output answer.

Step-by-step (what, why, example):

1. Prepare offline KV caches (experts)

• What happens: Each document in the corpus is encoded once to produce and store its KV cache with an embedding for retrieval.
• Why it exists: So at query time, we don’t pay the prefill cost for long inputs; we just load the caches.
• Example: For a sports wiki, each athlete’s page becomes a saved cache. Later, questions about matches can reuse those caches instantly.

1. Retrieve and score documents

• What happens: Given a query, a retriever grabs top-N candidate documents. A reranker refines the order. We normalize and fuse these into a single relevance score per document.
• Why it exists: Not every retrieved doc is equally useful; relevance scores help trust good docs and ignore distractors.
• Example: For “Where did Jeff Wood make his CART debut?”, retrieval might pull pages on Jeff Wood, the race, and unrelated drivers. The reranker promotes the on-topic pages.

1. Build parallel experts (N contextual + 1 amateur)

• What happens: We create N inputs that share the same question but differ by which document cache conditions them. We also include the amateur expert with no document, representing the model’s prior knowledge.
• Why it exists: This setup keeps documents modular and allows fast, batched decoding.
• Example: Expert 1 = Jeff Wood page, Expert 2 = Caesars Palace Grand Prix page, Expert 3 = unrelated page, Amateur = no-doc prior.

1. Get per-expert logits in parallel

• What happens: For each generation step, all experts run in one batched forward pass and output logits over the vocabulary.
• Why it exists: Running experts together keeps latency low and makes token-level switching possible.
• Example: For the next word after “He made his CART debut at the…”, each expert suggests candidates (like “Caesars”, “Indianapolis”, dates, etc.).

1. Retrieval-aware contrastive adjustment

• What happens: For each expert, subtract the amateur’s logits (contrast) to keep what the document adds, then add a small bonus from the document’s relevance score (retrieval prior). Two knobs control the balance: one for contrast strength (how much we subtract the amateur) and one for how strongly we use the relevance score.
• Why it exists: Without contrast, the model might lean on habits; without the relevance prior, noisy experts can distract.
• Example: If Expert 2’s document is highly relevant and makes “Caesars” much more likely than the amateur, “Caesars” gets a strong adjusted score.

1. Token selection via max across experts

• What happens: For each possible next token, take the maximum adjusted score among experts and pick the overall best token.
• Why it exists: Max allows sharp, decisive switching between experts when different docs are useful at different steps.
• Example: The model picks “Caesars” from Expert 2. On the next token, it might pick “Palace” from the same expert, or switch to Expert 1 for the year.

1. Share the chosen token with all experts

• What happens: Append the selected token to every expert’s generation history so they all stay in sync.
• Why it exists: Keeping experts aligned ensures their next-step predictions consider what’s already been written.
• Example: After writing “…Caesars”, all experts now condition on that partial answer when predicting the next word.

1. Repeat until the answer is complete

• What happens: Continue the parallel-contrast-select-share loop until the model outputs an end token or meets a stop rule.
• Why it exists: This loop gradually stitches evidence from multiple documents into one coherent answer.
• Example: The final answer reads the correct race and year, even if no single document contained everything.

Mini walk-through with data:

• Question: “At which 1983 race did Jeff Wood make his CART debut?”
• Experts: E1 (Jeff Wood bio), E2 (Caesars Palace Grand Prix 1983), E3 (distractor), Amateur (no doc).
• Step t: E1 boosts “Jeff Wood” details; E2 boosts “Caesars Palace Grand Prix”. After contrast + relevance, “Caesars” from E2 wins by max and is chosen. Token is shared with all experts. Next token “Palace” again from E2. Then the model finishes the phrase correctly.

The secret sauce:

• Contrast isolates what each document truly contributes over the model’s prior.
• Relevance scores downweight distractors so they can’t hijack the generation.
• Max aggregation enables crisp token-level expert switching, which is especially important for multi-hop reasoning.
• All of this is training-free and works with modular, precomputed caches—so it’s both fast and scalable.

Implementation notes (intuitive, no equations):

• Contrast strength is set adaptively at the first token, then held steady (works robustly across tasks).
• A single relevance weight in a comfortable mid-range tends to work best—too low lets noise in, too high can over-trust the retriever.
• Max aggregation is generally best for multi-doc hops; soft mixtures can be slightly better when one doc clearly has everything.

The tests and why:

• Multi-document QA (HotpotQA, MuSiQue, QAMPARI, QuEST): to see if the model can stitch clues across documents.
• Single/focused QA (NQ) and few-shot reasoning (ICL tasks like Web, Tracking, Date): to test using the right example or the right single source among many distractors.
• LongBench subsets and distractor settings: to stress-test robustness when irrelevant context is present.
• Latency: time-to-first-token (TTFT) and end-to-end speed, because users feel these immediately.

Competitors:

• Corpus in Context (Single vs. All): Standard long-prompt baselines stuffing top-1 or many docs into one input.
• KV merging (APE): Encodes docs separately, then tries to reconstruct cross-document attention.
• Agentic MapReduce: Summarize each doc, then do a final aggregation call.

Scoreboard with context:

• Accuracy: PCED consistently outperforms prior parallel methods (like APE) by large margins on multi-hop QA and often matches or beats long-context baselines despite not concatenating docs. For example, on Llama-3.1 8B for QAMPARI, APE scores around 7, while PCED jumps to the 70s—like going from an F to a strong A. On NQ, PCED hits mid-80s, improving over Single and All concatenation baselines.
• LongBench: Across single-doc and multi-doc QA, and summarization, PCED variants meet or exceed the all-in-context baseline, showing resilience to distractors.
• Efficiency: TTFT drops by up to roughly 180× (about 0.14s vs. 25.50s), which feels like answers appearing almost instantly. End-to-end latency also improves substantially on long-context workloads.

Surprising/insightful findings:

• Decode-time stitching really works: Visual traces show the model hopping experts mid-answer, aligning with the PCED design.
• Max vs. mixtures: Max aggregation is best for multi-hop reasoning (decisive switches), while soft mixtures can edge ahead when a single doc dominates.
• Both parts matter: Removing either the contrast or the retrieval prior causes big drops. Together they deliver the best stability and accuracy.
• Scales with top-k: Accuracy stays stable even as the number of experts grows (e.g., from 8 to 128), suggesting the relevance prior effectively suppresses noise.
• MapReduce can still win in some cases that need a global synthesis step, but it requires multiple LLM calls; PCED does it in one decoding stream.

Bottom line: PCED brings back most of the cross-document benefits of long prompts while avoiding their latency and attention-dilution problems—delivering strong accuracy with much faster starts.

Limitations:

• Needs logits: PCED must read full per-token logits from the model for each expert. That’s fine for open/self-hosted models, but many closed APIs don’t expose them.
• Retrieval quality: If the right document isn’t retrieved or is scored poorly, its expert won’t be picked. PCED can’t invent missing evidence.
• Storage trade-off: Saving many per-document KV caches consumes storage that scales with corpus size and model dimensions; this is ideal for static, read-heavy corpora but less so for fast-changing data.

Required resources:

• An open or self-hosted LLM that exposes logits.
• A retrieval stack (embedding retriever + reranker) producing normalized relevance scores.
• Storage for caches and a serving system that supports batched parallel decoding.

When not to use:

• Purely generative/creative tasks where there’s no external evidence to retrieve.
• Highly dynamic corpora that change constantly (cache churn may outweigh latency savings).
• API-only deployments without logit access.

Open questions:

• End-to-end training: What if models were trained to accept parallel experts directly and learn the switching rule themselves?
• Adaptive aggregation: Can we learn when to use max vs. mixtures per step or per task?
• Better priors: Can we refine how retriever and reranker scores are fused or make them token-aware rather than document-level?
• Compression: How far can we compress KV caches without harming accuracy, to reduce storage while keeping speed?

Three-sentence summary: PCED is a training-free decoding framework that treats each retrieved document as an expert, contrasts each expert against a no-doc prior, and uses retrieval scores to gate noise. By picking each next token from the best-supported expert and sharing it with all, PCED stitches evidence across documents without long-context attention. It matches or beats strong baselines on multi-document tasks while delivering dramatic speedups in time-to-first-token.

Main achievement: Showing that cross-document reasoning can be recovered at decode time—without joint attention—by combining contrastive calibration and retrieval-aware weighting over parallel, modular KV caches.

Future directions: Learn expert selection end-to-end, develop token-level or span-level priors, adaptively blend max and mixture aggregation, and compress caches for lighter storage footprints.

Why remember this: PCED reframes how we combine sources—don’t make the model read a phonebook; let it ask many librarians in parallel and pick the best-supported word at every step. That simple shift unlocks both accuracy and speed for retrieval-heavy AI.