固定状态，长距离影响：常量大小缓存在大规模块扩散中的优势

# 固定状态，长程覆盖：恒定规模缓存在大规模块扩散模型中的应用价值  
来源：https://arxiv.org/html/2609.11998  
Vaibhav Singh Pierre‑André Noël Torsten Scholak Eugene Belilovsky Oleksiy Ostapenko  
Mila 康考迪亚大学 ServiceNow Research  
††通讯作者：vaibhav\.singh@mila\.quebec  

###### 摘要  
扩散语言模型能并行解码词元，但其双向去噪机制阻碍了自回归快速推理所依赖的朴素键值（\(KV\)）缓存。**块扩散**通过逐块解码恢复了缓存能力，而目前部署的块缓存均依赖*注意力机制*：内存占用为 \(O(L)\)，且若作为免训练改造方案，仅能*近似*模型计算过程。这两个限制均可突破：序列混合器将已处理完毕的块汇总为可复用状态，从而支持块缓存；对应的块因果训练目标使缓存完全精确。我们在大规模场景下验证该方案，基于单一前沿目标在300B词元上预训练三种3B参数块扩散去噪器（注意力机制、Mamba与混合架构），并通过统一缓存接口解码三者。唯有状态空间缓存的内存与每步延迟在序列长度上保持 \(O(1)\) 恒定，而注意力缓存仍呈 \(O(L)\) 增长。在256k词元规模下（注意力机制占用达82GB且延迟29ms/步），Mamba缓存实现4.3倍延迟降低、11倍内存节省与2.6倍单流吞吐提升；且因其恒定开销，可通过批处理扩展至14倍聚合吞吐（注意力机制无法支持超出单流）。这种线性状态优势使Mamba与混合架构能在8–16倍训练长度外持续检索，而注意力机制在2倍长度时检索能力即崩溃，且质量无损。  

## 1 引言  
扩散语言模型（DLMs）已成为自回归（AR）范式的有力竞争者，而后者支撑着大多数大型语言模型\[5,1,8,34\]。掩码DLMs（如LLaDA\[24\]与Dream\[41\]）通过掩码破坏序列并学习去噪，实现多位置并行解码而非逐词元生成，从而降低延迟并提供更丰富的实例级监督\[23\]。然而实践中DLM推理常比AR更*慢*：正是使去噪成为可能的双向注意力机制，排除了AR解码器依赖的朴素键值（\(KV\)）缓存，导致每个去噪步骤需重新计算全序列注意力\[39,21\]。  

**块扩散**\[2\]通过介于两者间的方案解决矛盾：序列被划分为自回归生成的块，扩散去噪在每个*块内*运行。由于已定块形成干净前缀，其表示可被缓存复用，恢复了\(KV\)缓存并支持任意长度生成，这已成为近期快速解码系统的基础\[39,38,37\]。  

仍存在两项局限：其一，现有块缓存基于注意力骨干网络构建\[39,21,18,22\]；生成的\(KV\)缓存随上下文长度线性增长，底层注意力计算仍需二次复杂度，长上下文场景代价高昂。其二，多数方法是*免训练改造*全注意力模型，使块缓存仅能*近似*真实双向计算\[39\]。  

这两项局限并非块扩散的本质限制——仅要求已定块能汇总为可复用状态：任何具备此类状态的序列混合器均适用，且通过块因果目标训练可使缓存精确而非近似\[2\]。状态空间模型（SSMs）如Mamba\[12,9\]正是此类混合器：通过线性时间递归处理序列，其隐藏状态作为历史的固定大小摘要，与历史长度无关。  

我们基于DiffuMamba\[32\]的双向Mamba DLM去噪器构建系统，利用其递归状态作为块缓存：当块扩散定稿一个块时，前向SSM状态已用 \(O(1)\) 内存编码完整前缀。并行研究\[7\]近期在小规模（87M–350M参数）验证Mamba–注意力混合架构在反向Mamba扫描限于活跃块时支持此类缓存。但该方案在大规模预训练下的有效性、*纯*Mamba去噪器的性能，以及恒定状态对长上下文*检索*（而非仅吞吐量）的贡献尚未明确。  

因此，我们将SSM状态、注意力\(KV\)缓存及混合架构统一至同一解码接口，并在相同块因果目标下*预训练*所有3B参数模型，使缓存解码精确复现训练过程。具体而言，我们在300B词元上预训练并分析三种3B参数块扩散模型，它们共享数据、分词器、调度与解码预算，仅在去噪器上区分：Attn（全注意力）、Mamba（双向Mamba-2）与Hybrid（每五层Mamba插入注意力）。由于Mamba混合器具有更大投影层而MLP固定，Mamba和Hybrid参数量分别高出13%与11%；第5.4节表明这些额外参数在任意上下文长度下仅产生固定 \(O(d^2)\) 开销，而注意力分数计算成本为 \(O(Ld)\) 并随上下文增长。训练长度下单词元FLOPs相当，长上下文差距源于混合器而非参数预算。  

