AiPaper
论文状态:已完成

dLLM-Cache: Accelerating Diffusion Large Language Models with Adaptive Caching

发表:2025/05/17
原文链接PDF 下载
价格:0.10
价格:0.10
已有 4 人读过
本分析由 AI 生成,可能不完全准确,请以原文为准。

TL;DR 精炼摘要

本文针对扩散大语言模型(dLLMs)推理延迟高问题,提出未经训练的自适应缓存框架dLLM-Cache。其结合长间隔提示缓存与基于特征相似性的部分响应更新,实现了对中间计算的高效重用,在LLaDA 8B和Dream 7B上最高加速9.1倍,同时保证输出质量。

摘要

Autoregressive Models (ARMs) have long dominated the landscape of Large Language Models. Recently, a new paradigm has emerged in the form of diffusion-based Large Language Models (dLLMs), which generate text by iteratively denoising masked segments. This approach has shown significant advantages and potential. However, dLLMs suffer from high inference latency. Traditional ARM acceleration techniques, such as Key-Value caching, are incompatible with dLLMs due to their bidirectional attention mechanism. To address this specific challenge, our work begins with a key observation that dLLM inference involves a static prompt and a partially dynamic response, where most tokens remain stable across adjacent denoising steps. Based on this, we propose dLLM-Cache, a training-free adaptive caching framework that combines long-interval prompt caching with partial response updates guided by feature similarity. This design enables efficient reuse of intermediate computations without compromising model performance. Extensive experiments on representative dLLMs, including LLaDA 8B and Dream 7B, show that dLLM-Cache achieves up to 9.1 x speedup over standard inference without compromising output quality. Notably, our method brings dLLM inference latency close to that of ARMs under many settings. Codes are provided in the supplementary material and will be released publicly on GitHub.

思维导图

论文精读

中文精读

1. 论文基本信息

1.1. 标题

dLLM-Cache: Accelerating Diffusion Large Language Models with Adaptive Caching (dLLM-Cache:通过自适应缓存加速扩散大语言模型)

1.2. 作者

Zhiyuan Liu, Yicun Yang, Yaojie Zhang, Junjie Chen, Chang Zou, Qingyan Wei, Shaobo Wang, Linfeng Zhang。 作者主要来自上海交通大学人工智能学院和EPIC实验室,以及电子科技大学。

1.3. 发表期刊/会议

该论文发布于 arXiv 预印本平台。

1.4. 发表年份

2025年5月17日。

1.5. 摘要

长期以来,自回归模型(Autoregressive Models, ARMs)主导着大语言模型(Large Language Models, LLMs)领域。然而,最近出现了一种新的范式——基于扩散的大语言模型(diffusion-based Large Language Models, dLLMs),它通过迭代地去噪掩码(masked segments)片段来生成文本。这种方法展现出显著的优势和潜力。然而,dLLMs 存在推理延迟高(high inference latency)的问题。传统的ARM加速技术,例如键值缓存(Key-Value caching),由于dLLMs的双向注意力(bidirectional attention)机制,与dLLMs不兼容。

为了解决这一特定挑战,本文工作首先提出了一个关键观察:dLLM推理涉及一个静态的提示(static prompt)和一个部分动态的响应(partially dynamic response),其中大多数词元(tokens)在相邻的去噪步骤(denoising steps)中保持稳定。基于此,本文提出了 dLLM-Cache,一个免训练(training-free)的自适应缓存框架(adaptive caching framework),它结合了长间隔提示缓存(long-interval prompt caching)和由特征相似性(feature similarity)引导的部分响应更新(partial response updates)。这种设计能够在不损害模型性能的情况下,高效地重用中间计算(intermediate computations)。

在代表性的dLLMs(包括 LLaDA 8BDream 7B)上进行的大量实验表明,dLLM-Cache 在不损害输出质量(output quality)的情况下,实现了高达 9.1 倍的速度提升(speedup)。值得注意的是,在许多设置下,本文方法使dLLM的推理延迟接近ARM的水平。代码已在补充材料中提供,并将公开发布在GitHub上。

1.6. 原文链接

原文链接: https://arxiv.org/abs/2506.06295 PDF 链接: https://arxiv.org/pdf/2506.06295v1.pdf

2. 整体概括

2.1. 研究背景与动机

核心问题: 扩散大语言模型(dLLMs)在推理过程中面临严重的延迟问题,这极大地阻碍了其在实际应用中的广泛采用。

问题重要性:

  1. dLLMs的潜力: 尽管自回归模型(ARMs)目前占据主导地位,但dLLMs作为一种新兴范式,展现出巨大的潜力和优势,例如在处理“逆转诅咒”(reversal curse)等特定挑战上可能超越ARMs,预示着其在复杂语言任务中作为基础LLM的竞争力。
  2. 现有加速技术的局限性:
    • 与ARMs不兼容: 传统ARMs的加速技术(如键值缓存 Key-Value caching)依赖于因果注意力(causal attention),而dLLMs采用双向注意力(bidirectional attention)机制,导致这些技术无法直接应用于dLLMs。

    • 与视觉扩散模型差异: 针对通用扩散模型(主要为视觉领域)的加速技术,如减少采样步骤或通用特征缓存,通常采用统一的策略处理所有词元。然而,语言任务具有其独特的结构(静态提示 static prompt 和动态响应 dynamic response),使得这些通用策略在dLLMs上效果不佳。

      现有研究挑战或空白(Gap): 缺乏专门针对dLLMs独特结构(静态提示和部分动态响应)的、高效且不损害性能的推理加速方法。当前方法未能充分利用dLLMs推理过程中存在的计算冗余。

本文的切入点或创新思路: 本文通过对dLLMs推理过程的深入分析,发现其存在两种特定的计算冗余:

  1. 提示冗余: 输入提示词元在整个推理过程中保持不变,但其内部表示(如注意力输出 attention output)却在每个去噪步骤中被重复计算。

  2. 响应动态与冗余: 生成的响应特征虽然会演变,但这种演变是稀疏的。在相邻的去噪步骤之间,大部分响应词元的特征变化不大,这意味着并非所有响应词元都需要在每一步都完全重新计算。

    基于这些观察,本文提出了一种差异化的、自适应的缓存策略,以精确地针对这些冗余进行优化。

2.2. 核心贡献/主要发现

  1. 识别并刻画了dLLM推理中的双重计算冗余: 即准静态提示(quasi-static prompt)和动态响应(dynamic response)冗余。
  2. 提出了 dLLM-Cache 框架: 一个免训练(training-free)的自适应缓存框架,它结合了长间隔提示缓存和短间隔、由相似性引导的响应词元部分更新策略。
  3. 引入了 V-verify 机制: 一种轻量级但有效的机制,通过计算相邻去噪步骤间值(Value)向量的余弦相似度(cosine similarity)来识别变化最大的词元进行部分更新。这一机制基于经验发现:值向量的变化与整个词元演变(即注意力输出和前馈网络输出 FFN Output 的变化)高度相关。
  4. 实验验证了 dLLM-Cache 的有效性:LLaDA 8BDream 7B 等代表性dLLMs上进行广泛实验,结果显示在不损害生成质量的前提下,推理速度提升高达 9.1 倍,并显著降低了计算成本(FLOPs),使得dLLM的推理延迟在许多设置下接近于ARMs。

3. 预备知识与相关工作

3.1. 基础概念

3.1.1. 自回归模型 (Autoregressive Models, ARMs)

ARMs 是当前大语言模型的主流范式。它们通过在给定前序词元(tokens)的情况下,逐个预测下一个词元的方式来生成文本。这种生成方式是顺序的,并且通常依赖于因果注意力(causal attention)机制,这意味着模型在生成当前词元时只能看到之前的词元,而不能看到未来的词元。代表模型包括 GPT 系列和 LLaMA 系列。

3.1.2. 扩散模型 (Diffusion Models, DMs)

DMs 是一类生成模型,最初主要应用于连续数据(如图像)。它们的工作原理是学习如何逆转一个数据损坏过程。在训练阶段,模型逐渐向数据中添加噪声,直到数据变为纯噪声。在推理阶段,模型则学习如何从噪声中逐步恢复(去噪)出原始的干净数据。

3.1.3. 掩码扩散模型 (Masked Diffusion Models, MDMs)

