CoinRAG: Contextualized Information Nugget KV Cache Reuse for Long-Context RAG

Overview§

CoinRAG addresses the inefficiency of chunk-level KV cache reuse in Retrieval-Augmented Generation (RAG). Instead of decoding with long, noisy retrieved chunks, it identifies query-relevant information nuggets (sub-chunk semantic units such as sentences or phrases) and reuses their precomputed KV slices. By assembling these slices with chunk-level contextual information, CoinRAG achieves a better Pareto frontier between prefill latency and answer accuracy.

Core Methodology§

Information Nuggets and Offline Caching. Each retrieved chunk is decomposed into nuggets $\mathcal{N} = \{n_1, \dots, n_k\}$. For every nugget, the key/value vectors are precomputed using a frozen encoder and stored as KV slices indexed by token spans. This avoids recomputing representations at query time.

Two-Stage Retrieval.

1. Chunk retrieval: Standard dense retrieval selects top-$K$ chunks for the query $q$.
2. Nugget retrieval: Within each selected chunk, a lightweight cross-encoder scores each nugget by relevance to $q$:

$$ \text{score}(n_i, q) = \text{FFN}([\text{enc}(n_i); \text{enc}(q)]) $$

The top-$m$ nuggets are kept, discarding irrelevant parts of the chunk.

Contextualized KV Assembly. The selected nugget KV slices are combined with a compressed representation of the remaining chunk context to preserve global semantics. Let the chunk's KV cache be $(K^C, V^C)$ and selected nugget spans be $\{(s_j, e_j)\}$. The assembled cache is:

$$ K_{\text{ctx}} = [K^C_{\text{prefix}}; K^C_{s_1:e_1}; \dots; K^C_{s_m:e_m}; K^C_{\text{suffix}}] $$

with $V_{\text{ctx}}$ analogously. To avoid losing inter-nugget dependencies, a small contextualizer transformer applies a self-attention over these slices, allowing the nuggets to exchange information while keeping the total length far below the original chunk.

Implementation Sketch§

The system is straightforward to plug into existing RAG pipelines. The following PyTorch pseudo-code illustrates the core operations:

class CoinRAG:
    def __init__(self, encoder, selector):
        self.encoder = encoder      # frozen LLM or cross-encoder
        self.selector = selector    # nugget relevance scorer

    @torch.no_grad()
    def precompute_chunk(self, chunk_ids):
        out = self.encoder(chunk_ids, output_attentions=True)
        # store full KV and hidden states as a tensor of shape [L, H]
        return {"kv": out.past_key_values, "hidden": out.last_hidden_state}

    def select_nuggets(self, query_emb, chunk_data, spans, top_m):
        nugget_emb = [chunk_data["hidden"][s:e].mean(0) for s, e in spans]
        scores = self.selector(query_emb, torch.stack(nugget_emb))
        top_idx = scores.topk(min(top_m, len(spans))).indices
        return [spans[i] for i in top_idx]

    def assemble_kv(self, chunk_kv, selected_spans):
        K, V = chunk_kv  # shape: [L, H]
        slices = [K[s:e] for s, e in selected_spans]
        # optional: prepend a chunk summary token
        return torch.cat([K[:1]] + slices, dim=0), torch.cat([V[:1]] + slices, dim=0)

Results and Impact§

On LongBench multi-hop QA, CoinRAG outperforms chunk-level KV reuse baselines by an average 5.3% F1 while respecting low prefill latency budgets. The key insight is that not all tokens contribute equally; isolating and reusing only relevant information nuggets dramatically reduces the KV cache footprint without sacrificing answer quality. This makes it particularly attractive for serving large-scale RAG systems with tight latency constraints.