MHA-RAG: Improving Efficiency, Accuracy, and Consistency by Encoding Exemplars as Soft Prompts

TL;DR

Domain-adaptation of large foundation models with limited data remains costly, especially when using text exemplars in in-context learning. The authors introduce MHA-RAG, which encodes retrieved exemplars as soft prompts via a multi-head attention mechanism, with the number of heads as a tunable hyperparameter, enabling order-invariant aggregation and substantial inference-cost reductions. Across chemistry and medical QA benchmarks, MHA-RAG yields about a 20-point gain over standard RAG while reducing GFLOPs by roughly an order of magnitude, and it shows robust performance across exemplar counts and model families. This approach offers a practical, parameter-efficient alternative to fine-tuning for domain adaptation, with potential scalability to longer documents and broader task settings.

Abstract

Adapting Foundation Models to new domains with limited training data is challenging and computationally expensive. While prior work has demonstrated the effectiveness of using domain-specific exemplars as in-context demonstrations, we investigate whether representing exemplars purely as text is the most efficient, effective, and stable approach. We explore an alternative: representing exemplars as soft prompts with an exemplar order invariant model architecture. To this end, we introduce Multi-Head Attention Retrieval-Augmented Generation (MHA-RAG), a framework with the number of attention heads serving as a simple hyperparameter to control soft prompt-generation across different tasks. Across multiple question-answering benchmarks and model scales, MHA-RAG achieves a 20-point performance gain over standard RAG, while cutting inference costs by a factor of 10X GFLOPs-delivering both higher accuracy and greater efficiency, invariant to exemplar order.

Figures (5)

• Figure 1: Comparison of domain-adaptation methods: (a) RAG uses retrieved exemplars directly as context, while (b) xRAG, (c) xRAG-K, and (d) MHA-RAG derive soft-prompt representations from exemplars. The figure illustrates how each method forms its representation (single vector in xRAG, $K$ vectors in xRAG-K, multi-head representations in MHA-RAG) and highlights their trainable components. Among the four methods, only MHA-RAG---due to its use of scaled dot-product attention---is invariant to the order of exemplars.
• Figure 2: Comparison of inference compute in FLOPS with respect to $K$
• Figure 3: Varying number of heads in MHA-RAG. Performance averaged across 3 seeded runs.
• Figure 4: Total FLOPs for inference with encoder ChemBERTa-10M-MTR and foundation model Qwen3-4B using an increasing number of exemplars $c$ in the context, given $K=5$ and $K=10$ as input to create a soft prompt.
• Figure 5: Varying number of heads in MHA-RAG. Performance averaged across 3 seeded runs.

Theorems & Definitions (2)

• Claim 1: Exemplar-Order Invariance
• proof