Intermediate Layer Retrieval (ILRe)

• ILRe is a context compression method that selects an intermediate decoder layer to perform query-aware token retrieval, reducing long-context inefficiencies.
• It uses streaming chunked prefill for early layers and applies dense retrieval at one intermediate layer, lowering computational complexity from O(L²) to O(L).
• Empirical results show ILRe maintains high performance and memory efficiency in long-context benchmarks compared to full-context and other streamlined methods.

Searching arXiv for papers on Intermediate Layer Retrieval and closely related intermediate-layer methods. Intermediate Layer Retrieval (ILRe) is a training-free context compression pipeline for decoder-only causal LLMs that addresses long-context inefficiency by selecting a single intermediate decoder layer offline, encoding context only up to that layer with streaming chunked prefill, and recalling tokens through query-to-context attention computed in that layer (Liang et al., 25 Aug 2025). In the formulation described for long-context inference, ILRe is not external document retrieval in the usual retrieval-augmented generation sense; it is an in-prompt retrieval mechanism that compresses the prompt itself. The method is presented as a response to the practical limits of long-context LLMs, including short effective context length, quadratic prefilling cost, and high KV-cache memory overhead, while preserving or surpassing full-context performance on long-context benchmarks (Liang et al., 25 Aug 2025).

1. Definition and scope

ILRe denotes a specific long-context inference pipeline in which a decoder-only causal LLM uses one offline-chosen intermediate decoder layer as a retrieval layer, denoted $l_R$, to identify salient prompt tokens before standard generation proceeds on a compressed prompt (Liang et al., 25 Aug 2025). The pipeline is explicitly described as training-free and requiring neither additional post training nor operator development. Its stated purpose is context compression rather than model retraining, architecture redesign, or external knowledge augmentation.

The method’s scope is tightly defined. It assumes that a query is available at compression time, because token recall is query-aware. It also assumes that an intermediate layer with strong retrieval ability can be identified offline. The compressed output is then handed to standard inference engines such as vLLM or SGLang for ordinary downstream decoding (Liang et al., 25 Aug 2025). This distinguishes ILRe from approaches that retrieve external documents, from early-exit methods that decode directly from intermediate layers, and from robustness or OOD methods that attach auxiliary heads to intermediate representations.

A useful contextual contrast appears in work on multi-hop retrieval-augmented generation. "Optimizing Multi-Hop Document Retrieval Through Intermediate Representations" proposes Layer-wise RAG (L-RAG), which leverages intermediate representations from middle layers to retrieve external knowledge for multi-hop question answering (Lin et al., 2 Mar 2025). That abstract-level description places ILRe and L-RAG in a common intermediate-representation design space, but the supplied material for L-RAG does not provide a method description, equations, experiments, or analyses beyond the abstract and therefore does not support a fuller technical comparison (Lin et al., 2 Mar 2025).

2. Motivating problem: long-context bottlenecks

The motivating problem is the practical cost of full-context processing in long-context LLMs. Even when a model advertises 1M-token support, effective use of that context may deteriorate as input length increases. For Llama-3.1-UltraLong-8B-1M-Instruct, the reported RULER score under full-context processing drops from 83.2 at 32K to 70.2 at 128K (Liang et al., 25 Aug 2025). At the same time, full prefill uses self-attention with quadratic complexity in sequence length, and KV-cache memory scales prohibitively. The paper estimates that a 1M-token KV cache for Llama-3.1-8B is roughly 120 GB, beyond a single device (Liang et al., 25 Aug 2025).

ILRe is proposed to mitigate both latency and memory problems. In the complexity account given for FlashAttention-style full-context prefill, computation scales as $O(N_L L^2)$ and memory as $O(2N_L L)$, where the factor 2 corresponds to keys and values (Liang et al., 25 Aug 2025). ILRe reduces the prefilling complexity from $O(L^2)$ to $O(L)$ in the long-context regime by avoiding dense full-length attention in most layers and preserving full-length key states only in the retrieval layer (Liang et al., 25 Aug 2025).

This suggests that ILRe should be understood as a resource-allocation strategy over depth: dense global access is preserved only where it is most useful for retrieval, while earlier layers operate under a cheaper streaming approximation. A plausible implication is that the method depends not only on token selection quality but also on a particular claim about how retrieval-relevant information is distributed across layers.

3. Pipeline and algorithmic structure

The ILRe pipeline has four principal stages: offline retrieval-layer selection, streaming chunked prefill up to that layer, query-aware token recall at that layer, and compressed-prompt generation with a standard inference engine (Liang et al., 25 Aug 2025).

Offline layer selection

The retrieval layer $l_R$ is selected offline through a retrieval-style probing task based on a needle-in-a-haystack setup. A random passkey is embedded in a long background context, and the model must recover it from a question prompt. For each layer and pooling configuration, recall or accuracy is measured as key depth and key length vary. The layer with the highest recall is chosen; if multiple layers tie, the smallest index among them is selected (Liang et al., 25 Aug 2025). The paper notes that the optimal retrieval layer is usually in the first half of the network. Reported examples are $l_R = 3$ for Llama-3.1-UltraLong-8B-1M-Instruct with $N_L = 32$, and $l_R = 5$ for Qwen2.5-7B-Instruct (Liang et al., 25 Aug 2025).