MDMs 是专门为离散数据(如文本)设计的扩散模型变体。与图像扩散模型在连续空间中添加和移除高斯噪声不同,MDMs 通过迭代地掩盖(mask)和预测文本序列中的词元来工作。它们从一个完全被掩码的序列开始,在每个去噪步骤中预测并填充一部分被掩码的词元,逐步将序列转化为有意义的文本。LLaDADreamdLLMs 就是 MDMs 的具体实例。

3.1.4. 双向注意力 (Bidirectional Attention)

Transformer 架构中的注意力机制允许模型在处理序列中的某个词元时,同时考虑到序列中所有其他词元(包括其前序和后序词元)的信息。在dLLMs中,由于其迭代去噪的特性,模型通常采用双向注意力,这意味着在每个去噪步骤中,模型可以同时看到提示和当前响应的所有词元,从而更好地理解上下文以进行去噪。这与ARMs的单向(因果)注意力形成对比。

3.1.5. 推理延迟 (Inference Latency)

推理延迟是指模型从接收输入到生成完整输出所需的时间。在实际应用中,低延迟对于用户体验至关重要。dLLMs由于其迭代的去噪过程,通常比ARMs具有更高的推理延迟。

3.1.6. 键值缓存 (Key-Value (KV) Caching)

KV缓存是ARMs推理加速的关键技术。在Transformer模型中,每个词元的查询(Query, Q)、键(Key, K)和值(Value, V)向量在注意力计算中发挥作用。对于ARMs,在生成序列的每个新词元时,之前词元计算得到的KKVV向量可以被缓存并重用,而无需在每一步都重新计算整个序列的KKVV。这显著减少了后续词元生成时的计算量。

3.1.7. 浮点运算次数 (Floating Point Operations, FLOPs)

FLOPs 是衡量计算复杂度的指标,表示模型执行的浮点运算的总次数。在深度学习中,FLOPs 常用于评估模型的计算成本和效率。更低的FLOPs通常意味着更快的推理速度和更低的能耗。

3.1.8. 每秒词元数 (Tokens Per Second, TPS)

TPS 是衡量模型推理吞吐量(throughput)的指标,表示模型每秒能够生成的词元数量。更高的TPS意味着更快的生成速度,是衡量推理效率的关键指标。

3.2. 前人工作与技术演进

3.2.1. 大语言模型范式演进

  • 自回归模型 (ARMs) 的主导地位: ARMsRadford (2018)GPT-2 (Radford et al., 2019)GPT-3 (Brown, 2020)ChatGPT (OpenAI, 2022) 等,通过其强大的文本生成能力和对各种语言任务的通用性,成为了LLMs领域的主流范式。它们的成功主要归因于其在大规模文本数据上的预训练和基于Transformer的架构。
  • 扩散模型在语言领域的兴起: 扩散模型最初在图像生成领域取得了巨大成功 (Rombach et al., 2022, Peebles and Xie, 2023)。近年来,研究人员开始探索将扩散模型的思想应用于离散数据(如文本)。Masked Diffusion Models (MDMs) (Austin et al., 2021, Lou et al., 2023, Shi et al., 2024, Ou et al., 2024) 是这一方向的重要进展,它们通过迭代地去噪掩码片段来生成文本。 LLaDA (Nie et al., 2025)Dream (Ye et al., 2025) 是最新的代表性dLLMs,它们展现出与同等规模ARMs(如LLaMA3 8B)相媲美的性能,并且其固有的双向架构可能有助于解决ARMs的一些限制,例如“逆转诅咒” (Berglund et al., 2023)

3.2.2. 加速机制的发展

  • ARMs 的键值缓存 (Key-Value Caching for ARMs): 为了加速ARMs的推理,KV缓存 (Pope et al., 2023) 成为一种标准技术。它通过缓存前序词元的KeyValue状态,避免在生成后续词元时重复计算,从而显著提高推理效率。然而,由于dLLMs采用的是双向注意力机制,需要考虑所有词元(包括前序和后序)的上下文,这使得KV缓存与dLLMs不兼容。
  • 扩散模型的特征缓存 (Feature Caching for DMs): 在扩散模型领域,特征缓存技术旨在利用时间冗余性,即在不同的去噪步骤中重用中间特征。
    • 早期方法 (U-Net): DCace (Ma et al., 2024) 等方法最初关注于 U-Net 架构的扩散模型。
    • 扩散 Transformers 的发展: 随着扩散 Transformers (Peebles and Xie, 2023) 的出现,多种缓存范式被提出:
      • FORA (Selvaraju et al., 2024):对注意力层和MLP层采用静态间隔缓存。
      • ToCa 系列 (Zouetal.,2024a,b)(Zou et al., 2024a,b):通过动态更新缓存特征来减少信息损失。
      • TaylorSeer (Liu et al., 2025):引入“缓存-预测”(cache-then-forecast)范式,使用泰勒级数展开(Taylor series expansion)来预测特征而非直接重用。

3.2.3. 现有缓存技术对dLLMs的局限性

尽管扩散模型加速技术取得了进展,但现有的特征缓存方法对于dLLMs而言效果不佳。核心原因在于:这些方法大多是为视觉任务设计的,通常对所有词元应用统一的缓存策略,未能充分考虑到语言任务中独特的“静态提示”和“动态响应”结构。这种结构差异导致了在dLLMs高效加速方面存在关键的技术空白。

3.3. 差异化分析

本文提出的 dLLM-Cache 的核心创新点在于其差异化的自适应缓存策略,这与现有工作形成鲜明对比:

  • 针对dLLMs特定结构: dLLM-Cache 是第一个明确区分并利用dLLMs中“静态提示”和“部分动态响应”特性的缓存框架。它不像传统ARMKV缓存那样因双向注意力而失效,也不像现有通用扩散模型(主要为视觉设计)的特征缓存那样对所有词元一视同仁。
  • 长间隔提示缓存: 针对静态提示词元,dLLM-Cache 采用长间隔刷新策略,大幅减少了对这部分稳定特征的重复计算。这是以往方法所不具备的。
  • 自适应响应缓存与 V-verify 对于动态响应词元,dLLM-Cache 采用短间隔全刷新中间自适应部分更新相结合的策略。其核心是 V-verify 机制,通过轻量级的值(Value)向量相似性来高效识别并只更新那些变化最大的词元。这种基于特征相似性的动态选择性更新是其创新之处,确保了效率和准确性的平衡。
  • 免训练: dLLM-Cache 是一种免训练(training-free)的方法,可以直接应用于现有的dLLMs而无需重新训练模型,这大大降低了其应用成本和复杂性。

4. 方法论

4.1. 方法原理

dLLM-Cache 的核心思想是利用 dLLMs 推理过程中存在的两种主要计算冗余来加速推理:

  1. 静态提示的冗余: 输入的提示词元在整个生成过程中是固定不变的,其对应的中间特征在大多数去噪步骤中也高度稳定。

  2. 动态响应的稀疏性: 虽然响应词元在去噪过程中会演变,但这种演变并非均匀发生在所有词元上,而是高度稀疏的,即在相邻步骤中,只有一小部分响应词元会发生显著变化。

    基于这些观察,dLLM-Cache 采取了一种差异化的自适应缓存策略

  • 对于静态提示,采用长间隔缓存,大幅减少不必要的重复计算。

  • 对于动态响应,则结合短间隔的全量刷新基于特征相似性的自适应部分更新。其中,V-verify 机制通过计算值(Value)向量的余弦相似度来高效识别并只更新那些变化最大的响应词元。

    这种策略旨在最大限度地重用中间计算结果,同时确保模型性能不受影响。

4.2. 核心方法详解

4.2.1. dLLMs的推理过程概述

