PaperlayerView exact revision on arXiv

Source-linked AI summary

Retrieval from Within: An Intrinsic Capability of Attention-Based Models

Elad Hoffer, Yochai Blau, Edan Kinderman, Ron Banner, Daniel Soudry, Boris Ginsburg

arXiv:2605.05806v2cs.LG

TL;DR

RAG usually separates retrieval from generation, leaving open whether attention-based encoder-decoders can retrieve evidence internally. INTRA unifies both processes, and on multi-hop QA benchmarks matches or exceeds engineered RAG pipelines while enabling reuse of pre-encoded evidence.

  • Problem

    RAG separates retrieval from generation, motivating whether one attention-based encoder-decoder can select needed evidence and answer queries internally.

  • Method

    INTRA uses decoder attention to retrieve pre-encoded evidence chunks and directly reuse them as generation context within the same encoder-decoder.

  • Results

    INTRA matches or exceeds engineered RAG pipelines on multi-hop QA in evidence recall and end-to-end answer quality, while reusing pre-encoded evidence reduces prefilling costs.

  • Takeaways & Limitations

    Attention-based encoder-decoders can provide a unified alternative to modular RAG by carrying out retrieval within their representation space.

  • Takeaways & Limitations

    Evaluation is limited to a fixed context pool, one T5Gemma2-style encoder-decoder implementation, and short-answer text-QA benchmarks.

Abstract

from arXiv · show

Retrieval-augmented generation (RAG) typically treats retrieval and generation as separate systems. We ask whether an attention-based encoder-decoder can instead retrieve directly from its own internal representations. We introduce INTRA (INTrinsic Retrieval via Attention), a framework where decoder attention queries score pre-encoded evidence chunks that are then directly reused as context for generation. By construction, INTRA unifies retrieval and generation, eliminating the retriever-generator mismatch typical of RAG pipelines. This design also amortizes context encoding by reusing precomputed encoder states across queries. On question-answering benchmarks, INTRA outperforms strong engineered retrieval pipelines on both evidence recall and end-to-end answer quality. Our results demonstrate that attention-based models already possess a retrieval mechanism that can be elicited, rather than added as an external module.

1 Introduction

The paper asks whether attention-based encoder-decoder models can retrieve evidence from their own representations instead of relying on a separate retriever. It introduces INTRA, which couples evidence selection and answer generation through shared representations and attention-based retrieval.

  • RAG is the default design when relevant information is sparse because concatenating the full corpus into one prompt is computationally expensive.
  • Standard RAG separates retrieval from generation, despite attention already providing a query-conditioned mechanism for selecting and weighting relevant information.
  • INTRA uses a pretrained encoder-decoder to select evidence and generate answers over representations produced by its own encoder.
  • INTRA converts decoder cross-attention queries into chunk-level retrieval scores, using the same representation space for retrieval and generation.The paper frames attention and retrieval as query-conditioned matching over candidate states.
  • Shared encoded chunk states reduce retriever-generator mismatch and allow static evidence to be encoded once and reused across queries.The mechanism can also be adapted with lightweight decoder-side retrieval queries.
  • INTRA’s minimal recipe reuses native encoder chunk representations, performs coarse encoder-side late-interaction retrieval, and refines evidence with decoder-side queries.It avoids introducing a separate retriever or compression model.

2 Method

INTRA uses a shared encoder-decoder to pre-encode evidence chunks, retrieve them with decoder-side attention signals, and generate directly from the selected encoded context. Its retrieval pass scores the full corpus and is followed by generation over the selected chunks, while initialization may provide a starting context without restricting final retrieval.

  • Shared encoded context: The encoder maps each evidence chunk t_i to a precomputed representation k_i, and generation conditions on concatenated encoded context K(S) rather than raw retrieved text.The same pretrained encoder-decoder model encodes evidence and decodes the answer.
  • Intrinsic retrieval: INTRA converts decoder cross-attention’s token-level matching between query states and encoded chunks into chunk scores using a scaled ColBERT-style MaxSim late-interaction score.MaxSim takes the best-matching chunk token for each query token instead of applying a softmax over all tokens.
  • Intrinsic retrieval: It selects the top-n scored chunk indices as S_INTRA, then runs the ordinary decoder on K(S_INTRA) to generate the output.The selected set is defined by the largest retrieval scores.
  • Inference procedure: Inference uses two decoder passes: a retrieval pass exposes query states to score chunks, followed by a generation pass over the selected encoded context.Both passes operate over the pre-encoded chunk set K.
  • Initialization: The initial set S_0 may come from encoder similarity to the input or be empty, but final retrieval scores the full corpus and can select chunks outside S_0.Unlike reranking, INTRA does not restrict final selection to an initially retrieved candidate set.

3 Practical Implementation

