固定状态带来长上下文：Block Diffusion 如何突破缓存瓶颈

导语

扩散语言模型能够并行生成一组 token，但双向去噪机制也带来了一个工程难题：它不能直接复用自回归模型常见的键值（KV）缓存。Block diffusion 通过按块推进生成过程，重新提供了缓存的可能性；不过，已有方案大多依赖注意力，因此缓存仍会随上下文长度线性增长，而且作为训练后的改造方案时，缓存计算未必与模型原本的计算完全一致。

arXiv 论文《Fixed State, Long Reach》提出了一条不同路径：用序列混合器把已经完成的 block 压缩成可复用状态，并配套采用 block-causal 训练目标，使推理时的缓存成为模型计算的一部分，而不是近似补丁。

核心要点

• 统一比较设置。 研究者预训练了注意力、Mamba 和混合架构三种 3B 参数 block-diffusion 去噪器，训练数据规模为 3000 亿 tokens，并使用同一个 single-frontier 目标进行比较。
• Mamba 缓存保持恒定。 状态空间模型将历史 block 汇总为固定大小的状态，因此缓存内存和每步延迟理论上不随序列长度增长；注意力缓存则仍为 O(L)。
• 超长上下文优势明显。 在 256k tokens 的测试点，注意力缓存占用 82GB、单步延迟为 29 毫秒。Mamba 相比之下实现 4.3 倍更低的延迟、11 倍更少的内存，以及 2.6 倍更高的单流吞吐。
• 批量服务更容易扩展。 固定大小的状态使 Mamba 能够继续扩大批量，在该测试中聚合吞吐达到注意力方案的 14 倍；注意力缓存则无法运行超过单流的配置。
• 外推检索能力更强。 Mamba 和混合模型在训练长度之外仍能检索信息，范围达到训练长度的 8 至 16 倍；注意力模型在 2 倍长度时检索能力出现崩溃。摘要称，实验中未观察到相应的质量损失。

意义与影响

这项工作的重要性不只是把缓存做得更小，而是把“模型架构、训练目标与推理接口”放在同一个设计中考虑。对于扩散语言模型而言，block-causal 目标让已完成 block 的状态能够被精确复用，避免了训练后强行套用缓存所产生的计算偏差。

从部署角度看，固定状态尤其适合长上下文和多并发服务：上下文继续增长时，显存压力不会按序列长度同步放大，批量吞吐也更容易维持。与此同时，结果也显示，状态空间模型的归纳偏置可能有助于超出训练长度的检索。不过，论文比较的是特定的 3B 规模、训练方案和缓存接口，结果能否直接推广到更大模型、更多任务及不同 block 配置，仍需后续研究验证。

总体而言，该研究为 block diffusion 提供了一个较完整的系统方案：通过固定状态换取可扩展缓存，并用匹配的训练目标保证缓存的计算一致性。它说明，扩散语言模型能否真正走向长上下文部署，关键不只在并行解码本身，也在于如何设计可持续复用的历史状态。

arXiv