Streaming chunked prefill

For layers $1$ through $O(N_L L^2)$0, ILRe uses streaming chunked prefill rather than full dense attention. The input context is broken into chunks, and each streaming layer retains only a sink region of size $O(N_L L^2)$1 and a sliding window of size $O(N_L L^2)$2. The attention mask takes the $O(N_L L^2)$3-shaped pattern associated with streaming methods: tokens attend to a small fixed prefix sink and a local sliding window, but not the entire past (Liang et al., 25 Aug 2025). The default settings are $O(N_L L^2)$4, $O(N_L L^2)$5, and chunk size 1024, with the first chunk including the extra sink tokens.

Retrieval-layer attention and token scoring

At the retrieval layer $O(N_L L^2)$6, ILRe preserves the full context keys for the context segment and computes query states for the question part. For head $O(N_L L^2)$7 and layer $O(N_L L^2)$8, the query-to-context attention is defined as

$O(N_L L^2)$9

where $O(2N_L L)$0 are the query states, $O(2N_L L)$1 are the context keys, and $O(2N_L L)$2 is the head dimension (Liang et al., 25 Aug 2025). To obtain a one-dimensional importance score over context tokens, ILRe reduces across heads and query positions by

$O(2N_L L)$3

The compressed prompt is formed by selecting tokens from the original input IDs rather than only from cached states, so it can be re-fed to a standard inference stack (Liang et al., 25 Aug 2025).

Compressed-prompt generation

After token recall, the retained context tokens are concatenated with the query for standard generation. This makes ILRe a preprocessing and compression stage rather than a custom end-to-end runtime. The method is therefore compatible with ordinary downstream decoding engines (Liang et al., 25 Aug 2025).

4. Multi-pooling token allocation and semantic completeness

A central design choice in ILRe is the multi-pooling kernels allocating strategy for token recall. The motivation is that a single attention view or a single pooling granularity may not preserve semantic completeness, especially when relevant information is distributed over spans rather than isolated spikes (Liang et al., 25 Aug 2025).

The default configuration uses max-pooling kernels with sizes $O(2N_L L)$4 and average-pooling kernels with sizes $O(2N_L L)$5, yielding $O(2N_L L)$6 pooling combinations (Liang et al., 25 Aug 2025). Given a total token budget $O(2N_L L)$7, the per-combination budget is

$O(2N_L L)$8

where $O(2N_L L)$9 is the number of max kernels and $O(L^2)$0 is the number of average kernels (Liang et al., 25 Aug 2025).

The allocation procedure is described as follows: reserve sink tokens, compute pooled views of the attention vector $O(L^2)$1 for each max-pooling kernel and each subsequent average-pooling kernel, select top candidates from each pooled view, map those pooled indices back to original token indices, and allocate tokens evenly across combinations while skipping duplicates and moving to the next-best candidate when necessary (Liang et al., 25 Aug 2025).

The paper argues that small kernels capture short sharp cues, whereas larger kernels capture more diffuse or multi-token semantic units. This suggests that the retrieval target is not merely “high-attention tokens” but structurally coherent token subsets. The reported ablations support that interpretation: reducing the number of kernels lowers performance modestly, while collapsing to a single max/avg kernel pair, denoted “One-Max-Avg,” causes a substantial drop on longer-key retrieval tasks such as NS3 (Liang et al., 25 Aug 2025).

Related evidence from a different RAG setting points to a broader intermediate-layer interpretation. "LFD: Layer Fused Decoding to Exploit External Knowledge in Retrieval-Augmented Generation" reports a layer-specific functional demarcation in which shallow layers specialize in local context modeling, intermediate layers focus on integrating long-range external factual knowledge, and deeper layers increasingly rely on parametric internal knowledge (Sun et al., 27 Aug 2025). Although LFD is a decoding-time fusion method rather than an intermediate-layer retriever, it supports the general view that intermediate layers can preserve external or contextual information before later transformations overwrite or dilute it (Sun et al., 27 Aug 2025).

5. Complexity, memory, and empirical performance

The complexity analysis given for ILRe defines $O(L^2)$2 and $O(L^2)$3. The attention-compression cost is

$O(L^2)$4

for computation and

$O(L^2)$5

for memory (Liang et al., 25 Aug 2025). Because $O(L^2)$6, $O(L^2)$7, and $O(L^2)$8 are all much smaller than $O(L^2)$9 in long-context settings, the effective computational complexity is linear in context length, $O(L)$0 (Liang et al., 25 Aug 2025). The paper also states that ILRe’s memory footprint is roughly $O(L)$1 of full-context memory under the stated assumptions.

The reported empirical results are strongest on RULER-1M with Llama-3.1-UltraLong-8B-1M-Instruct and a 4K token budget. ILRe achieves 88.3 at 32K, 84.2 at 64K, 84.0 at 128K, 84.1 at 512K, and 79.8 at 1M (Liang et al., 25 Aug 2025). Full-context inference is reported as OOM at 512K and 1M. On the same benchmark family, ILRe is reported to outperform StreamingLLM, SnapKV, and RAG by large margins, with weaker degradation as context length increases (Liang et al., 25 Aug 2025).

