Skip to main content

Research

Retrieval is the attention.

lab358 serves language models whose attention is a top-k lookup over an index of key–value memory — memory that lives across GPU, CPU, and NVMe. There is no fixed context limit. Here is the published work we build on, what we integrate that is new, and how the approach differs from the common alternatives.

The idea

How it works

In a conventional model, every new token attends over the whole conversation held in a context window — and that window has a fixed ceiling. lab358 replaces that with retrieval: at each step, attention finds the top-k most relevant key–value entries and attends over those.

Those entries do not all have to sit in GPU memory. Recent context stays exact on the GPU; the long tail is retrieved from an approximate-nearest-neighbour index whose key–value memory is resident across a tier ladder — GPU, then CPU, then NVMe. The two are merged into a single result. Because the memory can spill to storage, the amount of context a model can draw on is bounded by disk, not by a fixed window — hence no fixed context limit.

Position is carried by a learned, per-head distance signal rather than a fixed positional scheme, which keeps precomputed memory position-free — so a document's key–value memory can be computed once and reused wherever it is needed, without re-encoding.

Lineage

The shoulders we stand on.

lab358 is an integration of ideas from a rich body of published research. A selection of the work most directly related to how our attention and retrieval are built:

  • RetrievalAttention: Accelerating Long-Context LLM Inference via Vector Retrieval

    Liu, Chen, et al. — arXiv:2409.10516, 2024

    Retrieves the most relevant key–value vectors per step from an approximate-nearest-neighbour index, and names the query-versus-key distribution mismatch that a retrieval index has to overcome.

    Read the paper →
  • Unlimiformer: Long-Range Transformers with Unlimited Length Input

    Bertsch, et al. — NeurIPS 2023, arXiv:2305.01625

    An early, clean statement of the core idea: let each attention head retrieve its top-k keys from an index instead of attending over everything.

    Read the paper →
  • Landmark Attention: Random-Access Infinite Context Length for Transformers

    Mohtashami & Jaggi — NeurIPS 2023, arXiv:2305.16300

    Makes block selection native to the attention mechanism itself rather than a bolted-on external retriever — the same philosophy of no fixed context window.

    Read the paper →
  • Query-Aware Sparsity for Efficient Long-Context LLM Inference (Quest)

    Tang, et al. — ICML 2024, arXiv:2406.10774

    Estimates which key–value regions matter for the current query and loads only those — the query-aware selection principle at the heart of our index.

    Read the paper →
  • Fast Accurate Billion-point Nearest Neighbor Search on a Single Node (DiskANN / Vamana)

    Subramanya, et al. — NeurIPS 2019

    The graph-based, disk-resident nearest-neighbour algorithm that makes billion-scale search practical on one machine — the retrieval substrate our tiered memory is built on.

    Read the paper →
  • RoarGraph: A Projected Bipartite Graph for Efficient Cross-Modal ANNS

    Chen, et al. — VLDB 2024, arXiv:2408.08933

    Builds the search graph under the query distribution to fix recall on out-of-distribution queries — how we keep retrieval accurate when queries and stored keys differ.

    Read the paper →
  • Train Short, Test Long: Attention with Linear Biases (ALiBi)

    Press, Smith & Lewis — ICLR 2022, arXiv:2108.12409

    Encodes position as a distance bias on attention scores rather than a positional embedding — the basis for the learned, per-head distance signal we use.

    Read the paper →
  • The Impact of Positional Encoding on Length Generalization in Transformers

    Kazemnejad, et al. — NeurIPS 2023, arXiv:2305.19466

    Shows decoder-only transformers without positional encodings generalise to longer inputs as well as, or better than, the alternatives — the evidence behind our no-fixed-position design.

    Read the paper →

A selection, not the full bibliography. Each ingredient above is well established in the literature; what lab358 contributes is the integration.

What's different

The integration is the novelty.

Each piece has strong prior art. Bringing them together into a single trained model and serving stack is where lab358 is distinctive.

Retrieval is the attention — and it is trained in

Attending means retrieving the top-k most relevant key–value entries. This is the model's native mechanism, learned during conversion — not a training-free patch layered onto a conventional model at inference time, and not retrieval bolted on in front of the prompt.

A query-aware search index

The nearest-neighbour graph is built to stay accurate for the queries attention actually issues, so the entries it returns are the ones the model would have chosen — closing the gap between approximate search and exact attention.

Memory across a hardware ladder

Key–value memory is resident across GPU, CPU, and NVMe tiers, with recent context kept exact on the GPU and the long tail retrieved from the index — so total context is bounded by storage, not by a fixed window.

Precompute a document once, reuse it everywhere

A document's key–value memory is computed a single time and reused across chats, agents, and users — no re-encoding and no prefill to redo, even days apart. Reuse is a first-class product primitive, not an accident of caching.

Approach comparison

How this differs from the common approaches.

A comparison of techniques, not products. Each row describes a way of giving a model access to more context, and where it runs into a wall.

The common approaches — and where each hits a wall

Retrieval into the prompt

Search a separate store for relevant text, then paste that text into the prompt for the model to re-read.

The wall: The retrieved text still has to fit inside a fixed context window, and it is re-encoded from scratch on every request.

Paged or bounded key–value attention

Keep the key–value cache in efficiently managed blocks, reusing exact shared prefixes where they occur.

The wall: Reuse is limited to verbatim prefixes, and attention still runs densely over a bounded cache — it does not selectively retrieve older material.

Fixed-window long context

Enlarge or slide the attention window, keeping a band of recent tokens (and sometimes a few initial ones) exact.

The wall: Content that falls outside the window is discarded, so the model cannot answer over what it has dropped.

The resolution

lab358 — top-k attention over tiered indexed memory

Attention retrieves the top-k most relevant key–value entries from an index that spans GPU, CPU, and NVMe, merging exact-recent and retrieved-distant memory into one result.

The result: No fixed context window and no re-encoding: memory is retrieved into attention directly, and precomputed document memory is shared across uses.

Build on durable context.

The research, in a product you can run — hosted on lab358 Cloud or in your own AWS account.

Join the waitlistTalk to us