Focus-dLLM：通过置信度引导的上下文聚焦加速长上下文扩散大语言模型推理 Focus-dLLM: Accelerating Long-Context Diffusion LLM Inference via Confidence-Guided Context Focusing
利用相邻步骤置信度相关性预测解码位置，实现无训练的29倍长上下文推理加速
前置知识
扩散大语言模型（dLLM）
扩散大语言模型是一种非自回归的文本生成范式，与传统从左到右逐个生成token的方式不同，dLLM通过迭代去噪的方式同时更新多个位置。给定一个固定长度的序列，初始时将待生成位置填充为[MASK]标记，然后在每一步去噪中，模型为每个被遮蔽位置预测一个条件概率分布，并根据置信度分数选择性地揭开遮蔽。这种方式可以并行更新多个token，结合双向注意力机制，能够实现更高的解码吞吐量。
理解dLLM的非自回归解码机制是理解本文的前提，因为Focus-dLLM正是针对这种特殊解码模式设计的加速方案
注意力汇聚（Attention Sink）
注意力汇聚是指在注意力计算中，某些特定token（通常是序列开头的token）会吸引不成比例的注意力分数，形成明显的垂直带状模式。这些token虽然可能语义上不重要，但对模型的生成质量起着关键作用。在dLLM中，注意力汇聚会随着去噪步骤动态变化，但研究发现它们在不同层之间具有位置一致性——即同一去噪步骤中，不同层的注意力汇聚点位置基本相同。
本文的核心创新之一就是利用注意力汇聚的跨层一致性来避免重复计算，从而大幅减少推理开销
近似KV缓存
由于dLLM使用双向注意力，传统自回归模型的KV缓存机制无法直接应用。近似KV缓存是一种妥协方案：在每个去噪步骤中，只对选定的token子集（如新解码的token）重新计算KV状态，而对其他token复用缓存的KV状态。形式化表示为：当token索引i属于刷新集合U(t)时，K_i^t = f_K(x^(t))_i（重新计算）；否则K_i^t = \tilde{K}_i（复用缓存）。
近似KV缓存是dLLM加速的基础机制，本文在此基础上进一步引入稀疏注意力来减少计算量
半自回归重遮蔽策略
这是dLLM推理中的一种混合解码策略：在每个块（block）内部采用非自回归方式同时解码多个token，但在块与块之间采用自回归方式从左到右依次处理。这种策略平衡了并行性和解码质量，是当前主流dLLM（如UltraLLaDA、Dream）采用的标准解码方式。
Focus-dLLM遵循这种策略，在块入口处进行全缓存刷新，其他步骤才应用稀疏注意力优化
研究动机
扩散大语言模型虽然在长上下文处理上展现出强大能力，但其推理效率面临严峻挑战。核心问题在于双向全注意力的计算开销：对于长度为L的序列，每一步去噪都需要计算L×L的注意力矩阵，时间复杂度为O(L²)。现有加速方法主要分为两类：近似KV缓存和稀疏注意力。KV缓存方法虽然能减少重复计算，但注意力仍然需要在整个缓存上下文上进行。稀疏注意力方法更为直接，但面临一个根本性困难：dLLM的解码位置是动态确定的——在去噪过程中，哪些位置会被揭幕（unmasked）在事先是未知的。这导致无法像自回归模型那样提前估计token的重要性。现有方法如Sparse-dLLM采用粗糙的块级指标进行缓存淘汰，SparseD在早期步骤仍依赖密集注意力，这些方法都没有充分利用dLLM的注意力模式特性，导致加速效果有限。
本文的目标是本文的具体目标是设计一个无需训练的注意力稀疏化框架，能够在保持生成质量的前提下显著加速长上下文dLLM推理。更具体地说，作者希望回答一个关键问题：能否准确预测未遮蔽token的位置，并且只保留必要的计算来实现更有效的长上下文推理加速？目标是在32K上下文长度下实现数量级的吞吐量提升，同时在LongBench等基准测试上保持或超越基线性能。
与已有工作不同的是，本文的独特切入角度基于两个关键观察发现：第一，相邻去噪步骤的token置信度存在强正相关性——当前步骤t中被解码的位置，在前一步t-1就已经展现出很高的置信度，平均召回率高达96.1%。这意味着可以通过前一步的置信度准确预测当前步的解码位置。第二，dLLM的注意力汇聚点表现出显著的跨层一致性——在不同深度的层中，注意力汇聚的位置基本相同。基于这两个发现，作者提出了一种全新的稀疏注意力方案：用历史置信度指导查询选择，用跨层一致性避免重复识别注意力汇聚点，从而实现精确且高效的注意力剪枝。
核心方法
Focus-dLLM的整体思路可以概括为预测-聚焦-剪枝三步走。直觉上，既然我们能通过前一步的置信度知道当前步哪些位置会被解码，那么只需要对这些位置附近计算注意力就够了；同时，既然注意力汇聚点在各层位置一致，那么在一个中间层识别出汇聚点后，深层直接复用即可，不需要每层重新计算。技术路线上，首先利用Past Confidence-Guided Indicator（PCGI）从前一步置信度预测当前步的解码候选位置，然后通过窗口扩展形成活跃查询集合；接着在指定深度的稠密层识别注意力汇聚点，利用跨层一致性复用到稀疏层；最后通过块级剪枝选择最相关的prompt块进行注意力计算，实现动态稀疏注意力。
核心创新点有两个：第一是过去置信度引导指示器（PCGI），它利用扩散解码中相邻步骤置信度高度相关的特性，从前一步的置信度分数预测当前步的解码位置。具体做法是从所有仍处于[MASK]状态的位置中，按置信度c_{t-1}排序选取top-k个作为候选集合I_focus，其中k = ⌊ρ·n(t)⌋，ρ是预定义的扩展因子，n(t)是当前步要解码的token数。第二是感知汇聚点的稀疏注意力，它在指定深度的稠密层通过聚合查询表示和token重要性分数识别注意力汇聚点，然后将这些汇聚点位置复用到更深层，避免重复计算。与已有方法的本质区别在于：Sparse-dLLM使用粗糙的块级指标进行淘汰，SparseD仍需在早期步骤进行密集计算，而Focus-dLLM通过精确的置信度预测和跨层汇聚点复用，实现了更精细的注意力剪枝。
方法步骤详情
方法的具体步骤如下：首先，在每个去噪步骤t，判断是否为块入口步骤。如果是块入口步骤，则进行全缓存刷新，对整个序列计算注意力（use_sparse = False）。如果不是块入口步骤，则进入稀疏注意力流程：（1）置信度引导的查询预测：从上一步t-1的置信度分数中，选取top-k个最高的[MASK]位置作为候选集合I_focus；（2）窗口扩展：将候选位置扩展为大小为w的局部窗口，形成活跃查询集合I_active，这是因为token表示强烈依赖于附近的语义上下文；（3）层间前向传播：对于前l_dense层执行全注意力，第l_dense层的注意力分布用于识别注意力汇聚点；（4）块级token剪枝：将prompt token划分为连续块，每个块计算代表性key（块内key的均值），然后计算候选查询与每个块的相关性分数R_b，选择top-C个最相关的块；（5）最终注意力计算：汇聚点集合I_sink与选中的prompt块取并集形成I_p，只对这些位置的key-value对计算稀疏注意力。
技术新颖性
技术新颖性体现在三个方面：第一，首次系统分析并利用了dLLM相邻步骤置信度的时间一致性，发现96.1%的高召回率，这为预测解码位置提供了可靠依据，而此前的方法（如Sparse-dLLM的动态缓存淘汰）使用的是粗糙的块级指标。第二，发现了注意力汇聚点的跨层一致性并加以利用——在指定深度的稠密层识别汇聚点后直接复用到所有稀疏层，这避免了Sparse-dLLM等方法中每层重新识别的开销。第三，设计了完整的置信度预测→窗口扩展→汇聚点识别→块级剪枝流水线，各组件协同工作：PCGI提供的精确查询使得块级相关性估计更准确，汇聚点保留保证了生成质量，块级剪枝提高了GPU利用率。这种组合设计比单独使用任何一种技术都更有效。
实验结果
实验结果全面验证了Focus-dLLM的有效性。在LongBench基准测试上，Focus-dLLM在UltraLLaDA模型上取得了45.14的平均分，不仅超过了Fast-dLLM（44.74）、Sparse-dLLM（44.86）和SparseD（44.70）等加速方法，甚至略高于原始的Vanilla基线（44.90）。在Dream-7B-Instruct模型上，Focus-dLLM达到42.82分，虽然略低于SparseD（43.59），但显著优于Fast-dLLM（42.75）和Sparse-dLLM（42.78）。效率方面，在32K上下文长度下，Focus-dLLM对UltraLLaDA实现了29.6倍的吞吐量提升（Vanilla为1.06 tokens/s，Focus-dLLM达到19.30 tokens/s），对Dream实现了27.6倍提升。速度提升随上下文长度增加而扩大：从8K的9.4倍增长到32K的29.6倍，这说明在更长序列上，冗余注意力计算的比例更大，Focus-dLLM的剪枝效果更显著。特别值得注意的是，在Niah（大海捞针）测试中，Focus-dLLM在最深层甚至超越了Vanilla基线，表明其稀疏注意力策略能有效保留关键信息。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| LongBench综合评分（UltraLLaDA） | Ave. Score | 45.14 | 44.90 (Vanilla), 44.74 (Fast-dLLM), 44.86 (Sparse-dLLM), 44.70 (SparseD) | 超越Vanilla 0.24分，超越Fast-dLLM 0.40分 |
| LongBench综合评分（Dream-7B） | Ave. Score | 42.82 | 43.03 (Vanilla), 42.75 (Fast-dLLM), 42.78 (Sparse-dLLM), 43.59 (SparseD) | 接近Vanilla性能，超越Fast-dLLM 0.07分 |
| 32K上下文吞吐量（UltraLLaDA） | tokens/s | 19.30 | 1.06 (Vanilla), 11.03 (Fast-dLLM) | 29.6x加速（对比Vanilla），1.75x加速（对比Fast-dLLM） |
| 32K上下文吞吐量（Dream-7B） | tokens/s | 15.20 | 0.55 (Vanilla), 9.57 (Fast-dLLM) | 27.6x加速（对比Vanilla），1.59x加速（对比Fast-dLLM） |
局限与改进
作者在论文末尾承认了两个主要局限性：第一，Focus-dLLM目前仅在文本任务上进行了验证，其向多模态推理的扩展尚未探索，这对于需要处理图像、视频等输入的实际应用场景是一个重要的缺失。第二，当前的超参数（如稀疏比α、扩展因子ρ、窗口大小w、稠密层数l_dense等）都是手动配置的，可能无法在所有专业领域达到最优性能。从我的观察来看，还有一些潜在问题值得关注：首先，置信度预测依赖于相邻步骤的强相关性，但在解码早期步骤（如从全遮蔽状态开始）这种相关性可能较弱；其次，块级剪枝的粒度是固定的（默认64个token），对于需要精细粒度关注的任务可能不够灵活；最后，虽然作者展示了在8B参数模型上的结果，但对更大规模模型（如70B+）的适用性尚不明确。
独立分析的弱点
从独立分析的角度，Focus-dLLM存在以下几个弱点：第一，置信度预测的准确性可能在解码初期受到影响——当大部分位置都是[MASK]时，置信度的区分度可能不够强，此时PCGI的预测召回率可能低于报告的96.1%。改进方向是设计一个渐进式的预测策略，在早期步骤采用更保守的窗口扩展或直接进行密集计算。第二，注意力汇聚点的识别依赖于指定深度的稠密层，但这个深度（l_dense）需要手动设定，论文中通过消融实验发现不同模型的最佳值不同（UltraLLaDA为6，Dream需要保留最后4层）。一个更优雅的方案是设计自适应的汇聚点检测机制，根据注意力分布的熵或集中度动态决定。第三，块级剪枝使用简单的均值池化作为块代表性key，这可能丢失块内的细粒度信息，特别是当块内包含语义差异很大的token时。可以考虑使用可学习的聚合函数或基于注意力分数的加权平均。
未来方向
作者提出的发展方向包括多模态扩展和自适应参数调整机制。基于本文的成果，还可以延伸出以下研究方向：第一，将PCGI的置信度预测思想应用到其他dLLM加速方法中，如改进Fast-dLLM的KV缓存更新策略——只对预测会被解码的位置进行缓存刷新。第二，探索注意力汇聚点的语义含义，理解为什么某些位置会成为汇聚点，这可能揭示dLLM的内部工作机制。第三，将Focus-dLLM与量化技术结合，因为稀疏注意力减少了需要计算的token数量，这为低精度计算提供了更好的条件。第四，研究Focus-dLLM在不同解码策略下的表现，如改变去噪调度器（scheduler）的类型，或者应用于连续空间扩散模型。第五，探索自动化的超参数搜索策略，根据输入序列的长度和特性动态调整稀疏比和窗口大小。
复现评估
论文的可复现性较好。代码已在GitHub公开（https://github.com/Longxmas/Focus-dLLM），这大大降低了复现门槛。实验使用OpenCompass评估平台，这是广泛使用的开源评测工具。所有实验在NVIDIA H200 GPU上进行，虽然这是高端硬件，但方法本身是训练无关的（training-free），不需要额外的训练过程。数据集方面，LongBench是公开的标准基准，GSM8K也是常用数据集。论文提供了详细的超参数设置（α=0.5, ρ=4, w=8, l_dense=6, N_sinks=0.01×M, 块大小=64），以及与各基线方法的公平对比配置。唯一的潜在障碍是实现细节——虽然论文提到了使用Triton和FlashAttention进行GPU内核优化，但具体的CUDA内核实现可能需要一定的系统优化经验。总体而言，对于有dLLM推理经验的研究团队，复现难度为中等。
论文图表
左图展示了相邻步骤置信度分数的相关性散点图，显示强烈的正相关关系（相关系数0.856）。右图展示了使用前一步top-4最高置信度的重遮蔽token预测当前步解码token的召回率，平均召回率达到96.1%，并随去噪步骤保持稳定。
这张图是整个方法的理论基础，证明了置信度的时间一致性是可靠的经验规律，为PCGI设计提供了直接依据