Aggressive RAG in small models: trading parameters for retrieval

· By José Manuel Román
Contents

This post belongs to the series on inference performance in small models. Its sibling piece, The roofline inverts in small models, explains why the compute-bound prefill is the bottleneck that shapes the whole discussion here. It is worth reading first: here we assume that adding more context is not free.

TL;DR

An SLM (say 1B–8B parameters) knows fewer facts than a 70B–700B model, simply because it has fewer weights in which to memorise them. But its ability to reason over text in front of it, following instructions, extracting, synthesising, comparing, degrades far less with size than its encyclopaedic knowledge does. The operational consequence is direct: use the SLM as a reasoning engine over curated context, not as a database. Move knowledge from the weights into the context via retrieval. The problem is that “aggressive retrieval” is often read as “put in lots of chunks”, and that collides head-on with three facts about SLMs: shorter context windows, worse use of long context (the lost in the middle effect is more severe the smaller the model) and a compute-bound prefill whose cost grows with the context length $C$, linearly in the projections and quadratically in attention. You cannot simply add tokens. Each retrieved token is paid for twice, in degraded quality and in TTFT, and the SLM is the worst placed to absorb either cost. The way out is not to retrieve less, but to retrieve better: reranking for precision over recall, compressing the context before injecting it, prefix caching the stable documents, semantic caching of answers, and structured output with external tools that stand in for internal knowledge. This post works through the maths and gives a TTFT number before and after compressing a context from 4,000 to 1,000 tokens on an RTX 4090.

The analogy: the open-book exam

Two students sit the same exam. The first has a prodigious memory: they have memorised the entire syllabus, paragraph by paragraph. The second has an ordinary memory, forgetting dates, confusing names, but is allowed to bring in a crib sheet.

If the second student’s crib sheet is a chaos of piled-up photocopies, they lose: they take too long to find what they are after, get distracted by irrelevant pages and run out of time. But if their crib sheet is excellent, trimmed to the essentials, reordered by relevance, with the important things at the top and no filler, then they not only do not lose: they often win, because they reason just as well as the first student and, on top of that, work from verified material instead of from blurry recollections they may be inventing.

The moral has three layers, and each maps to an engineering decision:

  • Memorising everything is expensive. The first student invested months. A large model invests parameters, and VRAM, and inference FLOPs, in memorising facts.
  • The crib sheet matters more than its size. A well-made one-page crib sheet beats ten badly organised pages. More retrieved context is not better context: the material’s precision beats volume.
  • Knowing how to search and synthesise is a different skill from knowing. It is the one the SLM retains. The whole strategy consists of leaning on that skill and subcontracting the memory.

The rest of the post is, essentially, how to build an excellent crib sheet under the constraint that the student (the SLM) reads slowly and gets tired with long texts.

The capacity argument: how many facts fit in the weights

Let us start by justifying the thesis with an order of magnitude, not with faith. How much factual knowledge really fits in a model’s weights?

There is a recurring empirical estimate in the interpretability and memorisation literature: a dense model can store on the order of 2 bits of memorised information per parameter before saturating (the exact figure varies by study and training regime; take it as an order of magnitude, not as a law). An 8B-parameter model then has an information storage ceiling on the order of:

$$8 \times 10^9 \text{ params} \times 2 \text{ bits/param} = 1.6 \times 10^{10} \text{ bits} \approx 2 \text{ GB of information}$$

And that budget is not only for facts: the vast majority goes on grammar, syntax, reasoning ability, code, formatting, and only a fraction is left for encyclopaedic knowledge. Compare it with the other side: a retrievable corpus of several million documents, a corporate wiki, a document repository, a technical knowledge base, easily takes up hundreds of GB to terabytes of text, indexed and queryable with millisecond latency. The asymmetry is two or three orders of magnitude in favour of the external corpus.

The conclusion is not that the weights are useless, they are where reasoning lives, which is the expensive thing to replicate, but that competing with an external index on fact capacity is losing by construction. A 70B model has ~9× more memorisation budget than an 8B one, but it is still negligible against the corpus. That is why the large model also does RAG in production. The difference is that the SLM needs it: without retrieval, its factual knowledge is too sparse and, worse, prone to hallucinating exactly in the gaps it did not memorise.

