Fixed-State Caching Gives Block Diffusion a Longer Reach
Introduction
Diffusion language models can decode multiple tokens in parallel, but their bidirectional denoising process prevents them from using the standard key-value cache that accelerates autoregressive inference. Block diffusion addresses this limitation by advancing generation one block at a time. Yet most existing block-cache designs remain tied to attention, so their memory footprint grows linearly with context length. When added after training, they may also approximate rather than reproduce the model’s original computation.
The arXiv paper Fixed State, Long Reach proposes a more integrated solution. A sequence mixer summarizes finalized blocks into a reusable state, while a matching block-causal objective teaches the model to use that state directly. The result is intended to make caching exact rather than a training-free retrofit.
Key findings
- A controlled comparison. The researchers pretrained three 3B block-diffusion denoisers—attention, Mamba, and a hybrid—on 300 billion tokens using the same single-frontier objective and decoded them through one cached interface.
- Constant-size state-space cache. Mamba compresses completed blocks into a fixed-size state. Its cache memory and per-step latency therefore remain constant with sequence length, whereas an attention cache remains O(L).
- Large gains at 256k tokens. At this context length, the attention cache used 82GB and required 29 milliseconds per step. Mamba reduced latency by 4.3x, memory use by 11x, and increased single-stream throughput by 2.6x.
- Better batching behavior. Because the Mamba footprint stays fixed, it can continue scaling with batch size. The reported aggregate throughput reached 14x that of attention, which could not run beyond a single stream in the test.
- Longer retrieval reach. Mamba and the hybrid model continued retrieving information at 8–16 times their training length, while attention retrieval collapsed at 2x. The abstract reports no measured quality cost for this behavior.
Why it matters
The main contribution is not simply a smaller cache. It connects architecture, training, and inference into one design. The block-causal objective makes the summary state part of the model’s intended computation, reducing the mismatch that can arise when caching is bolted onto a pretrained bidirectional denoiser.
For deployment, a constant-size state is particularly attractive for long-context and multi-stream serving. As prompts grow, memory pressure does not increase in proportion to the full history, and larger batches become more practical. The retrieval results also suggest that state-space inductive bias may help models use information beyond their training window.
The findings should still be interpreted within the reported setup: three 3B models, one training recipe, and a particular cached interface. Whether the same advantages persist at larger scales, across tasks, or under different block configurations remains an open question. Even so, the study presents a coherent route for making block diffusion more viable as a long-context inference system.
Comments
Checking sign-in status...
Loading comments...