Preloading Knowledge Into a Model Instead of Retrieving It (Part C)

Recap

In Part 13, we prefilled a corpus once and answered from the stored cache. In Part 14, we tried to make that cache smaller, and mostly could not.

Compression splits into two families.

• One shrinks the stored numbers, which is query-agnostic by construction and can therefore run offline, but it ran out of room between 2x and 4x on a real model.
• The other drops tokens judged unimportant, which reaches far higher ratios but has to decide what matters, and the dominant methods answer that using an observation window sitting exactly where the user's query sits.

That is fatal for preloading, since a score computed from the query cannot be computed before one exists.

We measured how much it matters and found the answer depends on depth, with kept token sets agreeing 89 percent of the time at layer 1 and only 37 percent at layer 16. Compressing with one question's selection and then asking a different question produced a fluent answer naming the wrong depot and the wrong person.

Both approaches in this part require a training run, and both escape that ceiling because of it. One changes the model so that independent blocks stop being a problem. The other abandons the cache the model produced and trains a much smaller one to behave as though the corpus were present.

We then close the series with what running any of this in production actually involves.

Let's begin!

As always, every notion will be explained through clear examples and walkthroughs to develop a solid understanding.

Let's begin!

Modular preloading

Naive preloading treats the corpus as one indivisible block. That works when the whole corpus fits, and every query needs all of it.

But neither of those two conditions is practically feasible for long documents.

The modular version prefills each passage separately, so passages can be assembled in any combination at request time. This is the same construction Part 12 examined and found broken, for three specific reasons.

There is a different way to address it. Instead of repairing the cache after the fact, change the model so the damage never occurs.

Making blocks independent by design

Block-Attention takes this route.

The input is divided into blocks, each block computes its keys and values without attending to any other block, and only the final block, containing the user query, attends across everything.

The model is then fine-tuned to operate under this attention pattern.

Because the training distribution now matches how the cache is assembled at inference, the mismatch that caused quality loss in Part 12 no longer exists.

Each passage becomes a genuinely reusable unit, and its cache can be computed once and dropped into any request.

The implementation involves three pieces, which are block segmentation, position re-encoding so each block's keys carry the position they occupy in the assembled input, and the fine-tuning step itself.

That middle piece is the same repair we performed by hand in Part 12.

A block prefilled on its own has its keys rotated for positions starting at zero. Dropping it into the middle of an assembled input means re-rotating those keys to the positions the block now occupies.

Block-Attention does not escape that correction in any way, but rather it moves it into the serving path so it happens on every assembly automatically.

The fine-tuning step addresses the other failure, which is the missing cross-attention.

Part 12 corrected position and found the answer still wrong, because a block prefilled alone never saw the blocks before it.

Selective recomputation solved that by recomputing a slice of tokens at request time:

Fine-tuning solves it by teaching the model that blocks are meant to be independent, so there is no longer a discrepancy left to repair.

If you see these two approaches side by side, the two routes split the problem identically and solve with different costs involved in both.

Position is fixed arithmetically in both.

Cross-attention is either bought back with computation on every single request, or trained away once and never paid for again.

Why the training step is necessary

The same paper has also provided a useful result about the training-free alternatives, which also answer a question we left open in Part 12.

The authors re-implemented Prompt Cache (the earlier approach that precomputes block caches without fine-tuning), and found its results identical to their own no-fine-tuning baselines.

They also tested Parallel Context Windows, a method that stretches the usable context by cutting the input into several windows which each attend only within themselves, while the test tokens attend across all of them.

Its purpose is extending context length rather than making retrieval cheaper, and when its attention pattern was applied directly to retrieval, the accuracy landed below those same weak baselines.

This is direct evidence that simply precomputing independent blocks and hoping for the best does not work in retrieval settings. Something has to give, either recomputation as in the selective approaches from Part 12, or a change to the model itself as here.

In summary, modular preloading buys order-independence, which naive preloading cannot offer. Getting it without quality loss means either paying recomputation at request time or paying a fine-tuning cost once.

Trained preloading

Every approach so far kept the cache the model naturally produced, either whole, in blocks, or with pieces removed.

The last approach in the spectrum trains a cache instead.

The idea comes from Cartridges paper at Stanford, which we discussed briefly in the last part.

Rather than storing the cache a corpus produces, they train a much smaller cache to behave as though the corpus were present.

The setup

The starting observation of the paper is the same as the one this chapter opened with.

Users issue many queries against the same corpus, whether that is a codebase, a set of filings, or a chat history. Any cost paid once per corpus is amortized across every query that follows.

That makes an offline training run affordable in a way it would not be per request. The authors point out that this training can run on idle or underutilized compute, such as overnight when user load is low.

The artifact produced due to training is called a cartridge, and at inference time it is loaded in place of the corpus cache.

Structurally, it is not a compressed copy of anything.

Instead, a cartridge is a set of trainable key and value vectors, and a size parameter fixes how many positions it holds.

That parameter is a hyperparameter you specify when deciding how much memory the artifact is allowed to occupy, and it is chosen up front.

The authors describe the construction as a simplified form of prefix tuning, which is the established technique of learning a short sequence of vectors that is positioned in front of the real input and steers a frozen model's behaviour.

Do note that the trainable key and value vectors mentioned above are not initialized randomly.

They start as the real key-value cache of the first stretch of the corpus, as many tokens as the cartridge has positions, so training begins from something the model already reads as text rather than from noise.

They did because when they measured the approach's effectiveness by starting with random initialization, it reached 29.9 percent on the LongHealth benchmark, against 55.3 percent when initialized from actual corpus tokens.

Why the obvious training approach fails

The natural way to train a small cache is next-token prediction over the corpus, which is how language models are trained in the first place.