On LongBench with Llama-3.1-8B-Instruct, ILRe is evaluated at 1K and 2K token budgets. At 2K, ILRe with full KV retention in the retrieval layer achieves especially strong results and outperforms SnapKV and StreamingLLM on most tasks, including retrieval-heavy question answering and summarization. Even with reduced retrieval layers such as $O(L)$2, the method remains competitive (Liang et al., 25 Aug 2025).

For efficiency, the paper measures Time To First Token on a Huawei Ascend 910B NPU and reports that ILRe can process a single 1M-token request in less than half a minute, with approximate speedup $O(L)$3 (Liang et al., 25 Aug 2025). A linear fit of ILRe latency versus length yields $O(L)$4, which is presented as consistent with the $O(L)$5 analysis. This empirical profile is central to the method’s significance: the claimed gain is not only accuracy retention under compression but a shift from impractical to deployable long-context processing.

6. Relation to adjacent intermediate-layer methods

ILRe belongs to a broader family of methods that exploit non-final representations, but its function is distinct from several neighboring lines of work.

One nearby line uses intermediate layers for external retrieval in RAG. The abstract of Layer-wise RAG states that middle-layer representations capture next-hop information for multi-hop question answering and can retrieve external knowledge with inference overhead similar to standard RAG (Lin et al., 2 Mar 2025). However, the supplied material explicitly states that no substantive paper text for that work is available, so only the abstract-level characterization is warranted (Lin et al., 2 Mar 2025).

A second line uses intermediate layers during decoding rather than retrieval. LFD fuses an intermediate layer’s decoded distribution with the final layer’s output to better surface retrieved context knowledge and selects the fusion layer using an Internal Knowledge Score in the latter half of layers (Sun et al., 27 Aug 2025). Likewise, LITE instruction-tunes LLaMA-2 with additional explicit losses from intermediate layers so that those layers acquire generation ability and can support dynamic confidence-based early exiting, yielding inference cost reductions of 37.86% for 7B and 46.35% for 13B while maintaining response quality (Varshney et al., 2023). These methods do not compress context by selecting prompt tokens, but they reinforce the broader claim that intermediate layers can contain operationally useful signals not fully preserved in the final layer.

A third line treats intermediate layers as alternative prediction spaces. Intermediate Layer Classifiers show that, for OOD generalization, intermediate representations can substantially outperform penultimate-layer representations, with zero-shot improvements such as Waterbirds 79.4% to 87.1%, CelebA 56.0% to 82.0%, and MultiCelebA 20.8% to 46.0% (Uselis et al., 7 Apr 2025). This suggests that non-final layers may be less over-specialized and less sensitive to shift. Although this is a classification setting rather than long-context compression, it is consistent with ILRe’s premise that retrieval-relevant information may be better exposed in intermediate depth than at the model output.

A plausible implication of these adjacent results is that “intermediate layer” is not a single functional category. In different settings, intermediate layers are used for external retrieval, decoding fusion, early exit, or shift-robust prediction. ILRe’s specificity lies in using one such layer as a query-aware in-context selector under streaming lower-layer computation.

7. Assumptions, limitations, and interpretation

ILRe makes several explicit assumptions. It is designed for decoder-only causal LLMs, relies on query-aware compression, and depends on offline identification of a good retrieval layer using a probing task (Liang et al., 25 Aug 2025). It therefore presupposes a deployment regime in which compression is performed after the query is known and before full decoding begins. This excludes purely query-agnostic prompt compression.

The method also assumes that ordinary inference frameworks can accept the compressed prompt after retrieval. ILRe is not presented as a modification of the model’s internal runtime for generation beyond the compression stage. Another explicit limitation is that layer choice is empirical: different models may favor different retrieval layers, as illustrated by $O(L)$6 for Llama-3.1-UltraLong-8B-1M-Instruct and $O(L)$7 for Qwen2.5-7B-Instruct (Liang et al., 25 Aug 2025).

Comparisons in the paper sharpen the interpretation of what ILRe is and is not. Relative to full-context inference, ILRe trades exact dense processing for lower cost and, in the reported experiments, often better long-context accuracy (Liang et al., 25 Aug 2025). Relative to StreamingLLM, it supplements sink-and-window attention with query-aware retrieval at an intermediate layer. Relative to SnapKV, it adds offline layer selection and multi-kernel pooling for semantic completeness. Relative to RAG, it does not retrieve external documents; it compresses the prompt itself (Liang et al., 25 Aug 2025).

The broader significance of ILRe is therefore methodological rather than merely engineering. It operationalizes the claim that one need not process all layers over all tokens to preserve task-relevant long-context information. Instead, long-context performance can be mediated through a selective intermediate representation, a retrieval-style attention reduction, and a compressed re-entry into standard generation. This suggests a general research direction in which layer locality, token locality, and query dependence are co-optimized rather than treated as separate concerns.