INTRA adapts a pretrained encoder-decoder by reusing pre-encoded chunk states for retrieval and generation. Its implementation moves key transformations to decoder queries, preserving attention scores while enabling a shared normalized chunk representation and efficient retrieval training and scoring.

  • Model adaptation: INTRA modifies T5Gemma2 decoder cross-attention so pre-encoded chunk states are reused directly for retrieval and generation.The implementation starts from T5Gemma2 and changes decoder cross-attention rather than adding a separate retriever.
  • Reverse-QWK: Reverse-QWK stores one normalized encoder representation and moves each layer’s learned key scale and projection matrix to the query side.This avoids repeatedly transforming stored encoder states while retaining the cross-attention computation.
  • Shared representation space: Retrieval and attention use the same normalized representation space, allowing decoder queries from different layers to score the same stored chunk pool.MaxSim uses the normalized chunk vectors, preserving the attention scores while sharing chunk representations across layers.
  • Retrieval supervision: With explicit retrieval supervision, a soft cross-entropy objective assigns equal target mass to oracle evidence chunks and updates retrieval tokens and layer aggregation weights.The decoder remains frozen while the induced decoder queries are trained to place probability mass on the oracle evidence set.
  • Efficient scoring: For efficient scoring, INTRA replaces each chunk’s Lc-token sequence with an Lp-token mean-pooled sequence where Lp ≪ Lc.The pooled vectors are fixed averages of the encoder’s own states, requiring no additional compressor or compression-specific training.

4 Benchmarks and Experimental Setup

The evaluation covers four Wikipedia-based QA benchmarks using a shared 759K-chunk retrieval pool under an approximately 100M-token budget. INTRA is compared with nine retrieval baselines, with experiments built from a T5Gemma2 4B-4B checkpoint and complete-evidence recall as a retrieval metric.

  • Benchmarks: Four benchmarks span bridge, comparison, two-hop, compositionally hard multi-hop, and single-hop open-domain QA.The benchmarks are HotPotQA, 2WikiMultihopQA, MuSiQue, and Natural Questions.
  • Benchmarks: The shared retrieval candidate pool contains 759K chunks under an approximately 100M-token budget.One pool is used across all four benchmarks.
  • Baselines and Metrics: INTRA is compared with nine baselines spanning sparse lexical, dense single-vector, reranking, hybrid RAG, and ColBERT-style late-interaction retrieval.The baselines include TF-IDF, BM25, BGE-large, Qwen3-Embedding-0.6B/4B, Jina reranker, RRF, and a ColBERT-style MaxSim baseline.
  • Baselines and Metrics: Complete-evidence recall is reported at k ∈ {5, 10, 20}, measuring the fraction of examples where all oracle chunks are retrieved.The evaluation also includes end-to-end QA, though the supplied passage does not provide its full metric definition.
  • Implementation Details: The model starts from a T5Gemma2 4B-4B checkpoint, while retrieval training updates only retrieval token embeddings ρ_i and layer aggregation weights α_l.The trainable components contain approximately 164K and 272 parameters, respectively; the encoder and decoder backbones remain frozen.

5 Results

INTRA improves complete-evidence retrieval and end-to-end question answering, especially on multi-hop benchmarks, by using decoder attention to align retrieval with generation. Reusing pre-encoded chunk representations also reduces generator-side prefilling costs as the number of retrieved chunks increases.

  • 5.1 Retrieval quality: Complete-evidence recall@k measures the fraction of examples whose annotated supporting chunks all appear within the top-k results.The paper treats this as a clearer retrieval-quality proxy than recovering only incomplete supporting evidence.
  • 5.1 Retrieval quality: INTRA is strongest on multi-hop retrieval benchmarks because decoder attention supports assembling multiple evidence pieces, while its advantage is less pronounced on NQ.The cited multi-hop benchmarks are HotPotQA, 2Wiki, and MuSiQue.
  • 5.1 Retrieval quality: Full-corpus INTRA scoring yields the largest complete-evidence-recall gains by recovering evidence absent from the initial candidate pool.Reranking the initial retrieval set is beneficial but does not match full-corpus scoring.
  • 5.2 End-to-end answer quality: With a fixed T5Gemma2 generator, INTRA surpasses all baselines on HotPotQA, 2Wiki, and MuSiQue, using a frozen generation-only decoder as its retrieval signal.Sharing the same decoder for retrieval and generation closes the largest average gap between random and complete-evidence contexts.
  • 5.3 Efficiency: INTRA avoids re-encoding selected evidence at query time by retrieving pre-encoded chunk states for decoder cross-attention, keeping prefilling small as retrieved-chunk count increases.Standard RAG becomes slower in the TTFT benchmark as k increases, whereas INTRA reuses stored chunk states.

6 Related Work

