1. 论文基本信息

1.1. 标题

dKV-Cache: The Cache for Diffusion Language Models

1.2. 作者

Xinyin Ma, Runpeng Yu, Gongfan Fang, Xinchao Wang* (通讯作者)

1.3. 隶属机构

National University of Singapore

1.4. 发表信息

发布时间 (UTC): 2025-05-21T17:32:10.000Z 原文链接: https://arxiv.org/abs/2505.15781 PDF 链接: https://arxiv.org/pdf/2505.15781v1.pdf 状态: 预印本 (arXiv)

1.5. 摘要

扩散语言模型 (Diffusion Language Models, DLMs) 被视为自回归语言模型 (autoregressive language models, ARs) 的有力竞争者。然而，DLMs长期以来一直受限于推理速度慢的问题。一个核心挑战是，它们的非自回归架构和双向注意力 (bidirectional attention) 机制排除了用于加速解码的键值缓存 (key-value cache, KV-Cache) 的直接应用。本文通过提出一种类似KV-Cache的机制——延迟KV-Cache (delayed KV-Cache)，来解决DLMs去噪过程中的这一瓶颈。该方法基于观察，即在整个扩散过程中，不同词元 (token) 具有独特的表示动态 (representation dynamics)。因此，本文提出了一种针对键 (key) 和值 (value) 状态的延迟且有条件的缓存策略。设计了两种互补的变体来逐步缓存键和值：(1) dKV-Cache-Decode，它提供了几乎无损的加速，甚至在长序列上改进了性能，这表明现有DLMs在推理过程中可能未充分利用上下文信息。(2) dKV-Cache-Greedy，它采用激进的缓存策略，生命周期缩短，实现了更高的加速比，具有二次时间复杂度，但会带来一定的性能下降。最终，dKV-Cache 在推理中实现了 $2 \text{ - } 10 \times$ 的加速，大大缩小了ARs和DLMs之间的差距。本文在多个基准测试中评估了 dKV-Cache，在通用语言理解、数学和代码生成基准测试中均实现了加速。实验表明，缓存机制也可以在DLMs中使用，甚至可以以无需训练的方式应用于现有DLMs。

2. 整体概括

2.1. 研究背景与动机

2.1.1. 扩散语言模型 (DLMs) 的潜力与挑战

扩散语言模型 (DLMs) 近年来作为文本生成领域的一个新兴范式而出现，其灵感来源于扩散模型 (diffusion models) 在图像、视频等连续域的成功。与自回归语言模型 (ARs) 逐词元 (token-by-token) 生成文本不同，DLMs通过逐步细化一个初始噪声序列或掩码 (masked) 序列来生成连贯的输出。DLMs的一个显著优势是它们能够并行解码任意数量的词元，这理论上比ARs的每次生成一个词元要高效。

2.1.2. 现有DLMs的推理效率瓶颈

尽管DLMs在理论上具有并行解码的潜力，但在实际应用中，它们仍然比ARs慢得多。这种低效率主要归因于两个因素：

1. KV-Cache机制的不兼容性 (Incompatibility with KV-Cache): DLMs无法直接利用广泛用于ARs的键值缓存 (KV-Cache) 机制。KV-Cache通过存储前面已计算的键 (Key) 和值 (Value) 状态来避免重复计算，从而显著加速解码过程。
2. 大量的网络评估 (Large number of network evaluations): DLMs的去噪过程通常需要大量网络评估。生成一个长度为 $L$ 的序列通常需要 $L$ 个去噪步骤，每个步骤都涉及一次完整的双向注意力 (bidirectional attention) 计算，导致时间复杂度为 $\mathcal{O}(L^3)$。相比之下，AR模型通过KV-Cache将每步复杂度降低到 $\mathcal{O}(L^2)$。

2.1.3. 现有研究空白与本文切入点

AR模型中KV-Cache的有效性建立在两个关键假设之上：

1. 因果注意力掩码 (Causal attention mask): 键和值状态在后续解码步骤中保持固定。AR模型通过因果注意力掩码实现这一点，确保每个词元只关注之前的词元。但DLMs采用双向注意力，所有词元都可相互关注，使得K/V复用变得困难。

2. 固定且从左到右的解码顺序 (Fixed left-to-right decoding order): AR模型按顺序计算下一个词元，可以有选择地计算当前解码位置的QKV状态。DLMs则支持灵活的生成顺序，任何位置都可能被更新，这使得预先确定哪些词元需要计算隐藏状态变得不可能。

   本文的切入点是解决DLMs中KV-Cache机制的缺失。通过深入分析DLMs中词元表示的动态变化，发现即使在双向注意力下，键和值状态也并非完全不可复用，而是需要延迟 (delayed) 和有条件 (conditioned) 的复用。

2.2. 核心贡献/主要发现

本文的主要贡献如下：

1. 首次为扩散语言模型提出了KV-Cache机制 (First KV-Cache for DLMs): 利用词元表示的演变动态，引入了一种与双向注意力兼容的延迟缓存策略。
2. 提出两种实用的dKV-Cache变体 (Two practical variants):
   • dKV-Cache-Decode: 实现长期的缓存复用，提供了几乎无损的性能加速，甚至在长序列上有所提升。
   • dKV-Cache-Greedy: 采用激进缓存策略，缩短缓存生命周期，以实现更高的加速比（将复杂度从 $\mathcal{O}(L^3)$ 降低到 $\mathcal{O}(L^2)$），但可能带来一些性能下降。
3. 广泛的实验验证 (Extensive experimental validation): 在现有7B规模的DLMs（包括LLaDA和Dream）上进行了广泛实验，在通用语言理解、代码生成和数学问题解决等多个基准测试中，实现了 $2 \text{ - } 10 \times$ 的推理加速，同时性能损失极小或可忽略不计。实验还证明该方法无需训练即可应用于现有DLMs。

3. 预备知识与相关工作

3.1. 基础概念

3.1.1. 扩散模型 (Diffusion Models)

概念定义： 扩散模型 (Diffusion Models) 是一类生成模型，其核心思想是模仿一个逐渐将数据（如图像、文本）转化为噪声的正向过程，并学习其逆过程，即从噪声中逐步恢复出清晰数据的去噪过程。这个逆过程通过一个参数化的马尔可夫链 (Markov chain) 来实现，其中模型在每个时间步预测并移除噪声。

在文本领域的应用： 对于文本生成，扩散模型通常分为两类：

• 连续扩散模型 (Continuous Diffusion Models): 在连续的词嵌入 (word embedding) 空间上进行扩散和去噪。
• 离散扩散模型 (Discrete Diffusion Models): 直接在离散的词元 (token) 空间上操作，通常通过掩码 (masking) 来引入噪声。本文主要关注离散扩散模型，特别是掩码扩散语言模型 (masked diffusion language models)。

3.1.2. 扩散语言模型 (Diffusion Language Models, DLMs)

概念定义： 扩散语言模型 (DLMs) 是将扩散模型的原理应用于文本生成任务的模型。它们不像传统的自回归语言模型 (ARs) 那样从左到右逐个词元生成，而是通过迭代地去噪或填充序列中的掩码词元来生成整个文本序列。

