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”?
During indexing:
Document → LLM → LLM-native latent representation → storage
During inference:
Query → retrieval → latent representation → LLM
The goal would be to bypass the normal tokenization/prefill/semantic-processing cost of retrieved documents.
The main challenges I see are:
- How to compress the latent representation without losing numbers, conditions, negations, names, etc.
- How to inject the latent memory into the appropriate Transformer layers.
- Whether the representation could be stored on NVMe/SSD and loaded only when retrieved.
- Whether this could substantially reduce RAG prefill latency.
I am not claiming this is a new idea. I would actually like to know what existing research is closest to this approach.
Is there existing work on treating an LLM’s internal representation or KV-like state as persistent external RAG memory?
You could use a sentence autoencoder for this—such as Meta’s SONAR—or check out this new paper on how to build your own: [2609.27248] Repurposing Pre-trained LLMs as High Fidelity Continuous Text Autoencoders To store the embeddings on disk rather than consuming large amounts of RAM, you can use a tool like the FAISS vector database.
Thank you for these helpful suggestions!
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 repo is open-source if you want to look at the implementation: Chromatix RAG Engine and Chromatix-Pixel-Standard
The short version: bypassing prefill by loading stored attention states works, but making it practical hits three walls very quickly.
1. 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.
2. 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.
3. 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 loading 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!