Superposition Prompting: Improving and Accelerating Retrieval-Augmented Generation

TL;DR

This work addresses the inefficiencies of large language models handling long contexts in retrieval-augmented generation by introducing superposition prompting, a DAG-based prompting paradigm that processes multiple prompt paths in parallel and prunes irrelevant ones without fine-tuning. By employing a ForkJoin topology, equilibrium path positioning, Bayesian path pruning, and lossless runtime optimizations such as path caching and parallelization, the method significantly reduces latency while enhancing accuracy on long-context QA tasks. Extensive experiments across OpenELM, BLOOMZ, and MPT models on NaturalQuestions-Open and MuSiQue demonstrate strong time efficiency gains (up to ~93x speedups) and notable accuracy improvements (often 12–43%), with ablations validating design choices like equilibrium positioning and Bayesian pruning. The approach thus offers a practical, model-agnostic enhancement for RAG that can be deployed on existing pre-trained LLMs without retraining, with potential implications for scalable AI systems reliant on grounded, long-context reasoning.

Abstract

Despite the successes of large language models (LLMs), they exhibit significant drawbacks, particularly when processing long contexts. Their inference cost scales quadratically with respect to sequence length, making it expensive for deployment in some real-world text processing applications, such as retrieval-augmented generation (RAG). Additionally, LLMs also exhibit the "distraction phenomenon", where irrelevant context in the prompt degrades output quality. To address these drawbacks, we propose a novel RAG prompting methodology, *superposition prompting*, which can be directly applied to pre-trained transformer-based LLMs *without the need for fine-tuning*. At a high level, superposition prompting allows the LLM to process input documents in parallel *prompt paths*, discarding paths once they are deemed irrelevant. We demonstrate the capability of our method to simultaneously enhance time efficiency across a variety of question-answering benchmarks using multiple pre-trained LLMs. Furthermore, our technique significantly improves accuracy when the retrieved context is large relative the context the model was trained on. For example, our approach facilitates a 93x reduction in compute time while *improving* accuracy by 43% on the NaturalQuestions-Open dataset with the MPT-7B instruction-tuned model over naive RAG.

Figures (7)

• Figure 1: Theoretical Maximum Speedup vs. Accuracy (Best EM Subspan) on NaturalQuestions-Open using the mpt-7b-instruct model bloomz__Muennighoff2023CrosslingualGT. Refer to \ref{['section:experiment:nq']} for experimental details. Plotted values are sourced from \ref{['table:nq_baselines']} and \ref{['table:nq_topk_ablation']}.
• Figure 2: Comparison of superposition prompting vs. the "classical" (Naive LLM-RAG) prompting paradigm. Squares represents a token, and arrows depict attention dependencies. Whereas the classical approach is a "linked-list" style DAG, superposition prompting arranges token dependencies such that all documents are processed independently. Due to this dependency structure, we can easily leverage the LLM logits to prune irrelevant context, improving long context reasoning. The dependency structure also allows for faster prompt processing, due to the new opportunities for caching and parallelism of the KV cache and logit computations (each gray box represents, logically, a "batch" that is processed by the LLM, reusing upstream KV caches).
• Figure 3: Implicit attention dependencies that must be computed during "online serving" (the colors in (b)-(f) correspond to the token segment colors in \ref{['figure:pdqtr_diagram']}). Note how the various optimizations reduce the computational burden required at online serving-time by pruning, precomputing, and parallelizing the work. It is worth re-emphasizing that in practice, inference is not sparse attention on one large sequence, but rather dense attention with many different shorter token segments.
• Figure 4: Visual intuition for our proposed equilibrium position assignment vs. left aligned (see \ref{['section:position_assignment']}).
• Figure 5: Sweeping values of superposition factor (SF) on the NaturalQuestions-Open dataset with a variety of models.