去噪过程： 给定一个噪声序列 $\mathbf{x}_{c(t)}$ (其中 c(t) 表示连续时间步 t/T， $T$ 为总去噪步数)，DLM学习一个神经网络模型 $\theta$，来预测原始的、未损坏的序列 $\mathbf{x}_0$，或者预测从 $\mathbf{x}_{c(t)}$ 到 $\mathbf{x}_{c(t-1)}$ 的逆向转换概率 $p_\theta(\mathbf{x}_{c(t-1)} | \mathbf{x}_{c(t)})$。在掩码扩散模型中，噪声通常表现为将一些词元替换为 [MASK] 词元。去噪过程通常涉及：

1. 模型预测一个去噪后的序列 $\mathbf{x}_0$。
2. 根据预测的 $\mathbf{x}_0$ 和当前噪声序列 $\mathbf{x}_{c(t)}$，重新掩码一部分词元，生成下一个时间步的序列 $\mathbf{x}_{c(t-1)}$。重新掩码的策略有很多种，例如随机重新掩码、保留置信度高的词元等。

3.1.3. 自回归语言模型 (Autoregressive Language Models, ARs)

概念定义： 自回归语言模型 (ARs) 是一种传统的文本生成模型，如GPT系列。它们通过在给定所有先前生成词元的情况下，预测下一个词元的方式来生成文本。这种生成方式是从左到右 (left-to-right)、顺序 (sequential) 的。

特点：

• 因果注意力 (Causal Attention): 在计算注意力时，每个词元只能关注它之前的词元，而不能关注它之后的词元。这通过一个因果注意力掩码实现。
• KV-Cache的优势： 由于生成是顺序的，并且每个词元只依赖于其左侧的上下文，因此在生成新词元时，前面词元的键 (Key) 和值 (Value) 状态可以被缓存并直接复用，无需重新计算。

3.1.4. Transformer 架构与注意力机制

概念定义： Transformer [44] 是一种基于注意力机制的神经网络架构，广泛应用于自然语言处理任务。其核心是自注意力 (Self-Attention) 机制。

