Idea
Core idea of context stitching is precomputing of KV values for different chunks of text, then to be loaded at inference time when required, and put together to form final context. Problem for this is assumed interdependence of chunks, since in vanilla attention, KV values for next chunk of text depends on previous chunk, which should create a bit of a problem.
Recap of slippery attention
In the previous post, “You don’t need to pay that much attention”, I described rather interesting property of attention mechanism to not break generation of text, if we replace vanilla attention pattern for already known tokens, with new sparse attention pattern, which shows that unless we need to predict a new token, we can actually skip a lot of computations.
Naturally, this is what enables the context stitching, since now we prove that KV values are mostly independent for chunks of text when there’s no new predictions to make.
In the original post, I’ve made an example of slippery attention with triangular attention pattern, but we’re actually not limited to that. It was just a display that even with such drastic changes in the attention pattern, LLM still remains capable of many tasks, with severe degradation only on some tasks.
Context stitching
What now we can do, is to take advantage of this property by precomputing KV values with normal attention pattern in “precompute time”, and then, at inference time, load these KV values into memory, to make predictions. As was shown in “You don’t need to pay that much attention” post, we will need to have full attention mask only for newly generated token.
Last problem for context stiching remains is positioning of these chunks in the context. For that, I will abuse property of a transformer to not care about continuity of positional embeddings. Basically, we compute KV values for each chunk starting at predefined positions, after which, context will look like it has multiple chunks decently spaced out from each other, with new tokens generated at a decent distance from last predefined token. That’s the main limitation of this, model should be taught to handle big contexts, as much as you need. Then, it’s only a matter of figuring out where you would want to put your chunks inside the context, and then, precompute KV values.
Theoretically, for prefill, we would need to load KV values for system prompt from memory/storage, compute user request KV values using slippery attention (with triangular mask, to save on compute), and only then predict new token with full mask. As an example, let’s look at following sequence:
<s>Use following table to help you answer the question:
2+2=4
2+3=5
Question: 2+2=
Answer:
Different colors of text are used to distinguish between different chunks of text and their origin. Blue and green are different chunks of text, loaded from storage, magenta comes for system prompt, red is for user request.
So, red is calculated with triangular mask, rest gets precomputed with with local full-attention mask. So, for sequence above, let’s build an (approximate) attention mask to show what I mean:
Notice how last token, attributed to system prompt, has attention to entire context. This is required to properly predict next token. All following generated tokens will attend to everything in the context, but note how system prompt, and chunks, all can have precomputed KV values since they don’t attend to each other. With little exception of token at index 0, which is also called “attention sink”(Xiao et al. 2024), but that can also be precomputed, since we always know what token it would be: it’s a sequence start token. User prompt only uses cheap “triangular” mask, but you can go as far as you want with this, we’ve saved enough here.
You might find it a bit similar to what’s happening in Quiet-STaR(Zelikman et al. 2024) paper, where they also precompute KV values for different chunks of text, but it’s done with motivation of running inference of different tokens in parallel. Here, however, motivation is precomputation of KV values to save time during inference. In example above, in order to get 4, we only need to compute KV values for single token, which will be just 4. Rest is just loaded. 25 vs (25*25)/2, isn’t it neat?
For now, this is just an idea, there’s no implementation YET. Will have to do it a bit later once I finish with other stuff.
Cite this as
@misc{avietisov2024contextstitching,
title={Unnatural abilities of context stitching},
author={Hlib Avietisov},
year={2024},
howpublished={Online, available at \url{https://havietisov.github.io/posts/context-stitching/}},
}References
Xiao, Guangxuan, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. 2024. “Efficient Streaming Language Models with Attention Sinks.” https://arxiv.org/abs/2309.17453.
Zelikman, Eric, Georges Harik, Yijia Shao, Varuna Jayasiri, Nick Haber, and Noah D. Goodman. 2024. “Quiet-STaR: Language Models Can Teach Themselves to Think Before Speaking.” https://arxiv.org/abs/2403.09629.