Clip transcript
stays fixed throughout the computation. To understand the utility of caching in auto-regressive models, let's look at a world without it. You know the drill. We feed the prompt followed by the special beginning of sequence marker. Internally, the LLM produces contextual embedding X for each input position. Then, it produces an embedding for the next token based on everything to its left. This is causal attention. On the next iteration, the predicted token is appended to the input. Now, a naive implementation with no caching would discard the internal state between iterations. So, all the input embeddings would have to be recomputed from scratch. So, it's clear that a cache would save a lot of duplicated work. Once input embeddings are produced, they are reused across steps, and the cache grows as more tokens are predicted. Now, you might have heard of the more specific name KV cache. In practice, instead of storing a single copy of the token embeddings, we store two linear transformations called keys and values. These are the model activations specific to the transformer architecture. KV caching is feasible because of the causal attention. Since tokens only depend on their left context, future predictions don't alter their embeddings, so they're stable across iterations. But, that is not the case for diffusion models. Here's a typical input where the response is fully masked. The LLM computes embeddings for the entire context window. Let's assume tokens 3 and 5 are selected to be unmasked into who and there. The refined response becomes the input to the next diffusion step. Now, what happens to the internal state? Naively, you might be tempted to only recompute X3 and X5. They were originally based on the masked token, and now they're outdated. But, they're not the only ones that change. Let's look at the first token. Because of the bidirectional attention in a diffusion model, it depends on everything, including that now outdated tokens. So, X0 is invalidated as well. So, the virus basically spreads to the entire context window. KV caching as implemented for auto-regressive models is simply not possible for the fusion models. And that's a big setback since KV caching can speed up inference on auto-regressive models by anywhere between three to 30 times. However, this doesn't mean the fusion inference has no redundancies to be optimized. They're just not as structurally obvious. While all internal token embeddings do get updated at each denoising step, not all updates have the same magnitude and impact. Consider the prompt portion. While its embeddings do depend on the ever-changing response, its fundamental semantics remain the same. Empirically, studies have shown that the prompt embeddings change very little across consecutive denoising steps. That's why some papers choose to cache them right from the start. I'm calling this cache approximate because the values are only an approximation of the fresh ones we're avoiding to recompute. Some papers like DALL-E um cache refresh the cache at regular intervals, say 100 steps, to mitigate the staleness. Researchers also noticed patterns in how the response tokens evolve. For instance, unmasking a token will significantly change its representation. That's why these similarities are low, but their embeddings will stabilize in subsequent steps. Here's what happens after another denoising step where we unmask token four. The newly decoded token is changes a lot since it's no longer a mask, but in the meantime, who and there stabilize compared to the previous step. This is an empirical observation made by multiple studies. One of them is DKV cache, where D stands for delayed, as in caching with one-step delay. Now, prompt caching and delayed response caching are ultimately still approximate. However, there is a compromise that enables precise KB caching. The trick is to bring back auto regression. Given a long context window, we could split it into equally sized blocks. Within a single block, tokens are generated with diffusion. But blocks themselves are generated sequentially from left to right. This hybrid method is called block diffusion or semi-auto regression. This enables a mixed attention pattern. A particular token within the current block attends causally all the tokens in previous blocks and bidirectionally to the other tokens in its own block. Causal attention means left blocks are not affected anymore by future updates to the response. Once a block is finalized, the activations in its last denoising step can be cached and reused by future blocks. This is no longer a heuristic or approximation like the methods we saw before. It's exact KB caching. Beyond speeding up inference through caching, block diffusion comes with another major advantage, variable length generation. Once the end of sequence is produced, sampling can stop early without filling in the entire context window. Block diffusion is halfway between pure diffusion and auto regressive LLMs. We see this compromise in most state-of-the-art diffusion research including seed diffusion and the LLaMA model family.