Where knowledge livesIn the weights (memorised)~2 GB of useful info in 8Bfixed, costly to update, hallucinates in gapsIn the context (retrieved)hundreds of GB – TB indexedfresh, citable, verifiable, no retrainingThe SLM as a reasoning enginereasons over the curated contextnot the database: the one who reads and synthesises it

The central tension: retrieving more is not putting in more

This is where most naive designs break. “Aggressive retrieval” sounds like a large top-k: if retrieval helps, retrieve 20 chunks instead of 5. But in an SLM that fails for two independent reasons, one of quality and one of cost.

(a) SLMs make worse use of long context

The lost in the middle effect (Liu et al., 2023) is well known: LLMs retrieve information placed at the beginning and end of the context better, and the information in the middle worse. What gets less emphasis is that the effect is more severe the smaller the model. An SLM has fewer attention heads, fewer layers and poorer internal representations for “tracking” a relevant fact buried at position 14 out of 20 chunks. Its nominal context window also tends to be shorter (4K–32K against the 128K+ of the large ones), and the effective window, the length beyond which quality collapses, is smaller still. Putting in 20 chunks does not mean the model reads all 20: it means it will probably ignore or misread the ones in the middle, while paying the cost of all of them.

(b) Prefill grows with the context and is compute-bound

This is the blow people underestimate. Prefill, processing the full prompt before emitting the first token, is the compute-bound phase of inference (unlike decode, which is memory-bound; the detail lives in The roofline inverts). Its cost grows with the context length $C$, and it determines TTFT (time to first token). More chunks → more prefill tokens → more TTFT and more compute cost per request. In an SLM, where prefill is proportionally more expensive relative to the model, this hurts especially.

The operational conclusion is uncomfortable but clear: you cannot compensate for fewer parameters simply by putting in more context. Each retrieved token is paid for twice, in degraded quality and in TTFT, and the SLM is the worst placed to absorb both costs. The way out is to retrieve less but better, and to compress what you retrieve.

The maths of prefill

Let us put numbers on “prefill grows with the context”. For a context of $C$ tokens, a transformer layer does two classes of work:

  1. Linear projections (QKV, attention output, FFN): each token is multiplied by fixed-size weight matrices. The cost is $O(C)$ in FLOPs, linear in the number of tokens.
  2. Attention ($QK^\top$ and the multiplication by $V$): each token attends to all the others. The cost is $O(C^2)$, quadratic in the number of tokens.

The total prefill cost per layer has the form:

$$\text{FLOPs}_{\text{prefill}} \approx \underbrace{a \cdot C}_{\text{projections}} + \underbrace{b \cdot C^2}_{\text{attention}}$$

with $a$ and $b$ constants that depend on the model dimension. For moderate contexts (a few thousand tokens) in an SLM, the linear term still dominates or is comparable to the quadratic one; the quadratic term becomes dominant at long contexts. What matters: if you compress the context $C \to C/k$, the linear term falls $\times k$ and the quadratic one falls $\times k^2$. Compression is the only lever that attacks both terms at once, and it attacks the worse one disproportionately.

Numerical example: TTFT before and after compressing, RTX 4090

Let us model TTFT as the time to process the prefill tokens at a given prefill throughput. Take an RTX 4090 (24 GB, Ada Lovelace) serving a quantised SLM, with a prefill throughput of ~5,000 tok/s (an illustrative figure; the real value depends on the model, the quantisation and the batch, so measure it, do not assume it).

Let there be a retrieved context of 4,000 tokens (8 chunks of ~500 tokens). Approximating TTFT as dominated by the context prefill:

$$\text{TTFT}_{\text{before}} \approx \frac{4000 \text{ tok}}{5000 \text{ tok/s}} = 0.80 \text{ s}$$

Now we compress that context to 1,000 tokens ($k = 4$). Prefill throughput is not constant with $C$, it drops a little at long contexts because of the quadratic term, but taking the conservative linear tokens/throughput approximation:

