Preloading Knowledge Into a Model Instead of Retrieving It (Part B)
Part 14: Shrinking a preloaded cache and the problems with those approaches, covered with implementations
Recap
In Part 13, we moved the retrieval work earlier.
Instead of waiting for a query and then reading the corpus, we read it once, kept the key-value cache, and answered from that cache afterwards.
Avi Chawla
We placed the idea on a spectrum, from ordinary retrieval through to a trained cache artifact, with each step moving more work offline and keeping a different artifact.
We then built the system, prefilling a corpus once and restoring it per query. The cache memory formula predicted the artifact to within 0.09 percent on a real model, and all eight answers were identical whether the corpus was prefilled fresh or restored from disk.
We also saw that preloading has many more benefits than just speed. A retriever answered a multi-hop question with the wrong name until it was fetching 40 percent of the corpus, because the section holding the second fact shared no vocabulary with the question.
We also understood two limitations from that part.
The advertised context window is not the ceiling, since a fact written in the question's own words was found at every length we tried while the same fact reworded started failing at 512 tokens.
And renting the cache from a provider turns preloading into a break-even calculation, where holding time sets the bar and idle entries still bill.
In this chapter, we will look at what happens when you try to make that stored cache smaller.
A preloaded cache is large. The 636-token handbook from Part 13 produced a 15.6 MB artifact, which scales to roughly 2.46 GB for a 100,000-token corpus on the same small model. Shrinking it before storage seems like the obvious next move, and there is a substantial research literature on doing exactly that.
We'll start with the two families in the literature, one that shrinks every stored number and one that removes tokens entirely, and understand why only one of them can run before a query exists.
We'll then examine the methods that dominate the second family, namely SnapKV, PyramidKV, and H2O, and see that all three decide what to keep using the query itself.
From there, we'll look at what happens when these methods are evaluated in the setting preloading actually operates in, and what a genuinely query-agnostic method has to look like instead.
Finally, we'll ask how far compression can be pushed at all, and find a ceiling that two completely different approaches run into from opposite directions.
As always, every notion will be explained through clear examples and walkthroughs to develop a solid understanding.
Let's begin!
Compressed preloading and the query-agnostic requirement
A preloaded cache is large.
Shrinking it before storage seems like an obvious next move, and there is a substantial research literature on doing exactly that.
Most of that literature cannot be used for preloading and understanding why that's the case is an important thing, because the reason is structural rather than a matter of tuning.
How cache compression usually works
Two broad families of approaches exist, and they behave differently under preloading.
The first shrinks each stored number.
More specifically, keys and values are held at reduced precision, so a cache kept at 8 bits instead of 16 occupies half the space.
Nothing is discarded, i.e., every token remains present, and the loss is arithmetic rounding spread evenly across the cache.
These types of approaches are naturally query-agnostic, since rounding a number does not require knowing at info about the query that may be asked.
The second family removes tokens entirely, keeping the cache entries classified as important and discarding the rest.
This reaches far higher compression ratios than precision reduction, because dropping a token removes all of its keys and values across every layer.
It also raises a question that the first family never faces, which is how importance gets computed in the first place.
Precision reduction is bounded by the format itself. Caches are normally held in float16, so dropping to 8-bit integers is a factor of two, and aggressive 4-bit schemes reach a factor of four before rounding error becomes visible in the outputs.
Since a number only has so many bits to give up, there's a ceiling on how far we can go before we start losing accuracy
Token eviction has no equivalent ceiling. Keeping one token in four is a factor of four, keeping one in ten is a factor of ten, and nothing in the method stops you going further.
This does not remove precision but rather evidence, so the quality question becomes whether the discarded tokens were ones some future tokens may have needed.
That question is where preloading runs into trouble.
Retention of 25 percent means three-quarters of the cache entries were evicted, which is a compression factor of four.
Lower retention means more aggressive compression, so retention and compression ratio move in opposite directions.
The prominent methods that answer this do it usually in the same way.
SnapKV starts from an observation about how attention behaves.
Take the last few dozen tokens of the prompt, which the authors call the observation window, and look at where those positions send their attention.
Each attention head turns out to focus on a fairly consistent small set of earlier tokens, and it keeps focusing on that same set as generation proceeds.
That consistency helps make the window usable as a predictor.
Every position inside the window carries an attention weight over every earlier token, so adding those weights down the window gives each earlier token a score, which is how much the window collectively cared about it. Rank the prefix by that score, keep the top slice, and evict everything else.
Two details are important to decide to see if this works well.
- The scoring and selection happen separately for each attention head rather than once for a whole layer, because different heads attend to different things and one shared choice would serve none of them properly.
- The winning positions are also pooled into contiguous clusters rather than kept as scattered individual tokens, which stops the method from shredding a passage into disconnected fragments that no longer read as a span.
PyramidKV keeps that scoring machinery and changes how the budget is spread across depth.
Its starting observation is that attention is scattered widely across the input in the lower layers, consolidates as you move up the stack, and ends up concentrated on a few critical tokens in the highest layers.
The authors call this pyramidal information funneling.
If that is how attention behaves, then a uniform per-layer budget is wrong at both ends.
A high layer attending to a handful of tokens does not need a large allowance.
A low layer spreading its attention broadly is ideally where a small allowance throws away information that was genuinely in use.
PyramidKV therefore allocates more cache to the lower layers and progressively less as it climbs, which is the pyramid the name refers to.
H2O's (short for Heavy-Hitter Oracle) starting observation is that attention mass is distributed very unevenly.