dLLMs 通过一个非自回归(non-autoregressive)、基于扩散(diffusion-based)的过程来生成文本。该过程迭代地将一个完全被掩码的序列去噪为目标输出。以 LLaDA 为例,其推理过程如下:

  1. 初始化: 给定一个提示 c=(c1,,cM)\mathbf{c} = (c_1, \ldots, c_M),模型生成一个长度为 LL 的响应 y=(y1,,yL)\mathbf{y} = (y_1, \ldots, y_L)。推理过程从一个完全由掩码词元 [MASK] 组成的序列 y(K)\mathbf{y}^{(K)} 开始,其中 KK 是总的去噪步骤数。 y(K)=([MASK],,[MASK]L times) \mathbf{y}^{(K)} = (\underbrace{[\mathtt{MASK}], \dots, [\mathtt{MASK}]}_{L \text{ times}})
  2. 迭代去噪: 在每个去噪步骤 kk (从 KK 递减到 0),一个掩码预测器 fϕf_\phi 估计给定提示 c\mathbf{c} 和当前中间状态 y(k)\mathbf{y}^{(k)} 下的干净序列分布: Pϕ(yc,y(k))=fϕ(c,y(k);ϕ) P_\phi(\mathbf{y} | \mathbf{c}, \mathbf{y}^{(k)}) = f_\phi(\mathbf{c}, \mathbf{y}^{(k)}; \phi)
  3. 预测干净序列: 从估计的分布中选择最可能的干净序列 y^(0)\hat{\mathbf{y}}^{(0)}y^(0)=argmaxyTLPϕ(yc,y(k)) \hat{\mathbf{y}}^{(0)} = \underset{\mathbf{y} \in \mathcal{T}^L}{\arg\max} P_\phi(\mathbf{y} | \mathbf{c}, \mathbf{y}^{(k)})
  4. 更新序列: 一个转换函数 SS 根据预测的 y^(0)\hat{\mathbf{y}}^{(0)} 和当前序列 y(k)\mathbf{y}^{(k)} 选择性地更新词元,生成下一个去噪步骤的序列 y(k1)\mathbf{y}^{(k-1)}y(k1)=S(y^(0),y(k),c,k) \mathbf{y}^{(k-1)} = S(\hat{\mathbf{y}}^{(0)}, \mathbf{y}^{(k)}, \mathbf{c}, k) 转换函数 SS 可能涉及基于置信度的重新掩码(remasking)或半自回归的块更新(semi-autoregressive block updates)。

这个迭代过程虽然灵活,但由于在每个步骤都需要对整个序列进行重复计算,导致推理延迟很高。

4.2.2. dLLM-Cache 框架概览