我们的贡献包括：  
- **受控的3B预训练研究**（第3、4节）：涵盖注意力、双向Mamba-2与混合骨干网络，均在300B词元上通过相同前沿块因果目标预训练。如BD3LM\[2\]（针对注意力）与近期Chaturvedi等人\[7\]（针对Mamba混合）所示，因模型在解码时使用训练目标，缓存完全精确；区别于免训练改造，缓存推理即为模型学习的功能。据我们所知，这是最大规模的研究，且是唯一包含*纯*状态空间去噪器的案例。  
- **基于线性状态骨干的恒定内存与长上下文支持**（第5.1、5.2节）。实测SSM缓存内存随长度保持 \(O(1)\)：Mamba在至256k的任意长度下占用约7.5GB并维持6.8ms/步延迟，而注意力机制需82GB；在256k规模下聚合吞吐达14倍。该骨干还能泛化至远超1024词元训练长度：Attn在NIAH与LongBench上仅支持2倍长度检索，而Mamba/Hybrid可外推至8–16倍。  
- **大规模长上下文评估**（第5节）。在除去噪器外所有条件固定下，我们覆盖至256k词元的缓存解码效率（延迟、内存与吞吐量）、NIAH与LongBench长度外推、八任务下游套件、MAUVE与生成困惑度，以及逐层FLOPs分析（将性能增益归因于*架构*而非参数量）。  

状态空间模型的这些特性并不新颖。固定状态\[12,9\]与长度外推\[10,28\]是*自回归*SSMs的既定属性，混合架构弥补了纯递归缺乏的召回能力\[35,16\]。但其能否迁移至每轮解码多个词元的扩散去噪器仍属开放问题。我们的结果构成了该AR文献在块扩散领域的对应；这种结合至关重要，因为块扩散在每次前向传播中揭示多个词元，而线性状态骨干使其能在永不增长的状态上执行。Singh等人\[32\]表明Mamba去噪器在训练长度上质量匹配注意力机制，Chaturvedi等人\[7\]验证其可在小规模下实现块缓存\[2\]；我们证明相同骨干在3B规模预训练后，能以恒定内存缓存*解码*长上下文并远超训练长度检索，将线性状态扩散LM的论证从训练扩展至推理，从吞吐量延伸至长上下文检索。  

## 2 背景与相关工作  
#### 掩码扩散语言模型。离散扩散模型定义逐步破坏词元序列的前向过程，并学习参数化反向去噪过程\[3,6\]。**掩码**（吸收态）实现对语言最有效：SEDD\[20\]、MDLM\[29\]及其改进\[31,30\]训练双向网络在噪声调度下恢复掩码词元，LLaDA\[24\]与Dream\[41\]将该方案扩展至数十亿参数。所有模型通过并行迭代取消掩码位置解码；均无法原生支持\(KV\)缓存，因每个去噪步骤需双向关注完整（部分掩码）序列。  

#### 块扩散。BD3LM\[2\]在AR与扩散间插值：词元被分组为左至右生成的块，块内执行扩散。已生成块作为干净前缀，其键值可在块间缓存，支持任意长度序列生成。我们的注意力模型遵循此方案；命名为Attn而非"BD3LM"以区分训练模型与引用方法。后续系统在此基础上构建快速解码器\[39,38,37\]。  

#### 扩散LM的缓存。由于双向去噪器阻碍标准\(KV\)缓存，一系列工作恢复了注意力DLM的缓存能力。Fast-dLLM\[39\]引入块级*近似*\(KV\)缓存加置信感知并行解码，通过缓存固定上下文并在块边界刷新，在LLaDA/Dream上实现大幅加速；dKV-Cache\[21\]提出兼容双向注意力的延迟\(KV\)缓存；dLLM-Cache\[18\]与弹性缓存变体\[22\]添加自适应淘汰与复用。近期研究专门针对*长上下文*场景：Prefilling-dLLM\[40\]缓存分块前缀\(KV\)并保留最相关的块；Focus-dLLM\[19\]利用注意力稀疏性实现动态缓存淘汰，在8–32k规模报告9–28倍加速。另一类方法通过缩放位置编码扩展注意力DLM的*可用上下文*：LongLLaDA\[17\]应用免训练NTK-RoPE（并指出"局部感知"滑动窗口偏差），UltraLLaDA\[14\]则*训练*至128k窗口。所有这些缓存均为注意力专用且内存占 \(O(L)\)，且因多为免训练改造，仅*近似*其替代的全注意力计算。例外是训练时块因果方案：BD3LM\[2\]（针对注意力）及与本工作并行的Chaturvedi等人\[7\]（针对Mamba-注意力混合），其缓存精确因模型在相同干净前缀计算上训练。我们的模型属于此范畴。  

#### 状态空间模型与混合架构。SSMs通过线性时间选择性递归建模序列：S4\[13\]、Mamba\[12\]与Mamba-2\[9\]在语言任务上匹配或超越注意力机制，且携带固定状态。相关结构化算子\[27,11\]共享线性时间特性。混合AR模型交错注意力与线性递归，结合全局召回与廉价长程混合\[16,10,36\]。DiffuMamba\[32\]