This section positions INTRA against modular RAG, late-interaction retrieval, memory-based approaches, and long-context modeling. Across these comparisons, INTRA uses decoder cross-attention to retrieve and consume evidence within generation while reusing pre-encoded representations.

  • Retrieval-Augmented Generation Pipelines: Most knowledge-intensive QA systems use modular retrieval-then-generation pipelines, whereas INTRA targets a single-pass retrieval block within such pipelines.Prior systems include REALM, RA-DIT, DPR, RAG, and Atlas; multi-pass agentic RAG instead interleaves reasoning with repeated retrieval.
  • Late Interaction and Representation-Space Retrieval: INTRA resembles late-interaction retrieval but replaces a dedicated retriever with decoder cross-attention that scores and then consumes matched representations during generation.ColBERT, ColBERTv2, and ColPali compare query and document tokens using MaxSim-style matching over multi-vector representations.
  • Memory, Latent Retrieval, and Unified Retrieval-Generation: Attention-based memory work frames attention as content-based access, while Memory Networks use differentiable lookup, RETRO injects retrieved chunks, and Titans argues for explicit memory mechanisms.The passage distinguishes these approaches by their source of memory, noting that Titans learns one.
  • Long-Context Modeling and Efficient Sequence Architectures: INTRA complements rather than replaces long-context modeling, addressing sparse evidence use while reusing pre-encoded evidence that standard RAG re-encodes before decoding.Dense attention has quadratic cost, motivating sparse, linear, and reduced-cost architectures; Figure 4 compares time to first token as retrieved chunks k increases, excluding retrieval time.

7 Limitations

INTRA’s evaluation is limited to a fixed context pool and one T5Gemma2-style encoder-decoder implementation, without establishing suitability for open-web retrieval or general-purpose retrievers. Its systems analysis isolates re-encoding versus reuse, so deployment trade-offs may differ when retrieval, indexing, storage, and data movement are included.

  • Scope: INTRA is evaluated on a fixed context pool, not as a replacement for RAG in open-web retrieval or webscale settings.The paper discusses billion-token practicality separately but does not position INTRA for those settings.
  • Scope: The experiments do not show that the observed retrieval mechanism generalizes from a pretrained encoder-decoder to a general-purpose retriever.The demonstrated behavior is explicitly limited to the evaluated regime.
  • Scope: The study focuses on a single implementation family: a T5Gemma2-style encoder-decoder with Reverse-QWK.No broader implementation-family comparison is reported in the limitation passage.
  • Systems evaluation: TTFT excludes retrieval time, while end-to-end deployment costs also depend on indexing, storage format, data movement, and memory size.Token-level memories can be substantially larger than compressed retrieval indices, so analyzed trade-offs may not hold end to end.

8 Conclusion

INTRA demonstrates that attention-based encoder-decoder models can perform retrieval within their own representation space, unifying retrieval and generation. It achieves competitive or superior question-answering performance while enabling reuse of pre-encoded evidence across queries.

  • Conceptual contribution: INTRA elicits retrieval within an encoder-decoder’s representation space, using the same pretrained model for retrieval and generation to eliminate representation mismatch.The framework unifies retrieval and generation through the model’s intrinsic attention-based retrieval mechanism.
  • Empirical results: On HotPotQA, 2WikiMultihopQA, and MuSiQue, INTRA competes with or exceeds engineered RAG pipelines in complete-evidence recall and end-to-end answer quality.These results cover multi-hop question-answering benchmarks.
  • Computational advantage: Static evidence can be encoded once and reused across queries, reducing prefilling costs and time-to-first-token during generation.This computational advantage follows from INTRA’s shared representation space.

Supplementary Material … D.3 INTRA