dLLM-Cache 是一个免训练的框架,通过以下方式管理缓存:

  • 缓存结构: 对于每个 Transformerll,它缓存键(Key, K(l)\mathbf{K}^{(l)})、值(Value, V(l)\mathbf{V}^{(l)})、注意力输出(Attention Output, AttnOut(l)\text{AttnOut}^{(l)})和前馈网络输出(FFN Output, FFNOut(l)\text{FFNOut}^{(l)})。这些特征被分别存储在提示缓存 Cp\mathcal{C}_p响应缓存 Cr\mathcal{C}_r 中。
  • 超参数: 框架由三个关键超参数控制:
    • 提示刷新间隔 KpK_p:控制提示缓存的刷新频率。
    • 响应刷新间隔 KrK_r:控制响应缓存全量刷新的频率。
    • 自适应更新比例 ρ[0,1]\rho \in [0, 1]:控制在自适应更新模式下,多少比例的响应词元需要重新计算。
  • 整体工作流程 (见 Figure 3):
    1. 初始化: 在去噪过程开始时(k=Kk=K),计算所有特征并填充提示缓存 Cp\mathcal{C}_p 和响应缓存 Cr\mathcal{C}_r

    2. 迭代步骤: 在每个后续去噪步骤 kk (从 K-11),对于每个 Transformerll

      • 提示部分: 如果当前步骤 kkKpK_p 的倍数(即 k0(modKp)k \equiv 0 \pmod{K_p}),则重新计算并更新 Cp\mathcal{C}_p;否则,重用缓存中的提示特征。
      • 响应部分: 如果当前步骤 kkKrK_r 的倍数(即 k0(modKr)k \equiv 0 \pmod{K_r}),则完全重新计算并更新 Cr\mathcal{C}_r;否则,执行自适应部分更新(详见下文)。
    3. 计算: 使用缓存或更新后的特征执行当前层的计算。

    4. 终止: 直到 k=0k=0 停止,并返回最终预测 y^(0)\hat{\mathbf{y}}^{(0)}

      下图(原文 Figure 3)展示了 dLLM-Cache 的整体流水线:

      Figure 3: The dLLM-Cache pipeline. Prompt features are updated with long intervals, while response features are updated adaptively based on the similarity between newly computed and cached \$\\mathbf {… Figure 3: The dLLM-Cache pipeline. Prompt features are updated with long intervals, while response features are updated adaptively based on the similarity between newly computed and cached V\mathbf{V} vectors. Response features of tokens with low similarity are updated, and the rest are reused.

4.2.3. 提示缓存管理 (Prompt Cache Management)

由于输入提示 c\mathbf{c} 在整个推理过程中是静态的,其对应特征具有高度的时间稳定性。dLLM-Cache 利用这一特性,通过维护一个长间隔提示缓存 Cp\mathcal{C}_p

  • 初始化: 在去噪过程的开始步骤 k=Kk=K 时,计算所有层的 Kp(l),Vp(l),AttnOutp(l),FFNOutp(l)\mathbf{K}_p^{(l)}, \mathbf{V}_p^{(l)}, \text{AttnOut}_p^{(l)}, \text{FFNOut}_p^{(l)},并将其存储在 Cp\mathcal{C}_p 中。

  • 周期性刷新: 在后续去噪步骤 kk 中,仅当 kk 是预设的提示刷新间隔 KpK_p 的倍数时(即 k0(modKp)k \equiv 0 \pmod{K_p}),才会重新计算并更新 Cp\mathcal{C}_p 中的提示特征。

    • 重要细节: 即使是提示部分的重新计算,其注意力输出 AttnOutp(l)\text{AttnOut}_p^{(l)} 和随后的 FFNOutp(l)\text{FFNOut}_p^{(l)} 的计算,也会考虑到响应缓存 Cr\mathcal{C}_r 中可能已经更新的键值特征。这确保了提示的表示能够感知到不断演变的响应。
  • 重用: 在其他步骤中(即 k≢0(modKp)k \not\equiv 0 \pmod{K_p}),直接重用 Cp\mathcal{C}_p 中缓存的提示特征。

    这种策略显著减少了处理静态提示的计算负担,尤其当 KpK_p 设置较大时。

4.2.4. 带有自适应更新的响应缓存 (Response Cache with Adaptive Updates)

与提示不同,响应特征 y(k)\mathbf{y}^{(k)} 会在不同去噪步骤中演变。然而,本文观察到,大多数响应词元的变化是渐进的,使得选择性更新成为可能。响应缓存 Cr\mathcal{C}_r 通过两种模式进行管理:

  • 全量刷新 (Full Refresh):
    • 在去噪过程的开始步骤 k=Kk=K 时,以及当 kk 是预设的响应刷新间隔 KrK_r 的倍数时(即 k0(modKr)k \equiv 0 \pmod{K_r}),会重新计算所有响应特征并更新 Cr\mathcal{C}_r
  • 自适应部分更新 (Adaptive Partial Update):
    • 在其他情况下(即 k≢0(modKr)k \not\equiv 0 \pmod{K_r}),dLLM-Cache 采用一种自适应策略:
      1. 识别变化(V-verify 机制): 对于每个响应词元 jj,计算其当前计算的值向量 vr,j(l)\mathbf{v}_{r,j}^{(l)} 与其缓存中对应的值向量 v~r,j(l)\tilde{\mathbf{v}}_{r,j}^{(l)} 之间的余弦相似度 sjs_jsj=(vr,j(l))v~r,j(l)vr,j(l)v~r,j(l) s_j = \frac{(\mathbf{v}_{r,j}^{(l)})^\top \tilde{\mathbf{v}}_{r,j}^{(l)}}{\|\mathbf{v}_{r,j}^{(l)}\| \|\tilde{\mathbf{v}}_{r,j}^{(l)}\|}

        • 符号解释:
          • sjs_j: 第 jj 个响应词元的当前值向量与其缓存值向量之间的余弦相似度。
          • vr,j(l)\mathbf{v}_{r,j}^{(l)}: 在 Transformerll 层,第 jj 个响应词元的当前计算出的值(Value)向量。
          • v~r,j(l)\tilde{\mathbf{v}}_{r,j}^{(l)}: 在 Transformerll 层,第 jj 个响应词元在上一个相关步骤中缓存的值(Value)向量。
          • \top: 向量转置运算。
          • \| \cdot \|: 向量的L2范数(欧几里得范数)。
        • 目的分析: 余弦相似度衡量了两个向量在方向上的接近程度。相似度越低(接近-1),表示两个值向量的方向变化越大,暗示该词元的特征发生了显著变化;相似度越高(接近1),表示方向变化越小,特征相对稳定。通过这种方式,V-verify 可以高效地量化每个词元的变化程度。
      2. 选择词元: 选取 ρL\lfloor \rho L \rfloor 个具有最低 sjs_j 值的词元。这些词元的索引集合构成 Tupdate\mathcal{T}_{\text{update}},表示它们是变化最显著的词元。

      3. 重新计算与重用: 仅对 Tupdate\mathcal{T}_{\text{update}} 中的词元,重新计算其 Kr(l),Vr(l),AttnOutr(l)\mathbf{K}_r^{(l)}, \mathbf{V}_r^{(l)}, \text{AttnOut}_r^{(l)}FFNOutr(l)\text{FFNOut}_r^{(l)}。对于不在 Tupdate\mathcal{T}_{\text{update}} 中的词元,则直接重用其缓存值。

      4. 更新缓存: 用新计算的值替换 Cr\mathcal{C}_r 中对应位置的缓存条目。

        这种自适应更新机制利用了时间稳定性来最小化计算量,并通过有针对性的特征刷新来保持生成质量。

4.2.5. V-verify 机制的依据

V-verify 机制之所以有效,是基于一个关键的经验观察,如 Figure 2 所示:响应词元的值(Value)向量或键(Key)向量在相邻去噪步骤中的变化(通过余弦相似度量化)与随后的注意力输出(Attention Output)和前馈网络输出(FFN Output)的变化之间存在强相关性。这意味着,监控早期阶段特征(如 VV 向量)的动态变化,可以有效地预测哪些更复杂的下游特征也可能发生了显著变化。

下图(原文 Figure 2)展示了响应词元的KKVV变化与其它特征变化的相关性:

Figure 2: Correlation of response tokens' K or V changes with other feature changes. We calculate the cosine similarity between the response tokens' \(\\mathbf { K }\) or \(\\mathbf { V }\) vectors and the… Figure 2: Correlation of response tokens' K or V changes with other feature changes. We calculate the cosine similarity between the response tokens' K\mathbf{K} or V\mathbf{V} vectors and their cached counterparts at adjacent steps, select the 25%25\% most dissimilar tokens, and compute the correlation between their similarity with (a) and (c) AttnOut, or (b) and (d) FFNOut across adjacent steps.

图中显示,VV向量与AttnOutFFNOut的变化有着密切的线性相关性,这支持了使用VV向量相似度作为识别变化词元的有效指标。

4.3. 算法流程 (Algorithm 1 概述)

本文的 dLLM-Cache 的核心算法 Algorithm 1 描述了整个推理过程,其中包含了上述的提示和响应缓存管理策略。

Algorithm 1 dLLM-Cache: Main Inference Algorithm (dLLM-Cache:主要推理算法)

  • 输入: 提示 c\mathbf{c},初始掩码序列 y(K)\mathbf{y}^{(K)},去噪步骤 KK,提示刷新间隔 KpK_p,响应刷新间隔 KrK_r,自适应更新比例 ρ\rho
  • 输出: 最终预测 y^(0)\hat{\mathbf{y}}^{(0)}
  1. 初始化缓存: 在步骤 k=Kk=K 时,通过 InitializeCache 函数(对应 Algorithm 2)计算 c\mathbf{c}y(K)\mathbf{y}^{(K)} 的所有特征,并存储在提示缓存 Cp\mathcal{C}_p 和响应缓存 Cr\mathcal{C}_r 中。
    • Algorithm 2 InitializeCache (初始化缓存):
      • 为每个Transformerll 初始化空的缓存结构 Kv_cacheattnmlp(对应K, V, AttnOut, FFNOut)。
      • 将提示 c\mathbf{c} 和初始掩码序列 y(K)\mathbf{y}^{(K)} 拼接作为第一个层的输入 xin\mathbf{x}_{\text{in}}
      • 遍历每个 Transformerll
        • 进行 LayerNorm,然后计算 QQ, KK, VV
        • KK, VV 分割为提示部分 (Kp,Vp\mathbf{K}_p, \mathbf{V}_p) 和响应部分 (Kr,Vr\mathbf{K}_r, \mathbf{V}_r),并分别存入 Cp[l][kv_cache]\mathcal{C}_p[l][\text{kv\_cache}]Cr[l][kv_cache]\mathcal{C}_r[l][\text{kv\_cache}]
        • 计算注意力输出 AttnOut,并分割为提示部分 (AttnOutp\text{AttnOut}_p) 和响应部分 (AttnOutr\text{AttnOut}_r),分别存入 Cp[l][attn]\mathcal{C}_p[l][\text{attn}]Cr[l][attn]\mathcal{C}_r[l][\text{attn}]
        • 计算 FFN 输出 FFNOut,并分割为提示部分 (FFNOutp\text{FFNOut}_p) 和响应部分 (FFNOutr\text{FFNOut}_r),分别存入 Cp[l][mlp]\mathcal{C}_p[l][\text{mlp}]Cr[l][mlp]\mathcal{C}_r[l][\text{mlp}]
        • 更新当前层的输出 xout\mathbf{x}_{\text{out}} 作为下一层的输入 xin\mathbf{x}_{\text{in}}
      • 返回初始化后的 Cp,Cr\mathcal{C}_p, \mathcal{C}_r
  2. 生成初始预测和更新序列: 使用模型 fϕf_\phi 生成基于 KK 步骤的预测 y^(0)\hat{\mathbf{y}}^{(0)},并通过转换函数 SS 得到 y(K1)\mathbf{y}^{(K-1)}
  3. 迭代去噪步骤:k=K1k=K-1 循环到 1
    • 将当前步骤的输入 xlayer_in\mathbf{x}_{\text{layer\_in}} 初始化为 [c;y(k)][\mathbf{c}; \mathbf{y}^{(k)}]
    • 遍历每个 Transformerll
      • 判断刷新条件:
        • refresh_prompt = (k(modKp)=0k \pmod{K_p} = 0)
        • refresh_response = (k(modKr)=0k \pmod{K_r} = 0)
      • 根据刷新条件选择缓存策略:
        • 情况 1:同时刷新提示和响应 (refresh_promptrefresh_response): 调用 FullRefresh 函数(对应 Algorithm 3),重新计算并更新所有提示和响应特征。
          • Algorithm 3 FullRefresh (全量刷新): 类似 InitializeCache,但使用当前层的输入 xin\mathbf{x}_{\text{in}} 进行计算,并更新相应的缓存。
        • 情况 2:只刷新提示 (refresh_promptnot refresh_response): 调用 RefreshPromptOnly 函数(对应 Algorithm 4),重新计算提示特征,重用响应特征。
          • Algorithm 4 RefreshPromptOnly (仅刷新提示):
            • 获取当前层的提示输入部分 xp_in\mathbf{x}_{p\_in}
            • 计算新的提示 QQ, KK, VV,并更新 Cp[l][kv_cache]\mathcal{C}_p[l][\text{kv\_cache}]
            • Cr[l][kv_cache]\mathcal{C}_r[l][\text{kv\_cache}] 中重用响应 KK, VV
            • 合并 KK, VV 向量,计算提示的注意力输出 AttnOutpAttnOut_p,并更新 Cp[l][attn]\mathcal{C}_p[l][\text{attn}]
            • Cr[l][attn]\mathcal{C}_r[l][\text{attn}] 中重用响应注意力输出 AttnOutrAttnOut_r
            • 合并 AttnOutp,AttnOutrAttnOut_p, AttnOut_r,计算残差连接后的 hh
            • 计算提示的 FFN 输出 FFNOutpFFNOut_p,并更新 Cp[l][mlp]\mathcal{C}_p[l][\text{mlp}]
            • Cr[l][mlp]\mathcal{C}_r[l][\text{mlp}] 中重用响应 FFNOutrFFNOut_r
            • 合并 FFNOutp,FFNOutrFFNOut_p, FFNOut_r,计算最终层输出 xout\mathbf{x}_{\text{out}}
        • 情况 3:只刷新响应 (not refresh_promptrefresh_response): 调用 RefreshResponseOnly 函数(对应 Algorithm 5),重新计算响应特征,重用提示特征。
          • Algorithm 5 RefreshResponseOnly (仅刷新响应):
            • 获取当前层的响应输入部分 xr_in\mathbf{x}_{r\_in}
            • Cp[l]\mathcal{C}_p[l] 中重用提示的 KK, VV, AttnOut, FFNOut
            • 计算新的响应 QQ, KK, VV,并更新 Cr[l][kv_cache]\mathcal{C}_r[l][\text{kv\_cache}]
            • 合并提示和响应的 KK, VV 向量,计算响应的注意力输出 AttnOutrAttnOut_r,并更新 Cr[l][attn]\mathcal{C}_r[l][\text{attn}]
            • 合并提示和响应的 AttnOut,计算残差连接后的 hh
            • 计算响应的 FFN 输出 FFNOutrFFNOut_r,并更新 Cr[l][mlp]\mathcal{C}_r[l][\text{mlp}]
            • 合并提示和响应的 FFNOut,计算最终层输出 xout\mathbf{x}_{\text{out}}
        • 情况 4:无须全量刷新(not refresh_promptnot refresh_response): 调用 AdaptiveUpdate 函数(对应 Algorithm 6),执行自适应部分更新。
          • Algorithm 6 AdaptiveUpdate (自适应更新):
            • Cp[l]\mathcal{C}_p[l] 中重用提示的 KK, VV, AttnOut, FFNOut
            • 如果 ρ>0\rho > 0
              • 计算当前响应输入的 VV 投影 Vrnew\mathbf{V}_r^{\text{new}}
              • Cr[l][kv_cache]\mathcal{C}_r[l][\text{kv\_cache}] 中获取缓存的响应 KK, VV
              • 使用余弦相似度 sj=(Vrnew[j])Vr[j]Vrnew[j]Vr[j]s_j = \frac{(\mathbf{V}_r^{\text{new}}[j])^\top \mathbf{V}_r[j]}{\|\mathbf{V}_r^{\text{new}}[j]\| \|\mathbf{V}_r[j]\|} 识别 ρy(k)\lfloor \rho |\mathbf{y}^{(k)}| \rfloor 个最低相似度的词元索引 Tupdate\mathcal{T}_{\text{update}}
              • 仅对 Tupdate\mathcal{T}_{\text{update}} 中的词元进行重新计算 (QQ, KK,注意VV总是使用VrnewV_r^new),并更新 Kv_cache 中的相应部分。
              • 计算 Tupdate\mathcal{T}_{\text{update}} 中词元的注意力输出 AttnOut_r_selected,并更新 attn 缓存中的相应部分。
              • 计算 Tupdate\mathcal{T}_{\text{update}} 中词元的 FFN 输出 FFNOut_r_selected,并更新 mlp 缓存中的相应部分。
            • 如果 ρ=0\rho = 0
              • 直接从 Cr[l]\mathcal{C}_r[l] 中重用响应的 AttnOutFFNOut
            • 合并所有输出,计算最终层输出 xout\mathbf{x}_{\text{out}}
      • 更新 xlayer_in\mathbf{x}_{\text{layer\_in}} 为当前层的输出,作为下一层的输入。
    • 使用模型 fϕf_\phi 生成预测 y^(0)\hat{\mathbf{y}}^{(0)},并通过转换函数 SS 得到 y(k1)\mathbf{y}^{(k-1)}
  4. 返回最终预测 y^(0)\hat{\mathbf{y}}^{(0)}

5. 实验设置

5.1. 数据集

实验主要在以下两个代表性的 dLLMs 上进行评估:

  • LLaDA 8B (Nieetal.,2025)(Nie et al., 2025)
  • Dream 7B (Yeetal.,2025)(Ye et al., 2025) 每个模型都包括 BaseInstruct 两种变体。

本文在多个基准测试(benchmarks)上进行了实验,这些基准测试涵盖了数学、科学和通用任务:

  • 数学与科学 (Mathematics & Science):
    • GSM8K (Cobbe et al., 2021):数学单词问题。
    • GPQA:通用问题回答。
    • Math:数学推理。
  • 通用任务 (General Tasks):
    • MMLU-pro:大规模多任务语言理解专业版。
    • MMLU:大规模多任务语言理解。
    • BBH:硬核大模型基准测试(Big-Bench Hard)。
  • 代码 (Code):
    • MBPP:手写代码问题生成(Mostly Basic Python Problems)。
    • HumanEval:代码生成和理解。
  • 长文本场景 (Long-Prompt Scenarios):
    • LongBench-HotpotQA (Baietal.,2023)(Bai et al., 2023):文档问答任务,用于评估长上下文理解能力。

具体实验设置 (取自原文附录 A Table 4 示例): 以下是 Instruct 模型在一些基准测试中的实验设置:

Task Steps Block Len Gen Len Few-shot
MMLU 3 3 3 5
MMLU-pro 256 256 256 0
Hellaswag 3 3 3 0
ARC-C 512 512 512 0
GSM8K 256 8 256 4
Math 256 256 256 0
GPQA 128 64 128 5
HumanEval 512 32 512 0
MBPP 512 32 512 3
  • Denoising Steps (Steps): 去噪迭代的总步数。
  • Block Length (Block Len): 在某些重掩码策略中,一次性处理的词元块的长度。
  • Generation Length (Gen Len): 生成序列的总长度。
  • Few-shot (Few-shot): 使用的少样本示例数量。
  • Remasking Strategy (Remasking Strategy): 所有模型默认使用低置信度重掩码策略,除非另有说明。

5.2. 评估指标

本文采用以下指标评估 dLLM-Cache 的加速效果和模型质量:

5.2.1. 推理效率指标

  1. 每秒词元数 (Tokens Per Second, TPS)

    • 概念定义: TPS 衡量模型每秒能够处理或生成的词元数量,反映了模型的推理吞吐量。它是一个直观的速度指标,TPS 值越高表示推理速度越快。
    • 数学公式: 无标准公式,通常计算为 Total Tokens Generated/Total Inference Time (seconds)\text{Total Tokens Generated} / \text{Total Inference Time (seconds)}
    • 符号解释:
      • Total Tokens Generated:模型在一次推理过程中生成的词元总数。
      • Total Inference Time:完成推理所需的总时间(以秒为单位)。
  2. 浮点运算次数 (Floating Point Operations, FLOPs)

    • 概念定义: FLOPs 衡量模型在一次推理过程中执行的浮点运算的总次数,用于量化模型的计算成本。FLOPs 值越低表示计算效率越高。
    • 数学公式: 无标准公式,通常通过特定的工具或方法对模型操作进行统计。
    • 符号解释: 无特定符号。
  3. 浮点运算速度 (Speed(FLOPs))

    • 概念定义: Speed(FLOPs) 指的是基于 FLOPs 减少的加速比,即基线模型的 FLOPs 除以优化后模型的 FLOPs
    • 数学公式: Speed(FLOPs)=Baseline FLOPsOptimized FLOPs\text{Speed(FLOPs)} = \frac{\text{Baseline FLOPs}}{\text{Optimized FLOPs}}
    • 符号解释:
      • Baseline FLOPs:基线模型(未加速)的浮点运算次数。
      • Optimized FLOPs:经过 dLLM-Cache 加速后模型的浮点运算次数。
  4. GPU 内存占用 (GPU Memory (GB))

    • 概念定义: 衡量模型在推理过程中占用的峰值 GPU 内存。较低的内存占用有助于在资源受限的环境中部署模型。
    • 数学公式: 无标准公式,通常通过系统监控工具获取。
    • 符号解释: 无特定符号。

5.2.2. 模型性能指标

  1. 任务性能分数 (Score/Accuracy/F1 Score)
    • 概念定义: 根据具体任务采用相应的评估指标。例如,GSM8K 任务使用准确率 (Accuracy) 来衡量模型解决数学问题的正确性;LongBench-HotpotQA 任务可能使用 F1 分数来衡量答案的召回率和精确率。
    • 数学公式 (以准确率 Accuracy 为例): Accuracy=Number of Correct PredictionsTotal Number of Predictions \text{Accuracy} = \frac{\text{Number of Correct Predictions}}{\text{Total Number of Predictions}}
    • 符号解释:
      • Number of Correct Predictions:模型正确预测的样本数量。
      • Total Number of Predictions:总的预测样本数量。
    • 数学公式 (以 F1 Score 为例): F1 Score=2×Precision×RecallPrecision+Recall \text{F1 Score} = 2 \times \frac{\text{Precision} \times \text{Recall}}{\text{Precision} + \text{Recall}} 其中, Precision=True PositivesTrue Positives+False Positives \text{Precision} = \frac{\text{True Positives}}{\text{True Positives} + \text{False Positives}} Recall=True PositivesTrue Positives+False Negatives \text{Recall} = \frac{\text{True Positives}}{\text{True Positives} + \text{False Negatives}}
    • 符号解释:
      • True Positives (TP):模型正确识别出的正例。
      • False Positives (FP):模型错误地将负例识别为正例。
      • False Negatives (FN):模型错误地将正例识别为负例。
      • Precision:模型识别出的正例中有多少是真正的正例。
      • Recall:所有真正的正例中有多少被模型正确识别。

5.3. 对比基线

本文主要将 dLLM-Cache 的性能与以下基线进行比较:

  • 标准推理 (Standard Inference): 即未应用任何加速技术的 dLLM 模型的原始推理过程。这是衡量 dLLM-Cache 提升效果的主要基线。
  • LLaMA3 8B: 一个代表性的 ARM 模型,用于在性能(如准确率)和推理延迟方面与加速后的 dLLM 进行对比,以展示 dLLM 在加速后与 ARM 的竞争力。
  • 更简单的缓存方法 (Simpler Caching Methods): 在消融研究中,通过调整 dLLM-Cache 的超参数(如设置 ρ=0\rho=0 或使用不同的选择策略)来模拟或对比更简单的缓存机制,以验证 dLLM-Cache 各组件的有效性。

实施细节:

  • 固定参数: 所有模型统一采用固定自适应更新比例 ρ=0.25\rho = 0.25
  • 可调参数: 提示刷新间隔 KpK_p 和响应刷新间隔 KrK_r 会根据具体任务和模型进行调整,以在加速和质量之间取得平衡。
  • 硬件: 实验在 NVIDIA RTX 4090H100 80GB SMX GPU 上进行。TPSFLOPs 的测试在单张 RTX 4090 GPU 上完成。

6. 实验结果与分析

6.1. 核心结果分析

实验结果一致表明,dLLM-Cache 在显著提高 dLLM 推理效率的同时,在大多数情况下保持了无损(lossless)的输出质量。

6.1.1. LLaDA 8B 模型上的表现

如 Table 1 所示,dLLM-CacheLLaDA 8B 模型上展示了卓越的性能提升。

  • GPQA 基准测试中,对于 LLaDA Instruct 模型,dLLM-Cache 实现了高达 8.08×8.08 \times 的速度提升,将 FLOPs22.07 T22.07 \text{ T} 降低到 2.73 T2.73 \text{ T},且性能无损。
  • GSM8K 任务中,LLaDA Base 模型在使用 dLLM-Cache 后,Speed(FLOPs) 达到 5.84×5.84 \times,而准确率仅下降 1.74%1.74 \% (69.06% 到 67.32%),或者通过调整参数 (Kp=25, Kr=5) 甚至略有提升 (+1.60%+1.60\%)。
  • MBPPHumanEval 等代码生成任务上,dLLM-Cache 也带来了显著的加速,例如 LLaDA InstructHumanEval 上实现了 6.07×6.07 \timesSpeed(FLOPs) 提升。

以下是原文 Table 1 的结果:

TaskMethodInference EfficiencyPerformance
Speed(TPS)↑FLOPs(T)↓Speed(FLOPs)↑Score↑
Mathematics & Science
GSM8KLLaDA Base + Cache (Kp = 100, Kp = 6) + Cache (Kp = 25, Kγ = 5)7.32 31.43+24.11 23.19+15.871.00× 4.29× +3.29 3.17×+2.1716.12 2.76_13.36 3.21 -12.911.00× 5.84× +4.84 5.02×+4.0269.06 67.32-1.74 70.66+1.60
LLaDA Instruct + Cache (Kp = 50, Kr = 7)6.95 29.75+22.801.00× 4.28×+3.2816.97 2.92 −14.051.00× 5.81×+4.8177.48 78.54+1.06
GPQALLaDA Base + Cache (Kp = 100, Kr = 8)5.12 25.23+20.111.00× 4.93×+3.9322.97 3.20_19.771.00× 7.18× +6.1831.91 32.81+0.90
LLaDA Instruct + Cache (Kp = 50, Kr = 6)5.33 28.01+22.681.00× 5.26×+4.2622.07 2.73-19.341.00× 8.08×+7.0829.01 29.01+0.00
MathLLaDA Base + Cache (Kp = 50, Kr = 8)8.31 33.92+25.611.00× 4.08×+3.0814.11 2.61-11.501.00× 5.41×+4.4130.84 29.84-1.00
LLaDA Instruct + Cache (Kp = 50, Kr = 1)23.65 31.02+7.371.00× 1.31×+0.315.16 3.961.201.00× 1.30×+0.3022.32 22.52+0.20
General Tasks
MMLU-proLLaDA Base + Cache (Kp = 100, Kr = 6)14.08 45.75+31.671.00× 3.25×+2.258.40 2.15-6.251.00× 3.91×+2.9124.26 24.69+0.43
LLaDA Instruct + Cache (Kp = 51, Kr = 3)14.01 39.63+25.621.00× 2.83×+1.838.46 2.625.841.00× 3.23× +2.2336.41 36.08-0.33
MMLULLaDA Base + Cache (Kp = 100, Kr = 6)8.09 33.52+25.431.00× 4.14×+3.1414.56 2.64-11.921.00× 5.52×+4.5263.99 64.26+0.27
LLaDA Instruct + Cache (Kp = 100, Kr = 7)10.12 21.23+11.111.00× 2.10×+1.1011.85 4.50-7.351.00× 2.63×+1.6361.24 62.82+1.58
BBHLLaDA Base + Cache (Kp = 50, Kr = 6)6.41 27.90+21.491.00× 4.35×+3.3518.29 3.09-15.201.00× 5.92× +4.9244.77 45.04+0.27
LLaDA Instruct + Cache (Kp = 100, Kr = 5)6.18 27.55+21.371.00× 4.46×+3.4618.98 3.08-15.901.00× 6.16×+5.1651.49 51.98+0.49
Code 7.87 1.00× 14.91
MBPP+ Cache (Kp = 50, Kr = 4) + Cache (Kp = 25, Kr = 4)30.74+22.87 24.61+16.743.91×+2.91 3.13×+2.132.99-11.92 3.07-11.844.99× +3.99 4.86× +3.8640.80 39.001.80 40.600.20
LLaDA Instruct + Cache (Kp = 100, Kr = 5) LLaDA Base7.55 31.73+24.181.00× 4.20×+3.2015.53 2.80-12.731.00× 5.55×+4.5539.20 39.60+0.40
HumanEval+ Cache (Kp = 100, Kr = 5) + Cache (Kp = 50, Kr = 5)19.98 51.96+31.98 30.68+10.701.00× 2.60×+1.60 1.54×+0.546.03 2.04-3.99 2.07-3.961.00× 2.96×+1.96 2.91×+1.9132.92 32.31-0.61 32.92+0.00
LLaDA Instruct + Cache (Kp = 50, Kr = 8) + Cache (Kp = 25, Kr = 5)10.57 47.29+36.72 44.77+34.201.00× 4.47×+3.47 4.24× +3.2411.10 1.83-9.27 2.05-9.051.00× 6.07× +5.07 5.41×+4.4138.71 37.19-1.52 39.02+0.31

6.1.2. Dream 7B 模型上的表现

如 Table 2 所示,dLLM-CacheDream 7B 模型上也取得了显著的效率提升。

  • GSM8K 基准测试中,Dream Base 模型与 dLLM-Cache 结合,实现了 6.90×6.90 \times 的速度提升,同时准确率保持不变 (76.95%76.95 \%)。
  • GPQA 任务中,Dream Base 模型的 Speed(FLOPs) 达到 7.15×7.15 \times,准确率甚至略有提升 (+0.23%+0.23 \%)。

以下是原文 Table 2 的结果:

TaskConfigurationInference EfficiencyPerformance
Speed(TPS)↑FLOPs(T)↓Speed(FLOPs)↑
TPS↑ Mathematics & Science
GSM8KDream Base + Cache (Kp = 100, Kr = 8)6.36 32.44+26.081.00× 5.10×+4.1019.59 2.84-16.751.00× 6.90×+5.9076.95 76.95+0.00
Dream Instruct + Cache (Kp = 25, Kr = 2)6.39 24.52+18.131.00× 3.84× +2.8419.57 4.24-15.331.00× 4.62×+3.6177.55 76.800.75
GPQADream Base + Cache (Kp = 100, Kr = 8)5.80 30.95+25.151.00× 5.33×+4.3321.66 3.03-18.631.00× 7.15×+6.1533.92 34.15+0.23
Dream Instruct + Cache (Kp = 10, Kr = 8)5.53 21.98+16.451.00× 3.97×+2.9722.63 4.69_ -17.941.00× 4.83×+3.8234.38 33.93-0.45
MathDream Base + Cache (Kp = 100, Kr = 4)9.40 36.89+27.491.00× 3.92× +2.9213.31 2.61-10.701.00× 5.10×+4.1038.68 37.940.74
Dream Instruct + Cache (Kp = 50, Kr = 1)8.85 23.52+14.671.00× 2.66×+1.6614.11 4.66-9.451.00× 3.03×+2.0338.28 37.620.66
General Tasks
MMLU-proDream Base + Cache (Kp = 25, Kr = 2)15.61 35.86+20.251.00× 2.30×+1.307.92 2.89-5.031.00× 2.74×+1.7424.13
Dream Instruct + Cache (Kp = 5, Kr = 1)15.401.00×7.981.00×23.860.27 43.79
Dream Base + Cache (Kp = 100, Kr = 2)23.98+8.58 9.10 31.07+21.971.56×+0.56 1.00×4.77-3.21 13.731.67×+0.67 1.00×43.96+0.17 73.49
Dream Instruct + Cache (Kp = 100, Kr = 8)8.45 38.01+29.563.41×+2.41 1.00× 4.50×+3.503.27-10.46 14.75 2.424.20×+3.20 1.00×73.200.29 73.40
Dream Base + Cache (Kp = 25, Kr = 4)7.24 29.61+22.371.00×-12.33 17.256.10×+5.10 1.00×73.42+0.02 52.25
BBHDream Instruct6.984.09×+3.09 1.00×3.35-13.90 17.905.15×+4.15 1.00×51.660.59 57.07
+ Cache (Kp = 10, Kr = 2)22.31+15.333.20×+2.204.8213.083.71×+2.7157.07+0.00
Code
MBPPDream Base + Cache (Kp = 25, Kr = 8)8.91 35.69+26.781.00× 4.01×+3.0114.06 2.66-11.401.00× 5.29×+4.2954.20 54.20+0.00
Dream Instruct8.461.00×14.651.00×57.00
+ Cache (Kp = 10, Kr = 8)29.77+21.313.52×+2.523.33-11.324.40×+3.4056.800.20
HumanEvalDream Base + Cache (Kp = 5, Kγ = 1)21.431.00×5.681.00×58.53
27.40+5.971.28×+0.284.17-1.511.36×+0.3657.31-1.22
Dream Instruct17.881.00×6.841.00×57.92
+ Cache (Kp = 50, Kr = 1)28.03+10.151.57×+0.573.94_2.901.74×+0.7456.091.83

6.1.3. 与自回归模型 (ARMs) 的对比

如 Table 3 所示,dLLM-Cache 使得 dLLM 的推理延迟接近 ARM,同时保持了更高的性能。

  • GSM8K 任务上,LLaMA3 8B (一个 ARM 模型) 提供了 47.73 TPS47.73 \text{ TPS} 的速度和 49.05%49.05 \% 的准确率。

  • 未加速的 LLaDA (使用 256 解码步骤) 仅有 7.37 TPS7.37 \text{ TPS},但准确率高达 69.06%69.06 \%

  • 应用 dLLM-Cache 后,LLaDATPS 提升到 25.02 TPS25.02 \text{ TPS} (3.39×3.39 \times 加速),准确率仍保持在 67.32%67.32 \%,远高于 LLaMA3 8B,且 TPS 接近 LLaMA3 8B 的一半。

  • 即使将 LLaDA 的解码步数减少到 128 (牺牲一部分准确率),dLLM-Cache 也能将其 TPS 提升到 49.15 TPS49.15 \text{ TPS} (3.33×3.33 \times 加速),这已经超越了 LLaMA3 8BTPS,同时准确率仍有 62.32%62.32 \%,显著优于 LLaMA3 8B

    这表明 dLLM-Cache 显著弥补了 dLLM 在推理速度上的劣势,使其在速度和质量的权衡上更具竞争力。

以下是原文 Table 3 的结果:

MethodTPS↑Speed(TPS) ↑Accuracy↑GPU Memory (GB)↓
LLaMA3 8B [Dubey et al., 2024]47.731.00×49.0516.06
LLaDA*7.37-40.361.00×69.06+20.0116.94
+ Cache* (Kp = 100, Kr = 6) LLaDA†25.0222.713.39×67.32+18.2717.93
14.77-32.961.00×64.14+15.0916.94
+ Cache† (Kp = 50, Kr = 8)49.15+1.423.33×62.32+13.2717.93

注:上标 * 表示 LLaDA 使用 256 解码步骤, 表示 LLaDA 使用 128 解码步骤。

6.1.4. 长提示场景下的有效性

dLLM-Cache 在处理长输入提示(如文档问答任务)时尤其有效。

  • LongBench-HotpotQA 任务上,对 LLaDA 8B Base 模型应用 dLLM-Cache,不仅实现了 9.1×9.1 \times 的速度提升,而且 F1 分数从 34.56 提升到 36.10,性能反而有所改善。
  • 这突出表明 长间隔提示缓存 机制在处理大量静态提示词元时的优势,能够最大限度地利用缓存,大幅减少冗余计算。

6.2. 消融实验/参数分析

6.2.1. 缓存刷新间隔 KpK_pKrK_r 的影响

下图(原文 Figure 4)展示了 LLaDA 8B Instruct 模型上缓存刷新间隔对效率和准确率的影响:

Figure 4: Effect of cache refresh intervals using LLaDA 8B Instruct.a) Varying `K _ { p }` EY with \(K _ { r } = 1\) \(\\rho = 0\) (b) Varying `K _ { r }` under two settings: baseline with \(K _ { p } = 1\)… 该图像是图表,展示了LLaDA 8B Instruct模型中缓存刷新间隔对性能的影响。图(a)中,随着缓存提示间隔KpK_p变化,准确率和TFLOPs的关系;图(b)中,分别比较了基线组(Kp=1,ρ=0)(K_p=1, \rho=0)和本方法(Kp=50,ρ=0.25)(K_p=50, \rho=0.25)下不同响应刷新间隔KrK_r对准确率和TFLOPs的影响。 Figure 4: Effect of cache refresh intervals using LLaDA 8B Instruct.a) Varying KpK_p EY with Kr=1K_r = 1 ρ=0\rho = 0 (b) Varying KrK_r under two settings: baseline with Kp=1K_p = 1 ρ=0\rho = 0 in gary and our setup Kp=K_p = 50, ρ=0.25\rho = 0.25 in Table 1.

  • 图 4(a) (提示间隔 KpK_p 的影响): 增加提示间隔 KpK_p(即更不频繁地刷新提示缓存)可以显著降低 FLOPs,且在很大范围内不会损害准确率。这验证了静态提示特征在去噪步骤中的稳定性,以及长间隔提示缓存的有效性。
  • 图 4(b) (响应间隔 KrK_r 和自适应更新的影响):
    • 基线(灰线,Kp=1,ρ=0K_p = 1, \rho = 0): 不进行提示缓存且不进行自适应更新。在这种设置下,当 KrK_r 增加(即响应缓存刷新变少)时,准确率急剧下降,表明简单地重用响应特征会导致严重的性能损失。
    • 本文设置(橙线和蓝线,Kp=50,ρ=0.25K_p = 50, \rho = 0.25): 结合长提示间隔和自适应响应更新。在这种设置下,即使 KrK_r 增加,模型也能保持较高的准确率,同时显著降低 FLOPs
  • 结论: 这项消融实验证实了 dLLM-Cache 中结合长间隔提示缓存和短间隔、自适应响应更新策略的有效性。