$$\text{TTFT}_{\text{after}} \approx \frac{1000 \text{ tok}}{5000 \text{ tok/s}} = 0.20 \text{ s}$$

TTFT falls from 0.80 s to 0.20 s, a $4\times$ reduction in the linear part. But the FLOPs arithmetic is even more favourable in the attention component: that part of the work falls $\sim k^2 = 16\times$. In practice total TTFT does not fall 16× because the cost is not purely quadratic at this scale, but the real reduction lies between 4× and something larger depending on how much attention weighed, and the aggregate compute saving (what pays the electricity bill and frees the GPU for another request) is substantially bigger than the plain 4× of the token count.

The argument generalises: compressing the context by a factor $k$ reduces TTFT by at least $\sim k\times$ and the attention cost by $\sim k^2\times$. For an SLM, where TTFT is often the SLA that matters, this is the difference between an assistant that answers instantly and one that feels slow.

The five levers to resolve the tension

The strategy is not “retrieve less and settle”. It is retrieve aggressively from the index and then distil aggressively what you retrieved before it reaches the SLM. Five levers, in order of application within the pipeline.

1. Aggressive reranking: precision over recall

The initial retriever (dense, sparse or hybrid) optimises recall: it brings back 50–100 candidates so as not to leave anything out. The reranker, a cross-encoder that sees the query and the document together, optimises precision: it reorders those candidates and you keep the 3–5 best. For an SLM this is not a luxury, it is structural: since the model uses long context badly, every chunk that goes in must earn its place. Better 4 chunks of very high relevance than 15 mediocre ones. The detail of hybrid retrieval and reranking is in Reranking and hybrid retrieval; here the rule is enough: maximise recall in the retriever, maximise precision in the reranker, and inject few.

2. Context compression: distilling the crib sheet

Once you have the best chunks, they still contain filler: padding sentences, redundancy, context irrelevant to the specific query. Context compression trims them before injection:

  • Extractive compression (LLMLingua / LongLLMLingua style, Jiang et al. 2023): a small model scores the perplexity or relevance of each token/sentence with respect to the query and removes the low-information ones, keeping the densest extractive subset. It reduces tokens without a second large generative model in the loop. LongLLMLingua adds position-aware reordering to mitigate lost in the middle.
  • Abstractive compression: a model summarises the retrieved chunks into shorter text. More aggressive in token reduction, but it introduces a generative step (cost and possible loss of fidelity).
  • Soft prompts / context distillation: compressing the retrieved context not into text, but into a handful of embeddings/soft tokens that the model consumes directly. It reduces the number of prefill tokens to a minimum, at the price of a trained, model-specific component.

The key point connects with the maths above: compressing what you retrieved by a factor $k$ reduces prefill tokens $\times k$, and therefore TTFT $\sim\times k$ and the attention cost $\sim\times k^2$. It is the lever with the best return when long context is the bottleneck.

3. Prefix caching of stable context

Not all the context changes between requests. System instructions, definitions, recurring reference documents, schemas: these are stable prefixes. Prefix caching stores the already-computed KV cache of those prefixes and reuses it, so that prefill only processes the new part (the query and the specific chunks). If 60 % of your context is stable, you save 60 % of that segment’s prefill on every hit. For it to work, the stable context must go at the start of the prompt (the KV cache is prefix-dependent) and it pays to maximise the hit rate; the hit-rate engineering detail is in Prefix cache hit rate. It combines especially well with RAG: retrieved documents that repeat across sessions get cached once.

4. Semantic caching of answers

A layer in front of the model: if a query is semantically equivalent to one answered before (embedding similarity above a threshold), return the cached answer and skip the model entirely, retrieval, prefill and decode included. In real workloads with long tails of repeated or near-repeated questions (FAQ, support), the saving is enormous because it eliminates the full cost, not just the prefill one. The trap is the threshold: too lax and you serve wrong answers to similar-but-different questions. The design is in Semantic caching for RAG.

5. Structured output and function calling: lean on tools, not on memory

