Retrieval on real text needs a better search

The retrieval circuit learned a synthetic lookup perfectly (the previous post). The real question is whether a library of text helps the model predict new text. This post is that test, run four ways. Along the way came one diagnosis, one fix that worked and one result we had to take back.

01First run: a bug

FineWeb-Edu bytes, a 5.2M model, a library of 200,000 chunks of 64 bytes from the training documents. Retrieval never returns chunks of the document being read.

Step	Library attached	Library removed
1,000	1.935	1.929
2,000	1.785	1.783
3,000	6.122	1.727

At step 3,000, scoring with the library attached blew up to 6.122 . The cause was a chain of three effects:

The search fetched 16 candidates, then dropped those from the document being read.
The nearest chunks are usually from that same document. So in training, the read often came back empty. The circuit's running output scale settled low.
Held-out documents are not in the library at all. At evaluation every read came back full. Divided by that low scale, it swamped the model.

The fix moves the same-document rule inside the search, so every query gets 4 results from other documents. A test now covers the case where the document being read fills most of the library.

02Then: no gain

With the fix, on the , training with retrieval is stable and no longer hurts. It does not help either:

Run	Library	Attached	Removed	Training time
Web text, hybrid with retrieval	training chunks	1.615	1.614	1,002 s
Web text, hybrid without (control)			1.613	465 s
Web text, hybrid with retrieval	25,795 Wikipedia articles	5.969	5.968	1,036 s
Web text, hybrid without (control)			5.970	571 s

The second pair is in , with the 32k tokenizer. Random chunks of web text are rarely relevant to the next bytes of a new document. Even a Wikipedia library gave nothing on general text, at 1.8 times the training time.

03Diagnosis with a knowledge task

So we built a task where the library is known to hold what the model needs. Each Wikipedia article is split in half. The first half goes into the library. The model reads and predicts the second half, which keeps mentioning names and facts from the first. It trains on 90% of the articles. It is scored on the other 10%, with a fresh library made of their first halves, so it has to retrieve rather than memorize.

Retrieval still gave nothing. To find out whether the search or the read was failing, we added an : a diagnostic that always looks inside the right article and returns its 4 chunks closest to the query.

Fig. 01

13.2%

right article in the top 4 (measured, seed 0)

Three ways to search the same library. Each dot is one chunk of a held-out article's first half, placed by its retrieval key. Orange dots belong to the article being read; hollow dots belong to other articles. The ringed point is the query. Lines go to the 4 chunks the search returns. Switch the search to see how often it lands in the right article (the measured hit rates, seed 0). Positions are illustrative; hit rates are measured.

With the oracle, retrieval helps: 5.242 bits per token against 5.282 for a model trained without retrieval, a gain of 0.040. At step 1,000, reading with the library beats reading without it by 0.145.
The learned search finds a chunk of the right article among its top 4 only about 12% of the time. A REALM-style training loss for the search (each chunk should find its neighbour) did not raise it.

The read works. The search does not, yet.

04Rare-word keys

A search that starts from nothing has to learn what "related" means from the language-model loss alone. Names, places and technical terms are the easiest signal of relatedness. A search that knows nothing yet ignores them. So each retrieval key got a second part: a fixed random code for every token in the chunk, weighted by how rare that token is in the library (its inverse document frequency). Chunks that share rare words get similar keys from the first step. The learned part trains as before.

Fig. 02

Learned keys13.2%

Learned keys + neighbour loss12.4%

Learned keys + strong neighbour loss12.9%

+ rare-word keys, seed 031.6%

+ rare-word keys, seed 129.0%

Oracle search (100% by construction)100.0%

How often the search returns a chunk of the right article among its top 4, on held-out articles at step 3,000. Learned keys alone stay near 12% with or without the neighbour loss. Adding rare-word keys raises it to about 30% on both seeds. The oracle is 100% by construction.

On seed 0, this gave the first retrieval gain on real text: 5.256 bits per token against 5.282 for the no-retrieval model, 0.026 better, 65% of what the oracle gets. We reported it, then queued seed 1.

05The correction

Seed	Rare-word keys, library attached	Same model, library removed	No-retrieval model
0	5.256	5.260	5.282
1	5.281	5.285	5.279

What holds on both seeds:

Rare-word keys raise the search's from about 12% to about 30%.
The library helps the model it was trained with, by 0.004 bits per token at test time, on both seeds.

What does not hold yet is a net gain over a model trained without retrieval. At 13M parameters and 3,000 steps, retrieval works mechanically. The search is much better. The gain on real text is not established. The oracle shows 0.040 bits per token is there to take. Getting it means a search that finds the right article far more often than 30% of the time.

06What we will try next

These are on the ExeaLM to-do list. Anyone can pick one up:

A better search: better pooling of a chunk into its key, hard negatives (related but wrong chunks) when training the search, and a bigger encoder.
Larger models. Published retrieval gains on ordinary text came with much larger models and libraries.
A retrieval eval: questions whose answers are in the library but not in the prompt.