6.2.2. 更新比例 ρ\rho 和选择策略的影响

下图(原文 Figure 5)展示了在 GSM8K 任务上,不同令牌选择策略在不同自适应更新比例 ρ\rho 下对准确率和计算成本的影响:

Figure 5: Effect of token selection strategy on GSM8K using LLaDA 8B Instruct model under varying update ratios \(\\rho\) . Figure 5: Effect of token selection strategy on GSM8K using LLaDA 8B Instruct model under varying update ratios ρ\rho.

  • 对比策略: 实验比较了三种令牌选择策略:V-verify(基于值向量相似度)、K-verify(基于键向量相似度)和 random selection(随机选择)。
  • 结果:
    • V-verifyK-verify 这两种基于相似度的策略,在各种 ρ\rho 值下都显著优于 random selection,这验证了利用特征动态性进行选择性更新的重要性。
    • V-verify 尤其表现出最佳的性能,在大约 ρ=0.25\rho = 0.25 的更新比例下,能够实现高准确率(接近无损性能),同时计算成本远低于全量重新计算(ρ=1\rho = 1)。例如,在 ρ=0.9\rho = 0.9 时达到了 77.48%77.48 \% 的无损准确率。
  • 结论: 适度的、有针对性的更新(例如 ρ0.25\rho \approx 0.25)能够在效率和输出质量之间取得良好的平衡。

