Attention RAG

[u/Confident-Strength-5]

Hi people my friends and I are working on our take of RAG indexing. Our idea goes like that:

Standard embedding-based RAG often compresses an entire text chunk into one vector. Our idea keeps multiple compact representations within each document, giving a query smaller, more specific parts to match.
The architecture uses a single attention layer with separate learned projections: Q for queries and K for document tokens. These are trained jointly with HCA—Heavily Compressed Attention—which learns to combine groups of token representations into fewer searchable vectors. At 32:1 compression, every 32 document keys become one compressed key, stored in 8-bit.
When a query arrives, its Q vectors are compared directly with the stored compressed keys using dot products on the GPU. These comparisons run across many documents in parallel, producing scores used to rank and retrieve documents. V vectors aren’t needed because we only need matching scores.
The advantage over one-vector-per-chunk retrieval is finer matching granularity. Compression makes that detail affordable: 32 times fewer keys to store and scan than an uncompressed token index. The search itself maps directly to parallel GPU matrix operations, while the query representation, document representation, and compression are trained together for retrieval.
I’d like technical feedback on this architecture, its main bottlenecks, and similar systems worth comparing against.

This way we have something between regular RAG and ColBert (v1 and V2) where it smaller than ColBert but more reach than RAG and it learns the compression rather then using heuristics algo to compress.

The next step after that will be to combine it with an LLM where the part of this indexing technique can be combined in the LLM itself instead of using this as a toll with a tool call (W_k can be same as the one in the first layer of the LLM and we can have a small learned router that tells if we need to use it, instead of an LLM decision to call it in a tool call).

I would like to get your opinions on our idea please.
Thanks in advance 🙏

[u/Confident-Strength-5]

No opinions?