EntropyCache: Decoded Token Entropy Guided KV Caching for Diffusion Language Models
作者: Minsoo Cheong, Donghyun Son, Woosang Lim, Sungjoo Yoo
分类: cs.CL
发布日期: 2026-03-19
🔗 代码/项目: GITHUB
💡 一句话要点
EntropyCache:利用解码Token熵引导扩散语言模型的KV缓存,实现高效推理。
🎯 匹配领域: 支柱九:具身大模型 (Embodied Foundation Models)
关键词: 扩散语言模型 KV缓存 熵 推理加速 近似计算 无训练 注意力机制
📋 核心要点
- 扩散语言模型由于双向注意力机制,无法直接使用KV缓存,导致推理效率低下,现有近似KV缓存方法存在决策开销随上下文增长的问题。
- EntropyCache利用解码token熵作为缓存失效的代理信号,并根据熵值决定是否重新计算最近解码的token,无需额外训练。
- 实验表明,EntropyCache在保证精度的前提下,显著提升了推理速度,决策开销极低,适用于多种扩散语言模型。
📝 摘要(中文)
基于扩散的大型语言模型(dLLMs)依赖于双向注意力机制,这阻碍了无损KV缓存,并需要在每个去噪步骤中进行完整的正向传播。现有的近似KV缓存方法通过选择性地更新缓存状态来降低这种成本,但它们的决策开销会随着上下文长度或模型深度而扩展。我们提出了EntropyCache,一种无需训练的KV缓存方法,它使用新解码的token分布的最大熵作为恒定成本信号,来决定何时重新计算。我们的设计基于两个经验观察:(1)解码的token熵与KV缓存漂移相关,为缓存陈旧度提供了一个廉价的代理,(2)解码的token的特征波动在unmask后持续多个步骤,促使重新计算最近解码的k个token。跳过或重新计算的决策每步仅需O(V)计算量,与上下文长度和模型规模无关。在LLaDA-8B-Instruct和Dream-7B-Instruct上的实验表明,EntropyCache在标准基准测试中实现了15.2倍-26.4倍的加速,在思维链基准测试中实现了22.4倍-24.1倍的加速,同时保持了具有竞争力的准确性,并且决策开销仅占推理时间的0.5%。代码可在https://github.com/mscheong01/EntropyCache获取。
🔬 方法详解
问题定义:扩散语言模型(dLLMs)推理过程中,由于双向注意力机制,无法像自回归模型那样直接利用KV缓存加速推理。现有的近似KV缓存方法,例如选择性更新缓存,虽然能降低计算成本,但其决策过程的计算开销会随着上下文长度和模型深度增加,成为瓶颈。
核心思路:EntropyCache的核心思想是利用解码token的熵值作为KV缓存失效程度的代理指标。高熵值通常意味着模型对当前解码结果的不确定性较高,缓存中的信息可能已经过时,需要重新计算。通过监控解码token的熵值,可以动态地决定是否需要重新计算KV缓存,从而在精度和效率之间取得平衡。
技术框架:EntropyCache的整体流程如下:在每个去噪步骤中,首先解码得到新的token及其概率分布。然后,计算该token概率分布的最大熵。如果熵值超过预设的阈值,则重新计算最近解码的k个token的KV缓存;否则,跳过重新计算步骤,直接使用缓存中的KV值。这个过程在每个去噪步骤中重复进行,直到完成整个序列的生成。
关键创新:EntropyCache的关键创新在于使用解码token的熵值作为KV缓存更新的触发信号。与现有方法相比,EntropyCache的决策过程计算复杂度为O(V),与上下文长度和模型规模无关,因此具有更好的可扩展性。此外,EntropyCache无需额外的训练,可以直接应用于现有的扩散语言模型。
关键设计:EntropyCache的关键设计包括:1) 使用最大熵作为缓存失效的指标;2) 重新计算最近解码的k个token的KV缓存,其中k是一个超参数,需要根据具体模型和任务进行调整;3) 熵阈值的设置,需要根据经验进行调整,以平衡精度和效率。
🖼️ 关键图片
📊 实验亮点
EntropyCache在LLaDA-8B-Instruct和Dream-7B-Instruct模型上进行了实验,结果表明,在标准基准测试中实现了15.2倍-26.4倍的加速,在思维链基准测试中实现了22.4倍-24.1倍的加速,同时保持了具有竞争力的准确性。决策开销仅占推理时间的0.5%,表明EntropyCache具有很高的效率。
🎯 应用场景
EntropyCache可应用于各种基于扩散的大型语言模型,尤其是在资源受限的场景下,例如移动设备或边缘计算平台。它可以显著提升dLLMs的推理速度,使其能够更快地生成文本、图像或其他模态的内容,从而促进dLLMs在自然语言处理、计算机视觉等领域的更广泛应用。
📄 摘要(原文)
Diffusion-based large language models (dLLMs) rely on bidirectional attention, which prevents lossless KV caching and requires a full forward pass at every denoising step. Existing approximate KV caching methods reduce this cost by selectively updating cached states, but their decision overhead scales with context length or model depth. We propose EntropyCache, a training-free KV caching method that uses the maximum entropy of newly decoded token distributions as a constant-cost signal for deciding when to recompute. Our design is grounded in two empirical observations: (1) decoded token entropy correlates with KV cache drift, providing a cheap proxy for cache staleness, and (2) feature volatility of decoded tokens persists for multiple steps after unmasking, motivating recomputation of the $k$ most recently decoded tokens. The skip-or-recompute decision requires only $O(V)$ computation per step, independent of context length and model scale. Experiments on LLaDA-8B-Instruct and Dream-7B-Instruct show that EntropyCache achieves $15.2\times$-$26.4\times$ speedup on standard benchmarks and $22.4\times$-$24.1\times$ on chain-of-thought benchmarks, with competitive accuracy and decision overhead accounting for only $0.5\%$ of inference time. Code is available at https://github.com/mscheong01/EntropyCache.