6.2.3. 相似性度量 Impact of Similarity Metric

V-verify 机制中,本文比较了余弦相似度(cosine similarity)和 L2 距离作为相似性度量。

  • GSM8K 任务上,使用 LLaDA 8B Instruct 模型,余弦相似度实现了 78.54%78.54 \% 的准确率。
  • L2 距离仅达到 55.95%55.95 \% 的准确率。
  • 结论: 余弦相似度能够更好地捕捉特征向量的方向性变化,从而更有效地识别词元语义上的改变,因此被选作 dLLM-Cache 的默认相似性度量。

6.2.4. V-verify 的成本和固定更新开销

下图(原文 Figure 6)展示了 TPS 随更新比例 ρ\rho 变化的趋势:

Figure 6: TPS versus \(\\rho\) A notable decrease in TPS persists at minimal \(\\rho\) values, highlighting fixed operational overheads associated with initiating any selective update. Figure 6: TPS versus ρ\rho A notable decrease in TPS persists at minimal ρ\rho values, highlighting fixed operational overheads associated with initiating any selective update.

  • V-verify 的计算开销: V-verify 机制本身通过计算轻量级的 VV 向量相似度来识别动态词元,其计算成本很低。
  • 固定操作开销: 即使当更新比例 ρ\rho 接近于零时,TPS 仍然存在显著下降。这表明,启动任何选择性重新计算(即 ρ>0\rho > 0)都会带来不可忽略的系统级固定开销,例如 GPU 内核管理和数据传输。这些开销不与更新的词元数量严格成比例。
  • 结论: 在极低的 ρ\rho 值下,这些固定开销会占据主导地位,限制了进一步的运行时节省。因此,需要选择一个最优的 ρ\rho 值来平衡这些固定成本和减少动态计算带来的收益。Figure 5 的结果表明 ρ0.25\rho \approx 0.25 提供了一个在效率和保真度之间有效的权衡点。

