I had a hunch about retrieval-augmented models and wanted to measure it rather than argue about it.
The hunch: if you are going to look a fact up and hand it to a model, it should be much easier for the model to use a fact placed right in front of it than to find that same fact buried in a long context. "Much easier" is vague, so I built the smallest experiment that could put a number on it.
Why it matters. Retrieval-augmented models are everywhere now: look something up, drop it into the context, let the model use it. The usual habit is to hand the model a big pile of retrieved text and trust it to find the part that matters. That pile is not free. It costs context length, it costs training compute and training data, and it leans on in-context recall, which recurrent models like Mamba are known to be bad at. So the question is worth real money: is the model actually learning from that haystack, or would it learn far faster from a single well-placed fact? If it is the second one, then the thing to invest in is putting the right fact in the right place, not training the model to be a better librarian. For someone building on one GPU instead of a data centre, that gap is the difference between a project that is tractable and one that is not.
This is the plain-language walkthrough. The code is the technical version.
The setup. A tiny language model, 1.86 million parameters, whose backbone is a stack of four Mamba-2 blocks. (Mamba-2 is a state-space model: it reads a sequence left to right and keeps a fixed-size running summary, rather than attending over the whole context like a transformer.) I train it on a made-up lookup task where the answer is a single token, and I measure one thing: how many training tokens it takes to learn the task to 95% accuracy. The pairings are random every time, so the model cannot memorise. It has to learn the skill, and "how long to learn it" is what changes by orders of magnitude between setups.
Everything trains in minutes on a single GPU.
Same model, same task, same optimiser. The only thing I change is the context in front of the answer:
- long: all K facts are laid out inline, and the model has to find the right
one.
FACT q1 v1 SEP q2 v2 ... qK vK Q q* A ? - inject: only the one needed fact is there.
FACT q* v* Q q* A ?
Both contain the answer. Both are scored on the exact same single token. The only difference is whether the model has to search.
| layout | facts (K) | training tokens to reach 95% |
|---|---|---|
| inject | 8 | 51,200 or fewer |
| inject | 32 | 51,200 or fewer |
| long | 8 | 3.7 million |
| long | 32 | 24 to 36 million |
Measured the same coarse way for both arms, handing over the one fact is about 70x cheaper to learn at K=8, and 470x to 690x cheaper at K=32. And that undersells it. The injected version converges so fast that checking only every 100 steps cannot catch the exact moment; measured properly (Chapter 3) it gets there around step 45, roughly 23,000 tokens, which puts the true K=32 gap closer to 1,000x to 1,500x. Either way, the injected version does not care how big the haystack would have been, because there is no haystack.
Why. The injected fact sits a few tokens before the answer, so the model just has to learn "copy the value that showed up next to the question." That is easy for a state-space model. The long version needs the model to stuff all K facts into its fixed-size running summary and then look one up by key at the end. That is in-context associative recall, which is exactly the thing state-space models are known to be worst at. The gap is not a quirk of my setup; it is that one setup asks for a hard skill and the other does not.
Run it: python experiments/01_inject_vs_long.py --full
Real questions often need more than one fact. "Who is the CEO's spouse?" is two
hops: find the CEO, then find their spouse. So I made the task two-hop
(e_t -> e_mid -> v) and asked whether injection still helps, in three flavours:
- inject_collapsed: the lookup was already solved for the model.
FACT e_t v ... - inject_2hop: the model is handed the chain, in order, and has to walk it.
FACT e_t e_mid SEP e_mid v ... - long_2hop: both banks of facts are inline and the model must chain them itself.
| layout | K | reached 95%? | best accuracy |
|---|---|---|---|
| inject_collapsed | 32 | yes, almost immediately | 1.00 |
| inject_2hop | 32 | yes, almost immediately | 1.00 |
| long_2hop | 32 | no | ~0.08 |
Walking a chain that is handed to you in order is just as cheap as being handed the answer. Doing the chain in-context is a different story: the long version never crossed 10% accuracy even after 202 million training tokens, roughly 2,900x the cost of the injected version, and it was not close to getting there.
The useful takeaway is a soft one about system design. A retrieval system does not have to pre-compute answers. It just has to return the relevant facts in chain order, and the model handles the rest cheaply. That is a much easier job to build than one that has to solve the question before the model sees it.
Run it: python experiments/02_two_hop.py --full
Chapter 1 said the injected fact is cheap because it sits right next to the answer. Chapter 3 tests that directly by pushing it away. I insert D filler tokens between the fact and the answer and watch what happens.
First, the baseline: with no filler, the model learns the task in about 45 optimiser steps. Tight and boring, which is the point.
Then I add filler:
| distance D | training steps to reach 95% | final accuracy |
|---|---|---|
| 0 | 45 | 1.00 |
| 4 | 250 | 1.00 |
| 16 | did not converge | 0.39 |
| 64 | never learned | 0.006 (chance) |
| 256 | never learned | 0.008 (chance) |
Four filler tokens make it 5.6x slower. Sixteen and it barely learns at all. Sixty-four and the signal is gone completely, stuck at random guessing.
Why. The model's short-range copy path reaches back only a handful of tokens (its convolution sees a window of four). Inside that window, copying is trivial. Past it, the fact has to survive in the running summary and be looked up at the end, which is the same hard recall skill from Chapter 1. So the cheap win is not really "injection." It is adjacency. Put the fact right before the answer and it is nearly free. Let it drift a dozen tokens upstream and you are back to paying full price.
Run it: python experiments/03_adjacency_decay.py --full
If you are building a retrieval-augmented system on a state-space (or similarly recurrent) model, three practical things fall out of this:
- Put the retrieved fact right before the answer. Not in a system prompt, not early in a long history. Adjacency is the whole win, and it decays fast.
- You do not need a smart retriever that pre-solves questions. One that returns the right facts in chain order is enough, and much cheaper to build.
- Do not count on the model to do multi-hop lookups in-context at this scale. It is bad at it, and no amount of training in my budget fixed that. Retrieval is not just cheaper here; it may be the only route that works.
None of this is a new architecture. It is a careful measurement of a known weak spot in state-space models, turned into design advice you can act on.
ssm_retrieval/
model.py the tiny Mamba-2 language model (with a CPU fallback for smoke tests)
tasks.py the synthetic lookup task in all its layouts (1-hop, 2-hop, padded)
train.py the shared training loop and the "steps to learn it" measurement
diagnostics.py accuracy plus confidence/margin, for reading eval curves
experiments/
01_inject_vs_long.py Chapter 1
02_two_hop.py Chapter 2
03_adjacency_decay.py Chapter 3
results/ JSON results land here at runtime
python -m venv .venv && source .venv/bin/activate
pip install -r requirements.txt
pip install mamba-ssm # needs a CUDA GPU; this is what makes the numbers real
python experiments/01_inject_vs_long.py --full
python experiments/02_two_hop.py --full
python experiments/03_adjacency_decay.py --full
Each script prints the table from its chapter. Use --smoke for a tiny, quick
version. Without a GPU the model uses a slow pure-PyTorch fallback so --smoke
still runs and you can see the pipeline work, but the real numbers need the
Mamba-2 kernels, which are CUDA-only. On a single modern GPU the full sweeps take
on the order of an hour or two; the two-hop long-context runs are the slow part.
- This is a synthetic lookup task with a small vocabulary, chosen so the experiment is cheap and the mechanism is clean. It is a controlled measurement, not a benchmark on real text.
- One architecture, one size (four Mamba-2 blocks, 1.86M parameters). The ratios are a structural argument, but I have not swept model size or state size here.
- The one-hop tokens-to-95% are measured on a coarse every-100-steps schedule, so the injected numbers are upper bounds; Chapter 3 re-measures the fast ones properly. Where a number is a bound I have said so.
- Seeds: the one-hop and adjacency results use two to three seeds each, and every arm of the two-hop sweep uses two seeds. The two seeds agree closely throughout, including the hard K=32 long-context arm, which lands at 0.06 and 0.08 (both well below the 10% mark), so the failure there is not a one-off.
I pulled this out of a larger private research project on training-efficient
state-space models. The model core here is the public mamba_ssm Mamba-2, the
same block the private version used under the hood, and re-running everything on it
reproduced the private results closely (the one-hop K=32 gap lands at 467x to 694x,
and the adjacency-decay curve is identical). That is the point: the finding is
about the architecture, not about my particular copy of it. Released under the MIT
license (see LICENSE).