The last lever changes what the SLM depends on. Instead of asking it to know a fact (its weak point), make it call a tool that knows it: a database query, an API, a calculator, a validator. Structured output (forcing JSON conforming to a schema) and function calling turn the SLM into an orchestrator that extracts arguments from the context and delegates the computation or the query. A reasonably capable SLM emits a well-formed tool call far more reliably than it recalls a specific fact. This reduces the pressure on parametric knowledge and on retrieval: for structured, fresh data (prices, inventory, states), querying beats retrieving text and beats memorising. The fundamentals are in Structured output and Function calling.

The complete pipeline

The five levers are not alternatives: they chain together. The flow, with the token counter falling at each step:

Retrieve aggressively, distil aggressively, reason cheaplyRetrieverhybrid, recall~80 chunksRerankerprecision5 chunks · 4,000 tokCompressionextractive k=41,000 tokPrefix cache+ semanticcacheminimal prefillSLMreasons +tool callsanswerThe prefill token counter falls along the pipelineretrieve: a lotrerank: 4,000 tok1,000TTFT on RTX 4090 at ~5,000 tok/s · 4,000 tok = 0.80 s → 1,000 tok = 0.20 sattention falls ~k² = 16× in that part of the compute

The order matters. Retrieving aggressively (high recall) before filtering guarantees that the right material is among the candidates; reranking and compressing afterwards guarantees that only the dense and relevant material pays the prefill toll; caching wraps everything so as not to repeat work. The SLM only sees the final crib sheet, short and ordered.

Implications for on-premise inference

The mental trap to avoid: treating the SLM as a large model with less quality. It is not. It is a different cost profile that rewards a different design. Three practical consequences:

  • The token budget is a first-class resource. With a large model with a 128K window, “putting in a bit more” is cheap relative to the model. With an SLM, every context token shows up in TTFT and in quality. Treat context size as a quantity to minimise subject to covering the answer, not to maximise.
  • The investment is worth it precisely because the model is cheap. Reranker, compressor and caches add complexity, but the model they serve is economical enough to run many replicas. The bottleneck shifts from the model to the data pipeline, which is exactly where you want it.
  • Retrieving does not replace adapting; they combine. For deep, recurring domain knowledge, adapting the SLM with LoRA (see the sibling QLoRA and aggressive multi-LoRA) can put part of the knowledge “into the weights” cheaply, reducing what has to be retrieved. Aggressive RAG and aggressive adaptation do not compete: the first gives freshness and citability, the second gives fluency and domain formatting. A good design uses both.

On the RTX 4090 (24 GB, Ada Lovelace)

The canonical scenario: a quantised SLM (4B–8B in INT4/FP8) fits with room to spare, leaving VRAM for a generous KV cache, essential for prefix caching, and for the reranker (a cross-encoder of a few hundred MB). The LLMLingua-style extractive compressor runs on a separate small model or on CPU. The TTFT calculation above (0.80 s → 0.20 s compressing 4× at ~5,000 tok/s) is representative of this card. The rule of thumb: if TTFT goes above your SLA, the first adjustment is to compress the context, not to change the model.

With 320 GB and native FP8 prefill is much faster, so the temptation is to relax the token discipline. That is not quite advisable: the lever changes from TTFT to aggregate throughput. Compressing the context not only speeds up each request but frees prefill compute to serve more requests per GPU; compute-bound prefill is exactly the resource that saturates first under load. Here prefix caching and semantic caching, shared across replicas, are what pay off most: at high QPS, the prefill work you avoid by caching is pure throughput you gain. The SLM is still the cheap reasoning engine; the difference is that you now run many in parallel and the data pipeline is what decides how many requests fit.

What we have not covered

  • Evaluating compression: how to measure that compressing $k=4$ does not throw away correct answers (faithfulness, answer recall over a set of questions with ground truth).
  • Query-aware versus query-agnostic compression: compressing before or after knowing the question changes what can be cached and what can be discarded.
  • Chunking and granularity: chunk size interacts with reranking and compression; that is left for the corpus curation post.
  • Multi-hop and agents: when a question requires several retrieval rounds, the token budget is split across hops and compression discipline becomes critical.

See also

References