6.3. 存储开销

dLLM-Cache 缓存了每层四个类型的中间特征:KKVVAttnOutFFNOut

  • 理论存储开销: 总缓存大小与词元数量 TT、嵌入维度 dd 和层数 LL 成正比,即 O(T×d×4×L)O(T \times d \times 4 \times L)。由于每层只缓存一个版本,因此总内存占用保持稳定。
  • 实际内存占用: 如 Table 3 所示,在 GSM8K 任务上,对于 LLaDA 8B Base 模型:
    • 无缓存时的峰值 GPU 内存占用为 16.94 GB16.94 \text{ GB}
    • 应用 dLLM-Cache 后的峰值 GPU 内存占用为 17.93 GB17.93 \text{ GB}
    • 作为对比,LLaMA3 8B(带有 KV 缓存的自回归模型)的峰值 GPU 内存占用为 16.06 GB16.06 \text{ GB}
  • 结论: dLLM-Cache 引入了适度的额外内存开销(约 1 GB1 \text{ GB}),但相对于标准 ARM 模型(如 LLaMA3 8B)而言,在获得显著加速的同时,实现了更高的准确率。

7. 总结与思考

7.1. 结论总结

本文提出了 dLLM-Cache,一个免训练且模型无关的缓存框架,旨在加速扩散大语言模型(dLLMs)的推理过程。该方法基于对 dLLMs 推理过程中“静态提示”和“部分动态响应”两种计算冗余的深入观察。dLLM-Cache 结合了长间隔提示缓存和由 V-verify 机制引导的自适应响应部分更新策略。