自注意力机制： 对于输入序列中的每个词元，Transformer 会计算三个向量：查询 (Query, $Q$)、键 (Key, $K$) 和值 (Value, $V$)。这些向量通过将词元对应的隐藏状态分别乘以三个不同的权重矩阵 $\mathbf{W_Q}, \mathbf{W_K}, \mathbf{W_V}$ 得到。 注意力机制的输出通过以下公式计算： $$\mathrm{Attention}(Q, K, V) = \mathrm{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$ 符号解释：

• $Q$: 查询矩阵，维度为 $L \times d_k$，其中 $L$ 是序列长度，$d_k$ 是键向量的维度。
• $K$: 键矩阵，维度为 $L \times d_k$。
• $V$: 值矩阵，维度为 $L \times d_v$，其中 $d_v$ 是值向量的维度。
• $K^T$: 键矩阵的转置。
• $\sqrt{d_k}$: 用于缩放点积的因子，防止点积过大导致 softmax 函数的梯度过小。
• $\mathrm{softmax}(\cdot)$: 归一化函数，将注意力分数转换为概率分布。

3.1.5. KV-Cache (Key-Value Cache)

概念定义： KV-Cache 是一种在Transformer解码器中用于加速自回归生成的技术。在自回归模型中，每生成一个新词元，模型都需要重新计算整个序列的键 (Key) 和值 (Value) 向量。KV-Cache通过存储前面已计算的词元的K和V向量，在生成后续词元时直接复用这些缓存值，从而避免重复计算，显著提高推理速度。

工作原理： 在每个解码步骤 $t$：

1. 计算当前词元 $t$ 的查询向量 $Q_t^{[t]}$。
2. 计算当前词元 $t$ 的键向量 $K_t^{[t]}$ 和值向量 $V_t^{[t]}$。
3. 将 $K_t^{[t]}$ 和 $V_t^{[t]}$ 与之前所有词元缓存的 $K_{t-1}^{[1:t-1]}$ 和 $V_{t-1}^{[1:t-1]}$ 进行拼接 (concatenate)，形成完整的 $K_t^{[1:t]}$ 和 $V_t^{[1:t]}$。
4. 使用 $Q_t^{[t]}$ 和拼接后的 $K_t^{[1:t]}, V_t^{[1:t]}$ 计算注意力输出。

3.2. 前人工作

3.2.1. 扩散语言模型 (DLMs) 的发展

早期的扩散模型主要集中在连续数据域，如图像生成 [15, 48]。随后，研究者将扩散模型扩展到离散数据，例如 [25, 3] 将DDPM (Denoising Diffusion Probabilistic Models) 扩展到分类数据，并定义了用于损坏和去噪的转移矩阵。

• 连续空间扩散： [29] 在连续词嵌入空间引入了连续时间扩散，[21] 通过在词汇表上定义一个单纯形 (simplex) 使生成质量与GPT-2相媲美。
• 离散空间扩散： [22] 训练BERT学习离散扩散过程的逆向过程，[40] 提出了简单掩码离散扩散模型 (Masked Discrete Diffusion, MDLM)，并展示了其与现有DLMs竞争的性能。
• 可扩展性与性能： [37, 19] 将掩码扩散语言模型扩展到数十亿参数规模，实现了与顶尖自回归大语言模型 (LLMs) 相当的性能。
• 非自回归与半自回归： [20] 是非自回归模型的早期工作。Block Diffusion [2] 将非自回归离散语言扩散模型扩展为半自回归模型，使其能够生成任意长度的序列。

3.2.2. 生成模型中的缓存机制 (Cache in Generative Models)

缓存是一种存储频繁访问数据的小而快的内存，以减少CPU从较慢内存中获取数据的时间。

• Transformer中的KV-Cache： KV-Cache [44] 是Transformer中的基础技术，用于缓存先前词元的键 (Key) 和值 (Value) 张量。为解决长上下文生成中的内存消耗问题，提出了多种改进技术 [18, 32, 26]。
• 图像扩散模型中的缓存： 图像扩散模型中也探索了缓存机制 [36, 46]，它们利用高级特征 [35, 9] 或注意力图 [54, 31] 之间的时间相似性来加速推理。
• DLMs中的KV-Cache探索：
  • [2] 在半自回归扩散语言模型中探索了KV-Cache，但这需要在训练中考虑KV-Cache，导致训练前向计算量翻倍，且形式仍受自回归限制。
  • [40] 也考虑了缓存，但有一个严格条件：没有新词元被计算。 本文的工作则旨在为非自回归的DLMs提供一种无需训练的KV-Cache机制。

3.3. 技术演进与差异化分析

DLMs在并行生成能力上理论优于ARs，但实际上由于其双向注意力机制和多步去噪过程，推理速度远不如ARs。ARs通过KV-Cache和因果注意力实现了高效的 $\mathcal{O}(L^2)$ 复杂度，而DLMs则受限于 $\mathcal{O}(L^3)$。

本文的差异化分析在于：

• 现有DLMs无法直接使用KV-Cache的根本原因： 并非仅仅是双向注意力，更重要的是K/V状态的“时间步变异性”和“非顺序解码顺序”这两个核心挑战。前人工作 [2, 40] 或是通过修改训练过程，或是施加严格限制，未能完全解决DLMs中KV-Cache的通用性和训练无关性问题。
• 本文的创新点： 观察到DLMs中词元表示的动态性，特别是词元一旦解码其表示趋于稳定，启发了“延迟且有条件”的缓存策略。这意味着即使在双向注意力下，K和V也并非完全不可复用，而是需要在适当的时机（延迟）和条件下（已解码词元）进行复用。本文提出的dKV-Cache是第一个专门为非自回归DLMs设计的训练无关的KV-Cache机制，它不依赖于自回归或半自回归结构，而是直接适应了DLMs的去噪过程特性。

4. 方法论

本文的核心是为扩散语言模型（DLMs）提出一种名为 dKV-Cache 的键值缓存（KV-Cache）机制，以解决其推理速度慢的瓶颈。该方法基于对DLMs中词元表示动态性的深入观察。

4.1. 方法原理

dKV-Cache 的核心思想是，尽管DLMs采用双向注意力机制，理论上与传统KV-Cache不兼容，但词元表示的键 (Key) 和值 (Value) 状态并非完全不可复用。通过分析发现：

1. K和V的跨时间步相似性： 尽管在不同时间步之间存在差异，K和V的嵌入在各个时间步之间仍表现出较高的一致性 (Figure 2(a))。

2. 解码后表示的稳定性： 一旦词元被解码，其表示在后续步骤中变得相对稳定，而那些仍被掩码的词元表示则会持续显著波动 (Figure 2(b))。

3. K和V最显著的变化： K和V最显著的变化发生在词元从 [MASK] 状态转换为其解码形式的步骤，以及去噪过程的早期阶段 (Figure 2(c))。

   这些观察表明，可以设计一种延迟且有条件的缓存策略：只缓存已解码且表示相对稳定的词元的键和值状态，并在需要时进行更新和刷新。

4.2. 核心方法详解

4.2.1. 扩散语言模型的采样过程

本文主要关注连续时间离散语言模型，特别是掩码扩散语言模型 [40, 37]。 给定一个初始噪声序列 $\mathbf{x}_1^{1:L}$ (通常所有词元都是 [MASK])。在每个去噪时间步 $t$，去噪模型 $p_\theta(\mathbf{x}_{c(t-1)} | \mathbf{x}_{c(t)})$ 会被调用。模型首先预测原始序列 $\mathbf{x}_0$，然后根据 $q(\mathbf{x}_{c(t-1)} | \mathbf{x}_0, \mathbf{x}_{c(t)})$ 重新掩码剩余的词元。未被掩码的词元会保持不变，或者通过选择置信度最高的词元 [7] 或选择两个最可能值之间边距最大的前 $k$ 个位置 [27] 来决定。

4.2.2. 自回归模型中KV-Cache的公式化

在Transformer解码器中，每个层都会将当前隐藏状态 $\mathbf{h}_t$ 投影成查询-键-值三元组 $(\mathbf{Q}_t, \mathbf{K}_t, \mathbf{V}_t)$。在步骤 $t$，只有第 $t$ 个词元的隐藏状态 $\mathbf{h}_t^{[t]}$ 被计算。递归的KV-Cache更新是将新的键值对追加到正在运行的缓冲KV-Cache中：

$$\mathbf{z}_t = \mathrm{softmax}\left(\frac{\mathbf{Q}_t^{[t]} \left(\mathbf{K}_t^{[1:t]}\right)^T}{\sqrt{d_k}}\right) \mathbf{V}^{[1:t]} \quad \mathrm{with} \left\{ \begin{array}{ll} \mathbf{K}_t^{[1:t]} = \mathrm{concat}\left(\mathbf{K}_{t-1}^{[1:t-1]}, \mathbf{K}_t^{[t]}\right) \\ \mathbf{V}_t^{[1:t]} = \mathrm{concat}\left(\mathbf{V}_{t-1}^{[1:t-1]}, \mathbf{V}_t^{[t]}\right) \end{array} \right. \tag{1}$$

符号解释：

• $\mathbf{z}_t$: 在时间步 $t$ 时注意力头 (attention head) 的输出。
• $\mathbf{Q}_t^{[t]}$: 在时间步 $t$ 时，当前词元的查询向量。
• $\mathbf{K}_t^{[1:t]}$: 在时间步 $t$ 时，所有已解码词元的键矩阵，包括当前词元。
• $\mathbf{V}^{[1:t]}$: 在时间步 $t$ 时，所有已解码词元的值矩阵，包括当前词元。
• $\sqrt{d_k}$: 用于缩放点积的因子，其中 $d_k$ 是键向量的维度。
• $\mathrm{concat}(\cdot)$: 拼接操作，将旧的缓存与当前词元的新键值对连接起来。
• $\mathbf{K}_{t-1}^{[1:t-1]}$ 和 $\mathbf{V}_{t-1}^{[1:t-1]}$: 在时间步 t-1 时缓存的键矩阵和值矩阵。

4.2.3. KV-Cache为何不能直接用于DLMs？

标准KV-Cache与扩散语言模型不兼容，原因有二：

1. 时间步变化的键值状态 (Timestep-variant key and value states)： 在自回归设置中，借助因果注意力掩码，每个时间步 $m$ 的 $\mathbf{K}_m^{[1:t-1]}$ 和 $\mathbf{V}_m^{[1:t-1]}$ 从 t-1 步开始都相同。然而，DLMs采用双向注意力掩码，使得 $\mathbf{K}_m^{[1:t-1]}$ 和 $\mathbf{V}_m^{[1:t-1]}$ 在不同时间步 $m \neq n$ 时也会不同。这打破了传统KV-Cache所依赖的全局复用假设。
2. 非顺序解码顺序 (Non-sequential decoding order)： DLMs的生成不遵循严格的从左到右顺序。模型在每个去噪步骤中根据计算出的概率动态填充掩码位置。解码词元的位置在模型前向传播后才确定，后续更新可能针对序列中的任何位置，而非顺序推进。这种不确定性阻止了预先确定哪个词元 $i$ 需要计算其隐藏状态 $\mathbf{h}^{[i]}$ 及其 $\mathbf{Q}^{[i]}, \mathbf{K}^{[i]}, \mathbf{V}^{[i]}$。

4.2.4. 广义非顺序KV-Cache (Generalized Non-sequential KV-Cache)

为解决上述问题，本文首先提出了一个更通用的非顺序KV-Cache公式。它用任意子集 $\mathcal{S}_t \subseteq \mathcal{T} = \{1, \dots, L\}$ 取代了 Eq.1 中连续的切片 $\mathrm{K}^{[1:t-1]}$，并将下一个词元位置从固定的 $t$ 变为 $D_t$。从先前步骤收集到的缓存键和值现在是 $\mathbf{K}_t^{\mathcal{S}_t}$。

$$\mathbf{z}_t = \mathrm{softmax}\left(\frac{\mathbf{Q}_t^{D_t} \left(\mathbf{K}_t^{\mathcal{S}_t \cup \{D_t\}}\right)^T}{\sqrt{d_k}}\right) \mathbf{V}_t^{\mathcal{S}_t \cup \{D_t\}} \mathrm{with} \left\{ \begin{array}{ll} \mathbf{K}_t^{\mathcal{S}_t \cup \{t\}} = \mathrm{concat.reorder}\left(\mathbf{K}_{t-1}^{\mathcal{S}_t}, \mathbf{K}_t^{D_t}\right) \\ \mathbf{V}_t^{\mathcal{S}_t \cup \{t\}} = \mathrm{concat.reorder}\left(\mathbf{V}_{t-1}^{\mathcal{S}_t}, \mathbf{V}_t^{D_t}\right) \end{array} \right. \tag{2}$$

符号解释：

• $D_t$: 在时间步 $t$ 被去噪（解码）的词元。

• $\mathcal{S}_t$: 在时间步 $t$ 时，已缓存的词元位置集合。

• $\mathbf{Q}_t^{D_t}$: 在时间步 $t$ 时，被去噪词元 $D_t$ 的查询向量。

• $\mathbf{K}_t^{\mathcal{S}_t \cup \{D_t\}}$ 和 $\mathbf{V}_t^{\mathcal{S}_t \cup \{D_t\}}$: 包含已缓存词元和当前去噪词元的键和值矩阵。

• $\mathrm{concat.reorder}(\cdot)$: 一种特殊的拼接和重排操作，用于高效处理非连续的键和值（详见 A.1 concat_reorder 设计）。

  本文将 Eq.2 进一步扩展，将查询输入 $\mathbf{Q}_t^{D_t}$ 从单词元解码推广到多词元和任意词元解码。具体地，在解码步骤 $t$，构建一个动态集合 $\mathcal{M}_t$，表示在生成过程中尚未确定（即仍为掩码）的词元子集。然后使用 $\mathbf{h}_t^{\mathcal{M}_t}$ 计算 $\mathbf{Q}_t^{\mathcal{M}_t}, \mathbf{K}_t^{\mathcal{M}_t}, \mathbf{V}_t^{\mathcal{M}_t}$。缓存延迟到词元被解码的步骤，而不是其在输入中追加的时间：

$$\mathrm{concat\_reorder}(\underbrace{\mathbf{K}_{t-1}^{\mathcal{S}_t}}_{\text{Cache from } t-1}, \underbrace{\mathbf{K}_t^{D_t}}_{\text{Calculated at } t}) \Rightarrow \mathrm{concat\_reorder}(\underbrace{\mathbf{K}_{t-1}^{\mathcal{T} \setminus \mathcal{M}_t}}_{\text{Cache from } t-1}, \underbrace{\mathbf{K}_t^{\mathcal{M}_t}}_{\text{Calculated at } t}) \tag{3}$$

符号解释：

• $\mathcal{T}$: 参与去噪过程的所有词元的集合。

• $\mathcal{M}_t$: 在时间步 $t$ 仍未被解码（即仍被掩码）的词元子集。

• $\mathcal{T} \setminus \mathcal{M}_t$: 在时间步 $t$ 已被解码的词元子集。

• $\mathbf{K}_{t-1}^{\mathcal{T} \setminus \mathcal{M}_t}$: 在时间步 t-1 已被解码的词元的缓存键。

• $\mathbf{K}_t^{\mathcal{M}_t}$: 在时间步 $t$ 重新计算的未解码词元的键。

  这个设计反映了核心观察：只有已解码的词元才符合缓存条件，而剩余的掩码词元在每个步骤都需要重新编码。此外，由于 $\mathcal{M}_t$ 在每个步骤都明确可知，该方法解决了需要预先定义或预测去噪顺序的问题。对应于 $\mathcal{T} \setminus \mathcal{M}_t$ 的缓存键和值在步骤之间复用，而未确定的位置在每个步骤重新计算以确保在双向注意力掩码下的正确性。

4.2.5. 一步延迟缓存 (One-step Delayed Caching)

根据图 2(c) 的分析，K和V最显著的变化发生在词元从 [MASK] 过渡到其解码形式的步骤。如果立即缓存，当前 $K_t^{D_t}$ 可能与 $K_{t+1}^{D_t}$ 显著不同。为了解决这个问题，本文引入了一步延迟缓存。在时间步 $t$，使用前一步骤的掩码状态 $\mathcal{M}_{t-1}$ 来确定哪些词元是可缓存的。这种方法，称为 dKV-Cache-Decode，最终形式化为：

$$\mathbf{z}_t = \mathrm{softmax}\left(\frac{\mathbf{Q}_t^{\mathcal{M}_t} \left(\mathbf{K}_t^{\mathcal{T}}\right)^T}{\sqrt{d_k}}\right) \mathbf{V}_t^{\mathcal{T}} \mathrm{with} \left\{ \begin{array}{ll} \mathbf{K}_t^{\mathcal{T}} = \mathrm{concat.reorder}\left(\mathbf{K}_{t-1}^{\mathcal{T} \setminus \mathcal{M}_{t-1}}, \mathbf{K}_t^{\mathcal{M}_{t-1}}\right) \\ \mathbf{V}_t^{\mathcal{T}} = \mathrm{concat.reorder}\left(\mathbf{V}_{t-1}^{\mathcal{T} \setminus \mathcal{M}_{t-1}}, \mathbf{V}_t^{\mathcal{M}_{t-1}}\right) \end{array} \right. \tag{4}$$

符号解释：

• $\mathbf{Q}_t^{\mathcal{M}_t}$: 在时间步 $t$ 时，当前仍被掩码的词元的查询向量。

• $\mathbf{K}_t^{\mathcal{T}}$ 和 $\mathbf{V}_t^{\mathcal{T}}$: 包含所有词元的键和值矩阵，其中一部分来自缓存，一部分重新计算。

• $\mathbf{K}_{t-1}^{\mathcal{T} \setminus \mathcal{M}_{t-1}}$ 和 $\mathbf{V}_{t-1}^{\mathcal{T} \setminus \mathcal{M}_{t-1}}$: 在时间步 t-1 时，已解码词元（即上一时间步的非掩码词元）的缓存键和值。

• $\mathbf{K}_t^{\mathcal{M}_{t-1}}$ 和 $\mathbf{V}_t^{\mathcal{M}_{t-1}}$: 在时间步 $t$ 重新计算的，上一时间步仍为掩码的词元的键和值。

  虽然这会稍微降低效率，但实验发现它对于 dKV-Cache 机制在DLMs中保持准确性和稳定性至关重要。

4.2.6. 缓存刷新机制 (Cache Refreshing Mechanism)

为了在解码过程中保持一致性和提高正确性，引入了缓存刷新机制。每隔 $N$ 个步骤，存储的缓存会被丢弃并重新刷新。此时计算会恢复到正常计算，即在 Eq.4 中用空集 $\varnothing$ 替换 $\mathcal{M}_{t-1}$。

4.2.7. dKV-Cache-Prefill 和 dKV-Cache-PD

• dKV-Cache-Prefill： 预填充 (prefill) 词元（即始终已解码的输入上下文词元）主要关注彼此，受后续词元影响有限。因此，dKV-Cache-Prefill 对预填充词元进行缓存而不进行刷新。
• dKV-Cache-PD： 基于 dKV-Cache-Prefill，dKV-Cache-PD 仅对新解码的词元进行间歇性刷新，而预填充词元的键和值则无需重新计算。

4.2.8. dKV-Cache-Greedy: 激进的缓存策略

上述 dKV-Cache-Decode 方法的复杂度仍为 $\mathcal{O}(L^3)$，因为它最初的 $\mathcal{M}_t$ 包含整个序列，直到最后才缩小到单个词元。为了将复杂度降低到 ARs 类似的 $\mathcal{O}(L^2)$，dKV-Cache-Greedy 采用了更激进的缓存机制。 它将 $\mathcal{M}_t$ 定义为仅包含三个组件：

1. 当前步骤的解码词元 $D_t$。
2. 前一步骤的解码词元 $D_{t-1}$ (受一步延迟缓存启发)。
3. 一个局部窗口 $\mathcal{W}(t)$。局部窗口定义为以当前解码词元或前一解码词元为中心，包含其自身和邻近词元的固定大小窗口。例如： $$\mathcal{W}_t = \left\{ x_i \ : \middle| \ : i \in \left[ D_t - \left\lceil \frac{w}{2} \right\rceil , D_t + \left\lfloor \frac{w}{2} \right\rfloor \right] \right\}$$ 符号解释：
   • $x_i$: 序列中的第 $i$ 个词元。
   • $D_t$: 在时间步 $t$ 解码的词元的位置。
   • $w$: 窗口大小。
   • $\lceil \cdot \rceil$: 向上取整。
   • $\lfloor \cdot \rfloor$: 向下取整。 通过限制缓存的词元数量（窗口大小 $| \mathcal{W}_t |$ 固定），dKV-Cache-Greedy 引入了额外的计算，但将整体时间复杂度保持在 $\mathcal{O}(L^2)$。

4.2.9. concat_reorder 机制设计 (附录 A)

concat_reorder 是 dKV-Cache 的实现细节，旨在提高其在DLMs中的速度。与ARs中标准KV-Cache的简单拼接不同，dKV-Cache 需要从任意位置收集和分散键值，引入了比连续空间拼接效率更低的索引操作。

挑战：

• 缓存步骤： 计算键和值后，需要收集非连续位置上缓存词元对应的状态。
• 复用步骤： 为了获得完整的键和值矩阵，需要将这些向量分散回序列的原始位置。
• 与ARs的对比： ARs中的KV-Cache仅需矩阵切片和拼接，效率更高。

解决方案： concat_reorder 的关键思想是在Transformer的前向计算过程中重新排列词元位置，将所有缓存的词元连续地放置在一侧（例如左侧），而新解码的词元放置在另一侧。这允许将部分索引操作转移到词元级别（形状为 [B, L] 的矩阵），而不是中间状态（形状为 [B, L, D] 的矩阵）。

核心步骤 (以 dKV-Cache-Decode 为例，Algorithm 1):

1. 输入准备： 接收当前时间步 $t$ 的序列 $\mathbf{x}$，掩码词元的索引 $\mathcal{M}_t$，以及上一步缓存的键 $\mathbf{K}_{t-1}^{\mathcal{T} \setminus \mathcal{M}_{t-1}}$ 和值 $\mathbf{V}_{t-1}^{\mathcal{T} \setminus \mathcal{M}_{t-1}}$。

2. 一步延迟： 使用 $\mathcal{M}_{t-1}$（前一步的掩码状态）来确定哪些词元需要重新计算。从 $\mathbf{x}$ 中提取这些词元形成 $\mathbf{x}' = \mathbf{x}[\mathcal{M}_{t-1}]$。

3. 位置嵌入重排： 重新排列位置嵌入 PE，将缓存词元（来自 $\mathcal{T} \setminus \mathcal{M}_{t-1}$）的 PE 放在左侧，未缓存词元（来自 $\mathcal{M}_{t-1}$）的 PE 放在右侧。

4. 计算 Q, K, V： 使用重排后的输入序列 $\mathbf{x}'$ 和位置嵌入计算 Transformer 的 Q, K, V 向量，得到 $\mathbf{Q}_t^{\mathcal{M}_t}, \mathbf{K}_t^{\mathcal{M}_t}, \mathbf{V}_t^{\mathcal{M}_t}$。

5. 拼接 K 和 V： 将上一步缓存的 K 和 V 与当前计算的 K 和 V 进行拼接：

   • $\mathbf{K}_t^{\mathcal{T}} \gets \mathrm{Concat}(\mathbf{K}_{t-1}^{\mathcal{T} \setminus \mathcal{M}_{t-1}}, \mathbf{K}_t^{\mathcal{M}_{t-1}})$
   • $\mathbf{V}_t^{\mathcal{T}} \gets \mathrm{Concat}(\mathbf{V}_{t-1}^{\mathcal{T} \setminus \mathcal{M}_{t-1}}, \mathbf{V}_t^{\mathcal{M}_{t-1}})$ 这形成了包含所有词元（缓存和新计算）的完整 K 和 V 矩阵。

6. K 和 V 矩阵重排： 根据新的索引 $I'$ （即 $\mathcal{T} \setminus \mathcal{M}_t$ 在重排后的序列 $\lfloor \mathbf{x}[\mathcal{T} \setminus \mathcal{M}_{t-1}] \ ; \ \mathbf{x}[\mathcal{M}_{t-1}] \rfloor$ 中的位置）对完整的 $\mathbf{K}_t^{\mathcal{T}}$ 和 $\mathbf{V}_t^{\mathcal{T}}$ 进行重排，以获取缓存。

7. 注意力计算： 使用 $\mathbf{Q}_t^{\mathcal{M}_t}$ 和重排后的 $\mathbf{K}_t^{\mathcal{T}}, \mathbf{V}_t^{\mathcal{T}}$ 计算注意力输出 $p'$。

8. 结果散布： 将 $p'$ 散布回原始位置，得到最终的概率分布 $p$。

   位置编码 (Positional Encoding) 的处理： 词元位置的变化会影响位置编码 (Positional Encoding, PE)。解决方案是也重新排列位置嵌入。由于位置嵌入的重排只需每模型评估一次，并可在层之间共享，因此开销不大。

一步延迟的额外好处： 由于一步延迟缓存，步骤 $t$ 缓存词元的位置对应于步骤 t-1 解码的词元位置。这种对齐允许跟踪需要缓存哪些键值条目，而无需存储整个键值矩阵。

以下是 concat_reorder 的伪代码：

Algorithm 1 concat_reorder
Require: Sequence L at step t (Simplied as x), position index of masked tokens Mt, cached Key Kt−1T\Mt−1 and Value Vt−1T\Mt−1
1: x′ ← x[Mt−1] ▹ Use Mt−1: t−1 for one-step shift
2: PE′ ← [PE[T \ Mt−1]; PE[Mt−1]] ▹ Positional embeddings: cached on left, uncached on right
3: QtMt, KtMt, VtMt ← T(x′) ▹ T: Calculation in Transformer to get Q, K and V
4: KtT ← Concat(Kt−1T\Mt−1, KtMt), VtT ← Concat(Vt−1T\Mt−1, VtMt) ▹ Get all K and V
5: KtT ← Reorder(KtT, I′), VtT ← Reorder(VtT, I′) ▹ I′: The index of T \ Mt in the [x[T \ Mt−1] ; x[Mt−1]]
6: p′ ← A(QtMt, KtT, VtT) ▹ A: Attention calculation
7: p ← Scatter(p′, Mt) ▹ Put the token logits back to the original position
8: Return p, KtT, VtT

4.2.10. Dream模型的特殊设计 (附录 B)

Dream模型是从预训练的自回归模型改编而来，其输出位置与下一个词元的概率对齐，而非传统掩码扩散模型中当前词元的位置。这导致了缓存策略的差异。

三种缓存策略变体：

1. Un-Shift (无位移)： 缓存未位移的词元表示。对于第 $t$ 个词元，将其键和值 $\mathbf{K}^t, \mathbf{V}^t$ 存储在位置 $t$。

2. Right-Shift (右位移)： 由于隐藏状态对输入变化高度敏感，探索了一种右位移变体。对于第 $t$ 个词元，缓存其 $\mathbf{K}^{t+1}$ 和 $\mathbf{V}^{t+1}$。

3. Un&Right-Shift (无位移与右位移组合)： 更严格的变体，缓存的条件是词元的输入稳定且已被解码。对于第 $t$ 个词元，仅在其输入固定且被解码后才缓存其特征。

   一步延迟在此处仍然适用。例如，在右位移变体中，第 $t$ 个词元在下一步以 $t+1$ 位置输入模型，然后缓存其输出 $\mathbf{K}^{t+1}$ 和 $\mathbf{V}^{t+1}$。 实验（Table 4）显示 Un&Right-Shift 性能最佳，但与 concat_reorder 不兼容，因此主实验中采用 Un-Shift。

5. 实验设置

5.1. 数据集

本文在 LLaDA [37] 和 Dream [52] 的原始评估基准上进行了综合评估，涵盖了以下任务类型：

• 通用语言理解 (General Language Understanding): MMLU [23], GPQA [39]。
• 数学推理 (Mathematical Reasoning): GSM8K [10], Math500 [30]。
• 代码生成 (Code Generation): HumanEval [8], MBPP [4]。

评估方式：

• LLaDA: 采用零样本 (zero-shot) 评估。与直接比较词元概率的多项选择评估不同，本文要求模型生成答案字母，并将生成的答案与真实答案匹配。
• Dream: 遵循Dream的评估设置，采用少量样本上下文学习 (few-shot in-context learning) [6] 方式。

5.2. 评估指标

对论文中出现的每一个评估指标，都将提供概念定义、数学公式（如果适用）和符号解释。

5.2.1. 准确率 (Accuracy)

概念定义： 准确率用于衡量模型预测结果与真实标签（或真实答案）一致的比例。在多项选择题或生成特定答案的任务中，它表示模型给出正确答案的频率。对于生成任务，通常需要通过匹配生成文本与标准答案来判断。 数学公式： $$\text{Accuracy} = \frac{\text{Number of Correct Predictions}}{\text{Total Number of Predictions}}$$ 符号解释：

• Number of Correct Predictions: 模型给出正确答案的样本数量。
• Total Number of Predictions: 总共进行预测的样本数量。

5.2.2. Tokens/s (每秒词元数)

概念定义： Tokens/s 是衡量模型推理速度的指标，表示模型每秒能处理（生成或解码）的词元数量。它直接反映了模型的吞吐量和效率。 数学公式： 该指标通常直接由模型推理时间计算得出。 $$\text{Tokens/s} = \frac{\text{Total Tokens Processed}}{\text{Total Inference Time (seconds)}}$$ 符号解释：

• Total Tokens Processed: 在一次推理中处理的词元总数（例如，生成序列的总长度）。
• Total Inference Time (seconds): 完成词元处理所需的总时间（以秒为单位）。

5.2.3. 缓存率 (Cache Ratio)

概念定义： 缓存率是衡量缓存机制有效性的一个指标，表示在整个去噪过程中，有多少比例的词元键值对是从缓存中复用而不是重新计算的。高缓存率意味着更多的计算被节省。 数学公式： $$\text{Cache Ratio} = \frac{1}{T} \sum_{i=1}^{T} \frac{\left| \mathcal{T}_i^{\mathrm{cache}} \right|}{\left| \mathcal{T}_i \right|}$$ 符号解释：

• $T$: 去噪过程中的总步数。
• $i$: 当前去噪步骤的索引，从 1 到 $T$。
• $\left| \mathcal{T}_i \right|$: 在步骤 $i$ 处理的词元总数（通常是整个序列的长度 $L$）。
• $\left| \mathcal{T}_i^{\mathrm{cache}} \right|$: 在步骤 $i$ 中，其KV对从缓存中复用的词元子集数量。在 dKV-Cache 中，这对应于 $|\mathcal{T} \setminus \mathcal{M}_i|$ 或 $|\mathcal{M}_i|$，具体取决于缓存策略。

5.3. 对比基线

本文选择了几种步数较少的采样方法作为基线，这些方法的采样速度与本文方法相当或更慢：

• Half-Steps (一半步数): 使用 $50\%$ 的去噪步数。
• Few-Steps (少量步数): 使用 $62.5\%$ 的去噪步数。 这些基线用于展示，在相似的速度下，本文的方法能否提供更好的性能。

5.4. 配置细节

• 硬件平台： LLaDA 在 A6000 GPU 上进行测试，Dream 在 H20 GPU 上进行测试。

• LLaDA 实验配置 (Table 5): 以下是原文 Table 5 的结果：

  Remasking	Base (random / confidence) Configuration	Few-Steps (random) Steps T	dKV-Cache-Greedy (random) Cache Interval	dKV-Cache-Greedy (random) Window Size	Half-Steps (confidence) Steps T	dKV-Cache-Decode (confidence) Cache Interval
  MMLU	L=32, T=32, B=16	T=20	2	4	T=16	8
  GSM8K	L=256, T=256, B=32	T=160	2	4	T=128	8
  Math500	L=256, T=256, B=64	T=160	2	4	T=128	8
  GPQA	L=128, T=128, B=64	T=80	2	4	T=64	8
  HumanEval	L=512, T=512, B=32	T=320	2	4	T=256	8
  MBPP	L=512, T=512, B=32	T=320	2	2	T=256	8

  符号解释：

  • Remasking: 重新掩码策略（随机或基于置信度）。
  • $L$: 解码序列长度。
  • $T$: 去噪总步数。
  • $B$: 批处理大小 (Batch Size)。
  • Cache Interval: 缓存刷新间隔（每隔多少步刷新一次）。
  • Window Size: dKV-Cache-Greedy 中局部窗口的大小。

• Dream 实验配置： 遵循Dream的原始评估流程，并采用少量样本上下文学习 (few-shot in-context learning)。由于Dream模型是从自回归模型改编而来，其输出位置与下一个词元的概率对齐，而非传统掩码扩散模型中当前词元的位置，这使得缓存策略有所不同。

6. 实验结果与分析

6.1. 核心结果分析

6.1.1. LLaDA 模型上的性能与速度

以下是原文 Table 1 的结果：

分析：

• dKV-Cache-Greedy (随机重新掩码): 在大多数基准测试中，dKV-Cache-Greedy 的性能优于少步数的基线模型（如Few-Steps），除了 HumanEval。在 dKV-Cache-Greedy w/ Window 中，引入轻量级缓存窗口带来了显著的性能提升，而计算开销可忽略不计。
• dKV-Cache-Decode (基于置信度的重新掩码): 实现了几乎无损的性能，并具有较高的缓存率和较少的刷新步骤。在所有策略中，dKV-Cache-Decode 提供了最佳的权衡，在不牺牲性能的情况下验证了KV-Cache在DLMs中的可行性。
• 速度提升：两种变体都实现了显著的推理速度提升，例如 dKV-Cache-Decode 在 MMLU 上达到了 $2.35 \times$ 的加速。
• 由于 dKV-Cache-Greedy 依赖预定义的解码顺序（例如随机），在低置信度重新掩码时准确率下降明显，因此后续实验更侧重 dKV-Cache-Decode。

6.1.2. Dream 模型上的性能与速度

以下是原文 Table 2 的结果：

分析：

• dKV-Cache-Prefill 的显著优势： 由于Dream模型采用少量样本上下文学习 (few-shot in-context learning)，需要较长的输入上下文 (prefill length)。在这种设置下，dKV-Cache-Prefill 提供了巨大的速度提升，例如在 MMLU 和 GPQA 上实现了高达 $10 \times$ 的加速 (见 Table 3)。
• 解码长度和去噪步数的影响： 随着去噪步数的增加，本文方法相对于基线模型的优势更明显。例如，在 GSM8K (解码长度 $L=256$)，当去噪步数从 $T=256$ 减少到 $T=128$ 时，dKV-Cache-Decode 的加速比从 $2.09 \times$ 增加到 $4.13 \times$，dKV-Cache-PD 的加速比从 $3.55 \times$ 增加到 $7.11 \times$。
• 性能提升：在某些情况下，例如 GSM8K (解码长度 $L=256$, $T=128$)，dKV-Cache 不仅加速了 $6.6 \times$，还将性能从基线的 46.63 显著提升到 63.31 (提升了 16.68)。

6.1.3. 长预填充设置下的性能

以下是原文 Table 3 的结果：

分析：

• 在长预填充 (long prefill) 设置下，dKV-Cache-Prefill 提供了最显著的速度提升，例如在 GPQA (5-shot, $L=128$, $T=64$) 上实现了 $18.91 \times$ 的加速。这证明了其在处理长上下文时的优越性。
• dKV-Cache-Decode 也提供了可观的加速，同时性能损失很小。

6.2. 消融实验/参数分析

6.2.1. 一步延迟缓存的影响

下图（原文 Figure 3）展示了一步延迟缓存 (one-step delayed caching) 的影响：

分析：

• 无延迟缓存 (Without delay): 在缓存率较低时性能尚可，但随着缓存率的增加，性能迅速下降，最终准确率趋近于零。这表明立即缓存未完全稳定的表示会引入噪声，导致生成质量崩溃。
• 一步延迟缓存 (With one-step delay): 引入一步延迟显著稳定了生成质量。即使在很高的缓存率下，模型也能保持几乎无损的性能。这证实了一步延迟对于 dKV-Cache 在扩散语言模型中有效运行至关重要。

6.2.2. 解码长度、去噪步数和刷新步数的影响

下图（原文 Figure 4）展示了 dKV-Cache-Decode (左图) 和 dKV-Cache-Greedy (右图) 在 GSM8K 上，针对不同设置 (解码长度 $L$, 采样步数 $S$, 刷新间隔 refresh intervals 和窗口大小 window size) 的影响：

分析：

• 解码鲁棒性 (Decoding robustness): 模型的性能受解码步数的影响不大，表明 dKV-Cache 在不同生成长度和采样步数下都具有强大的鲁棒性。
• 增强长文本生成 (Enhanced long-form generation): 在涉及较长输出的任务中（例如 $L=512$），dKV-Cache 甚至优于基线模型，提升了生成质量（例如 GSM8K 从 $80.97\%$ 提高到 $83.13\%$，HumanEval 从 $39.63\%$ 提高到 $46.34\%$）。这暗示现有DLMs在双向注意力中可能存在上下文信息利用不足的问题。
• 少量刷新步骤的有效性 (Effectiveness with few refreshes): 即使刷新频率不高（例如每16步刷新一次），性能下降仍然很小。
• 局部窗口的有效性 (Effectiveness of local windows): dKV-Cache-Greedy 中整合局部窗口显著提升了性能，且只引入了极小的额外计算成本。

6.2.3. 内存和速度分析

下图（原文 Figure 5）展示了 dKV-Cache-Decode 和 dKV-Cache-Greedy 的速度 (Tokens/s) 和内存 (Peak Memory) 使用情况，针对不同的解码和预填充长度，以及刷新步数（2和8）：

分析：

• 显著加速： dKV-Cache 实现了显著的推理加速，范围从 $1.75 \times$ 到 $3.3 \times$。
• 内存适度增加： 内存使用量仅适度增加。
• dKV-Cache-Greedy 的潜力： dKV-Cache-Greedy 展现了更高的推理加速潜力，而 dKV-Cache-Decode 的加速上限较低。在相同的刷新间隔下，dKV-Cache-Greedy 始终比 dKV-Cache-Decode 实现更高的加速比。

6.2.4. Dream模型中不同缓存策略的比较

以下是原文 Table 4 的结果：

分析：

• Un&Right Shift 策略（输入稳定且解码完成后才缓存）在 MMLU 和 GSM8K 上表现最佳。
• Right-Shift 策略导致了显著的性能下降，尤其是在 GSM8K 上。
• 尽管 Un&Right-Shift 表现最优，但由于其与 concat_reorder 机制不兼容，本文主实验中使用了 Un-Shift 策略。

6.2.5. 批处理大小对速度的影响

下图（原文 Figure 7）展示了批处理大小 (batch size) 对解码速度 (Tokens/s) 的影响。评估在单个NVIDIA H20上对LLaDA模型进行，预填充长度固定为100个词元。

分析：

• 单批次效率低下： 在批处理大小为1时，由于索引操作（收集和散布）产生大量小且不连续的内存访问，导致推理过程受内存限制 (memory-bound)。此时，GPU的计算单元空闲等待数据，有缓存的推理甚至可能比无加速的基线更慢。
• 大批次加速效果显著： 增加批处理大小可以有效解决这个问题，实现更大、更稳定的加速比。这表明 dKV-Cache 在高吞吐量场景下能更好地发挥优势。

6.3. 案例研究 (Case Study)

6.3.1. LLaDA-8B-Instruct (Table 6)

以下是原文 Table 6 的结果：

分析： 在数学推理问题中，不同 dKV-Cache 配置（包括 dKV-Cache-Decode 和 dKV-Cache-Greedy 及其刷新间隔）都能生成正确的答案，并清晰地展示了解题步骤。这表明 dKV-Cache 在加速推理的同时，能够保持模型生成内容的准确性和逻辑性。

6.3.2. Dream-Instruct-7B (Table 7)

以下是原文 Table 7 的结果：

分析： 在故事生成任务中，Dream-Instruct-7B 基线模型和两种 dKV-Cache 变体 (dKV-Cache-Decode, dKV-Cache-Prefill) 都能生成连贯且以“Joey和Rachel最终结婚”结尾的故事。这进一步证明 dKV-Cache 在保持生成质量方面表现出色。

6.3.3. Dream-Instruct-7B (Table 8)

以下是原文 Table 8 的结果：

分析： 与 LLaDA 模型的案例研究类似，Dream-Instruct-7B 在数学问题上，无论基线还是 dKV-Cache 变体，都能准确计算并给出正确答案，且推理过程清晰。这再次强调 dKV-Cache 在保持模型性能的同时，实现了显著的推理加速。

6.4. 总结

实验结果有力地验证了 dKV-Cache 在DLMs上的有效性。它不仅在多种任务和模型上实现了显著的推理加速（2-10倍），而且在大多数情况下保持了与原始模型相当甚至更好的性能，尤其是在长序列生成中。一步延迟缓存机制和针对预填充词元的特殊处理是其成功的关键。虽然 dKV-Cache-Greedy 在某些情况下会有性能下降，但其在降低时间复杂度方面的潜力使其仍具有吸引力。批处理大小对加速效果有显著影响，表明系统级优化仍有潜力。

7. 总结与思考

7.1. 结论总结

本文首次提出了一种针对扩散语言模型 (DLMs) 的键值缓存 (KV-Cache) 机制——延迟KV-Cache (dKV-Cache)，解决了DLMs因非自回归架构和双向注意力机制而导致的推理速度慢的问题。该方法的核心洞察是，通过观察词元表示在扩散过程中的动态变化，发现即使在双向注意力下，键 (Key) 和值 (Value) 状态也可以通过延迟和有条件的方式进行复用。

本文设计了两种互补的 dKV-Cache 变体：

1. dKV-Cache-Decode: 提供了几乎无损的性能加速，在长序列生成任务上甚至能提升性能，暗示现有DLMs可能未充分利用上下文信息。

2. dKV-Cache-Greedy: 采用更激进的缓存策略，以牺牲部分性能为代价，实现了更高的加速比和 $\mathcal{O}(L^2)$ 的时间复杂度。

   通过在LLaDA和Dream等7B规模的DLMs上进行广泛实验，dKV-Cache 在通用语言理解、数学推理和代码生成等多个基准测试中实现了 $2 \text{ - } 10 \times$ 的推理加速，显著缩小了DLMs与自回归语言模型 (ARs) 之间的推理速度差距。值得注意的是，该方法无需对现有DLMs进行重新训练即可应用。

7.2. 局限性与未来工作

7.2.1. 局限性

本文指出的主要局限性在于其主要关注算法设计本身。尽管 dKV-Cache 在算法层面引入了有效的缓存机制，但扩散语言模型在系统层面仍有巨大的优化空间。

7.2.2. 未来工作

作者认为未来的研究方向应该结合算法创新与系统级优化，以释放DLMs的进一步效率提升和性能改进。这包括：

• 内存管理 (Memory Management): 优化缓存的内存占用和访问模式。
• 并行化 (Parallelism): 探索更高效的并行计算策略。
• 硬件感知执行 (Hardware-aware Execution): 考虑底层硬件特性进行优化，例如针对GPU架构的内存访问模式优化。

7.3. 个人启发与批判

7.3.1. 个人启发

1. 突破思维定式： 这篇论文的思路非常具有启发性。传统观点认为双向注意力模型无法有效利用KV-Cache，但作者通过深入分析词元表示的动态性，找到了“延迟和条件缓存”这一巧妙的解决方案。这表明在看似不兼容的技术之间，通过对底层机制的深刻理解，仍有可能找到创新的结合点。
2. DLMs的潜力： 过去DLMs因推理速度慢而备受诟病，限制了其在实际应用中的竞争力。dKV-Cache 极大地提升了DLMs的推理效率，使其有望在更多场景下与ARs竞争，尤其是在需要并行生成或对全局上下文有更高要求的任务中。
3. 无训练加速的价值： dKV-Cache 无需重新训练现有DLMs即可部署，大大降低了应用门槛和成本。这对于已经投入大量资源训练的DLMs而言，是一个巨大的福音，能够快速实现性能提升。
4. 长序列性能提升： dKV-Cache-Decode 在长序列上甚至能提升性能的发现非常有趣。这可能暗示在扩散模型的去噪过程中，即使是原始模型，其对长上下文信息的利用也存在冗余或不足，而通过缓存机制的结构化信息复用反而能带来更好的效果。这为未来DLMs的架构设计和训练策略提供了新的研究方向。

7.3.2. 批判与潜在改进

1. dKV-Cache-Greedy 的性能下降： dKV-Cache-Greedy 在追求极致速度的同时带来了性能下降。这表明激进的缓存策略可能引入了过多的“陈旧”信息或丢失了关键的动态信息。未来的工作可以探索更智能的贪婪策略，例如，基于不确定性或信息增益来选择哪些词元可以被激进缓存，从而在速度和性能之间找到更好的平衡。
2. 系统级优化的重要性： 作者明确指出算法设计是目前的重点，系统级优化是未来工作。然而，concat_reorder 机制中涉及的“收集”和“散布”操作在实际实现中可能仍是性能瓶颈，尤其是在大模型和长序列下。未来的研究应深入探讨如何结合编译器优化、新的硬件指令集或更底层的GPU编程技术来提升这些非连续内存访问的效率。例如，是否可以设计一种硬件友好的数据结构，使得K/V存储本身就是半连续的，减少运行时重排的开销。
3. 刷新策略的动态性： 当前的缓存刷新机制是基于固定间隔 $N$。能否开发一种动态刷新策略，根据模型在去噪过程中的置信度、词元表示的变化程度或预测误差来自适应地决定何时刷新缓存？这将可能进一步提升效率和性能。
4. 对不同类型扩散模型的通用性： 本文主要关注掩码离散扩散语言模型。dKV-Cache 是否能同样有效地应用于其他类型的DLMs，例如连续嵌入空间上的扩散模型，或者更复杂的生成范式？这需要进一步的探索和验证。
5. 内存占用与扩展性： 尽管本文提到内存占用适度增加，但对于超长上下文（如百万词元级别）的生成，即使是“适度增加”也可能成为瓶颈。结合现有的KV-Cache压缩技术 (如量化、稀疏化) 到 dKV-Cache 中，将是一个有前景的方向。