The supplementary material develops INTRA’s Reverse-QWK attention, retrieval training, ablations, and efficiency analysis. It shows how shared encoder representations support ANN retrieval and avoid repeated evidence re-encoding at query time.

  • A Reverse-QWK: Reverse-QWK moves layer-specific key projection and scaling onto decoder queries, making logits identical while enabling one shared ANN index over encoder representations.The transformation preserves standard attention scores but avoids separately indexing layer-specific keys.
  • A.1 Dimensionalities and Group-Query Attention: GQA is compatible with Reverse-QWK because a head-agnostic encoder pool is shared across query heads, reducing key materialization from O(N nh dh) to O(Nd).Head-group replication occurs on the query side rather than by expanding encoder keys.
  • A.2 Implementation Details: Implementation changes only the key path: queries receive the key transformation, scores use the shared normalized pool, while values and output attention retain the standard GQA path.The pseudocode applies query-side Wk and γk, broadcasts the shared pool, and computes output as softmax(scores) @ V.
  • A.3 Practical Considerations: Reverse-QWK stores one positionally invariant N × d pool, supports precomputation and quantization, and yields approximately 30× cross-attention KV-cache compression for the T5Gemma2 4B-4B model.For L = 34 and nkvdh ≈d/2.5, the compression ratio is on the order of ∼30×, and cached evidence is not re-encoded during cross-attention prefill.
  • B.1 Retrieval Training Details: INTRA is warm-started from a CLaRa-adapted T5Gemma2 4B-4B checkpoint and jointly trained on HotPotQA, 2Wiki, MuSiQue, and NQ using pre-encoded evidence.Retrieval training updates 64 retrieval-token embeddings, approximately 164K parameters, and 272 layer-head aggregation parameters; scoring uses retrieval-token positions.
  • B.2 Baseline Details: The evaluation compares INTRA with sparse lexical, dense single-vector, reranked, and hybrid RAG baselines, while end-to-end EM and F1 use a fixed T5Gemma generator with 95% confidence intervals.Baselines include TF-IDF, BM25, BGE-large, Qwen3-Embedding-0.6B, Qwen3-Embedding-4B, Jina reranking, and hybrid retrieval.
  • C.1 Ablation Study: Ablations vary initial context, similarity metric, context size, pooled chunk length, retrieval-token count, and generation chunk selection, showing that initial context and learned retrieval-token scoring are important.Table 7 reports CE-recall@5 and EM on HotPotQA and 2Wiki with k = 5 retrieved chunks.
  • D.3 INTRA: Compared with long-context prompting and standard RAG, INTRA reuses offline-encoded chunk representations, with query-time Prefill = O(Lq(Lq+kLc)) and Generation = O(Lg(Lq+kLc+Lg)).The efficiency analysis characterizes retrieval as O(MLqLc) and identifies standard RAG’s quadratic encoder-side re-encoding as the main distinction from INTRA.

D.4 Additional Timing Results and Benchmark Details · D.5 Compute Resources

D.4 shows that INTRA reduces generator-side time-to-first-token by reusing stored chunk states, with measured scaling favoring INTRA as retrieval depth increases. D.5 reports the H100-based training scale and compute costs for pretraining and benchmark-specific runs.

  • D.4 Additional Timing Results and Benchmark Details: The efficiency benchmark fixes Lq = Lc = Lg = 128 and N = 65,536, excludes retrieval time, and gives RAG and INTRA identical top-k chunks.It measures generator-side cost after evidence selection.
  • D.4 Additional Timing Results and Benchmark Details: Long-context prompting incurs O((Lq + N)^2) prefill, standard RAG incurs O((Lq + kLc)^2), and INTRA pays only O(Lq^2) before the first token.INTRA reuses stored chunk states rather than re-encoding selected evidence.
  • D.4 Additional Timing Results and Benchmark Details: As k increases from 1 to 500, INTRA TTFT grows from 12.8 ms to 65.7 ms, while standard RAG grows from 23.1 ms to 1.25 s.Long-context prompting is about 1.31 s in the same setting.
  • D.4 Additional Timing Results and Benchmark Details: Figure 5 reports a generation-throughput benchmark sweeping k across FC, RAG, and INTRA, with a smaller gap than for time-to-first-token.All methods use the same Reverse-QWK decoder.
  • D.4 Additional Timing Results and Benchmark Details: Figure 6 reports a time-to-first-token benchmark sweeping chunk length across FC, RAG, and INTRA.The measured trend follows the same comparison as chunk count increases.
  • D.5 Compute Resources: Training used NVIDIA H100 GPUs with 80GB per GPU, up to 160 GPUs, and approximately 2,000 GPU-hours for the longest QA pretraining run.The QA pretraining stage took approximately 12 hours, while each benchmark-specific retrieval and generation run used less than 200 GPU-hours.

E Dataset and Pool Details

The experiments use a deduplicated end-to-end QA dataset spanning four benchmarks and a shared retrieval pool built under an approximately 100M-token budget. The resulting dataset contains 617K QA examples, while the pool contains 758,500 chunks covering 514,999 unique oracle chunks.

  • Dataset: The dataset combines training and evaluation splits from HotPotQA, 2WikiMultihopQA, MuSiQue, and Natural Questions.It is derived from the union of the per-benchmark datasets constructed in CLaRa and deduplicated end to end.
  • Pool construction: The shared retrieval candidate pool merges and globally deduplicates benchmark-specific chunks under an approximately 100M-token budget.It includes all oracle chunks referenced by QA examples and adds uniformly sampled non-oracle chunks without replacement until reaching the budget.
  • Dataset statistics: 617K QA examples and 1.6M document references cover 514,999 unique oracle chunks across all splits.These are the aggregate saved-dataset statistics reported for the experiments.
  • Pool statistics: 758,500 total chunks comprise 514,999 oracle chunks and 243,501 sampled non-oracle chunks.The pool contains 69,734,081 whitespace-delimited words, averaging 91.94 words per chunk with a 106-word median.
Loading 2605.05806v2…