实验结果表明,dLLM-CacheLLaDA 8BDream 7B 等主流 dLLMs 上实现了显著的推理速度提升(最高达 9.1×9.1 \times),同时保持了无损或接近无损的生成质量。值得注意的是,通过 dLLM-Cache 的加速,dLLM 的推理延迟在许多情况下能够与自回归模型(ARMs)相媲美,甚至在某些任务上以更低的延迟实现更高的性能。此外,该方法在长提示场景下表现出尤其显著的优势。

7.2. 局限性与未来工作

作者指出了当前研究的一个主要局限性:

  • 模型规模限制: 目前公开可用的 dLLMs 主要限于 8B8 \text{B} 参数规模。这意味着 dLLM-Cache 尚未在更大规模的模型(如 33B33 \text{B}70B70 \text{B} 参数)上进行评估。
  • 未来工作方向: 作者相信 dLLM-Cache 在更大规模的模型上将带来更显著的加速收益。因此,在更大规模的 dLLMs 上进行评估将是未来的重要研究方向。

7.3. 个人启发与批判

7.3.1. 个人启发

  1. 从领域特性出发的优化: dLLM-Cache 的成功在于其深刻理解了 dLLMs 在语言任务中的独有特性(静态提示与动态响应),而非简单套用通用加速技术。这启发我们,在设计高效系统时,领域知识和对数据流、计算模式的细致分析是至关重要的。
  2. “稀疏性”与“自适应”的威力: 响应特征演变的稀疏性是关键洞察。V-verify 机制利用轻量级相似度度量来识别最需要更新的词元,体现了“按需计算”和“自适应”的强大优化潜力。这种思路可以推广到其他具有局部动态性但整体稳定的迭代生成模型中。
  3. 工程与学术的结合: 该工作不仅提出了理论上的优化方法,更提供了实际的加速效果,并详细分析了内存开销和固定操作成本,展现了扎实的工程实践与严谨的学术研究相结合的典范。

7.3.2. 潜在的问题、未经验证的假设或可以改进的地方

  1. V-verify 的泛化能力: 尽管 VV 向量相似度与下游特征变化存在强相关性(Figure 2),但这种相关性是否在所有 dLLMs、所有任务类型以及不同的去噪步长设置下都保持一致?是否存在某些边缘情况,VV 向量变化不显著但 AttnOutFFNOut 却发生了重要变化,从而导致缓存错误并影响生成质量?虽然实验结果表明无损,但这仍是一个潜在的理论弱点。
  2. 固定操作开销的优化: Figure 6 指出,即使 ρ\rho 极小,仍然存在显著的固定操作开销。这可能与 GPU 内核启动、数据传输、缓存管理机制本身的实现效率有关。未来的工作可以探索系统级的优化,例如通过定制化的 CUDA 内核或更高效的内存管理,进一步降低这部分固定开销,从而在更低的 ρ\rho 值下实现更高的 TPS
  3. 动态调整 ρ\rho 论文中 ρ\rho 值是固定的 (0.25)。然而,在不同的生成阶段或根据生成内容的复杂性,可能需要不同比例的更新。例如,在生成早期,序列变化可能更剧烈,需要更高的 ρ\rho;而在生成后期,序列趋于稳定,可以采用更低的 ρ\rho。未来的工作可以探索一种动态调整 ρ\rho 的策略,使其根据模型状态或预测的不确定性进行自适应调整。
  4. 对模型性能的深层影响: 尽管论文报告了“无损”的输出质量,但这通常是基于如准确率、F1 分数等宏观指标。在某些对细微语义变化敏感的任务中,部分词元的缓存重用是否会引入“次优”的局部特征,从而影响生成文本的流畅性、创造性或事实准确性,但这些影响在现有评估指标中难以体现?进行更细粒度的文本质量评估可能会揭示更多信息。
  5. 缓存一致性与并行计算: 论文主要关注单序列推理。在批量推理(batch inference)或并行计算场景下,dLLM-Cache 如何维护缓存的一致性,并确保高效的并行数据访问和更新,将是一个值得研究的问题。

相似论文推荐

基于向量语义检索推荐的相关论文。

暂时没有找到相似论文。