Current RAG systems encode documents into embeddings for retrieval, but after retrieval, the embedding is discarded and the original text is passed to the LLM again.
This seems inefficient: the LLM has to process and understand the retrieved text again, even though a semantic representation was already computed during indexing.
What if the LLM itself had an “embedding/memory mode”?
LLMAE looks particularly relevant to the idea I had in mind. I’ll need to study it more carefully.
My main motivation is to make RAG more efficient for local LLMs, especially on hardware with limited VRAM. Ideally, instead of retrieving text and processing it all over again, we could retrieve useful information in a form the LLM can directly reuse.
I’m not a researcher, so I may not be able to implement this myself, but your references give me a much better starting point. Thanks again for pointing me in the right direction!
I went through this exact process while building a local RAG engine for technical and scientific PDFs (Docling OCR, dense math, LaTeX). I ended up implementing persistent KV-cache caching on disk and ran it on local hardware (RTX 4070 Ti, testing Gemma-4 and Ministral).
The short version: bypassing prefill by stored attention states works, but making it practical hits three walls very quickly.
- Raw cache size kills the I/O advantage
Storing standard FP16 or BF16 KV tensors across every layer and head produces massive files. A dense technical page easily eats tens of megabytes. If your query latency gets eaten up by streaming gigabytes off an NVMe over the PCIe bus, you gain nothing over normal token prefill.
To keep it lightweight, I had to serialize states page-by-page into a custom binary container (.cgc):
Stored keys and values quantized in 4-bit NormalFloat (NF4), followed by zstd compression.
Downsampled the sequence dimension by a factor of 2 during extraction (step_seq=2), then used 1D linear interpolation (torch.nn.functional.interpolate) on the fly when injecting into the GPU to realign the attention shapes.
In practice, a 108-page document came down to 42.3 MB on disk (~390 KB per page). Total VRAM stayed flat around 6 GB while keeping formulas and structural layout intact.
- Attention layers and RoPE mechanics
You can’t just inject raw tensors into the forward pass and expect the model to continue smoothly:
RoPE collision: If you load a cached context, the prompt tokens must have their position_ids manually offset starting at past_len (past_len -> past_len + new_len). Otherwise, positional embeddings collide and the output degrades into garbage.
SDPA shape enforcement: PyTorch / HF standard attention backends expect strict contiguous tensor layouts. I had to monkeypatch SDPA attention forwards to accept dynamically aligned, injected caches without triggering kernel panics or forced recompilations.
- Multi-page selection isn’t just string concatenation
With regular RAG, joining three relevant paragraphs is trivial string manipulation. With cached states, concatenating attention blocks requires careful offset management.
I initially tested an attractor-based dynamic router (SEH) with Winner-Take-All dynamics to suppress irrelevant pages before them. In practice, strict competitive inhibition suppressed pages that held complementary information (two separate pages that only made sense together). Switching to orthogonal Walsh-Hadamard bases with proportional reranking and a hard survival floor worked much better.
Practical takeaways
Tight architectural coupling: A crystal cache is locked to the specific model, layer count, head dimension, and quantization. If you switch models, you re-index from scratch.
The real bottleneck isn’t the theory: The hard part isn’t indexing or retrieval logic. It’s building a high-density binary serialization format that reads off storage faster than the GPU can run native prefill on raw tokens.
Thanks for the suggestions! LLMAE looks very close to what I had in mind, especially the idea of using continuous latent representations to preserve text information.
I’m still trying to understand how such representations could be injected directly into a retrieval-augmented LLM, rather than decoded back into text first. That seems to be the key challenge for reducing RAG prefill latency.
I’m not sure I can implement this myself yet, but you’ve given me a very useful direction to explore. Thanks again!
@PerryThe3rd The mental block usually comes from assuming an LLM strictly needs raw words as input. It really doesn’t.
Text is only needed for the very first step. Internally, each transformer layer converts those words into raw attention states (the KV-cache, basically its working memory). Once computed, the model no longer cares about the original text—it only uses these memory states to predict the next token.
In practice, here is how you connect the dots:
At index time: You compute a standard dense embedding for search, generate the internal memory state (the “crystallized” cache), and save that cache as a file on disk. In your vector DB, you simply store the embedding alongside the file path to that cached memory.
At query time: The vector DB matches the user query to the best embedding and gives you the file path.
Injection: You load that binary state straight into the model’s cache slots and feed only the user’s question as text. No source text ever touches the model.
To the LLM, it behaves exactly as if it had read the document a split second ago. It answers right away without having to re-read or re-process a single token from the source text. That’s how you bypass prefill entirely.
Interesting idea. One challenge is that hidden states and KV caches are often tied to a specific model, layer structure, and context position, which makes them harder to store and reuse than ordinary embeddings.
I’d look into work on recurrent memory, cache reuse, and approaches like Transformer-XL or Infini-attention. They may offer useful starting points, even if they don’t solve exactly the same problem.
That model-specific limitation seems central: a latent cache may need a key for model or checkpoint, layer, position or length, and tokenizer/config, plus invalidation when weights or context assumptions change. Is your idea to store transferable memory, or an opaque cache bound to one exact runtime? I work in marketing for BOMTOBE, makers of orke; this is a question from the discussion, not an implementation claim.
First, I apologize for the late reply. My regular work has been keeping me very busy, and as a non-expert, I needed extra time to understand the technical details in your comments. I used ChatGPT to help explain the concepts so I could better understand what you were describing.
My original question came from a misunderstanding, but also from an idea I was curious about. I assumed that the vector data used by RAG might be related to the internal representations an LLM uses during inference, and wondered whether we could retrieve and reuse those internal representations instead of the original text.
What I was trying to express with the word “weights” was not that the model’s actual parameters change when a prompt is entered. Rather, I was thinking that the same words might be represented and interpreted differently inside the model depending on the surrounding context and the user’s instructions. In other words, even if a document is preprocessed in isolation, its internal representations may differ from those produced when the document is processed in the context of a user’s question.
Thanks to your explanations and references, I now have a much better understanding of the distinction between retrieval embeddings and a model’s internal KV cache. I also understand that a KV cache is not a general-purpose memory representation that can simply be retrieved and reused without constraints. Its reuse depends on factors such as the model, token sequence, positional information, and how the cached states are integrated with the subsequent prompt.
I also learned that adding a user prompt does not necessarily invalidate all previously computed KV states. If the preceding token sequence remains identical, parts of the cache may still be reusable. However, making this work as a general RAG mechanism is more complicated than I initially imagined.
Thank you all for sharing your knowledge, practical experience, and useful references. Although I am not a researcher, this discussion has significantly improved my understanding of LLM internals and the challenges involved in using latent representations for retrieval.
Your final clarification gets at an important distinction: reusing an unchanged prefix is different from composing independently cached documents into a new context.
For a causal decoder, a document cached before the question can remain reusable when only the later question changes. But B cached after A generally is not the same state as B cached in isolation: its later-layer keys/values can depend on the preceding tokens. Fixing position offsets alone does not recreate that missing contextual computation. That also makes “exactly as if it had read the document” a stronger claim than I would make for a cache that has been quantized, sequence-downsampled and interpolated. Those can be useful approximations; they need a quality comparison, not just successful injection.
A small comparison I would find informative is: fresh text prefill of A+B, exact reuse of the same A+B prefix, and composition of separately cached A and B. Ask questions whose answers need both documents, including an exception in B to a rule in A; then swap document order. Report answer quality alongside cold-cache load/decompression/transfer time and total time to first token. That separates the benefits of exact prefix reuse from the additional approximation involved in composable memory.
I have not tested the implementations linked here; this is a suggested comparison, not a result about them. Your original efficiency question still seems worthwhile even once those two cases are separated.
The distinction in the last replies is helpful: exact prefix reuse is not the same as retrieving a compressed memory for a new question. For local use, I would frame the comparison around three costs rather than latency alone: storage and time for the cached state, accuracy loss from compression or from stitching states together, and the engineering cost of keeping the cache tied to one model version. A useful test would be the same set of documents answered from fresh text, from exact cached prefix, and from a retrieved compressed state, with answer quality checked on numbers, negations, and multi-document reasoning. That would show whether the savings come from reuse itself or from avoiding repeated reading of irrelevant text.