Accelerating Diffusion Language Models via Structured Suffix Modeling¶

会议: NeurIPS 2026
arXiv: 2608.23167
OpenReview: https://openreview.net/forum?id=SkQxO613nL
代码: https://github.com/zifengcheng/SSM
领域: LLM效率
关键词: 扩散语言模型、并行解码、后缀建模、训练免除加速、软嵌入混合

一句话总结¶

针对扩散语言模型块级去噪时后缀注意力计算冗余与表征静态重置问题，本文提出免训练的结构化后缀建模（SSM），将后缀划分为局部、中部与尾部三区域差异化分配保留预算，并引入软后缀嵌入混合上步去噪信号，实现最高 72.81× 的端到端推理加速且维持或提升生成精度。

研究背景与动机¶

扩散语言模型（Diffusion Language Models, DLMs）依托全序列或块级的双向注意力机制，能够在单一去噪步内并行预测多个 token，为打破自回归生成的延迟瓶颈提供了全新范式。然而，在以块为单位的离散扩散生成过程中，模型在去噪当前块的同时，必须计算所有被掩码覆盖的后缀 token 的隐藏状态与注意力分数。随着序列生成长度的增加，这部分仅承载未来潜在上下文信号的后缀计算开销急剧放大，迅速成为制约扩散语言模型推理吞吐与端到端延迟的核心瓶颈。

现有的后缀丢弃（Suffix Dropout）方法大多将这一问题简化为基于局部窗口的修剪任务。例如 DPad 通过设定局部窗口并在窗口内施加距离感知丢弃，Streaming-dLLM 则简单保留相邻后缀 token 与句尾 token。这种局部化启发式修剪存在两项根本性缺陷：其一，它们将后缀 token 视作同质化的冗余容器，忽视了不同后缀区域在全局去噪过程中承担的异构功能角色，未能精细化建模结构与边界骨架；其二，后缀 token 在每一去噪时间步均被粗暴地重新初始化为完全相同的纯静态 [MASK] 嵌入，彻底遗弃了前序迭代中逐步积累并演化的宝贵预测分布与语义演化信号。

本文的切入角度是实证分析后缀注意力在不同空间跨度上的分布演化：注意力得分随后缀距离增加呈现先骤降、后平稳、再在尾部显著回升的“U 型”结构。核心 idea：将后缀解构为局部细粒度上下文、中部轻量结构骨架与尾端边界锚点三个异构功能区并差异化分配保留预算，同时利用软后缀嵌入将前序解码概率融合至当前步输入，使极精简的后缀 token 能够自适应传递渐进演化的去噪信号。

方法详解¶

整体框架¶

SSM 是一种即插即用、无需重新训练的离散扩散语言模型推理加速框架。其整体处理流程分为三个阶段：首先根据注意力衰减与回升特性，将当前解码块之后的全量被掩码后缀块解构为局部区域、中部区域与尾部区域；接着按各区域承担的功能角色执行非对称的 token 保留策略，并保留各 token 在原始序列中的绝对/旋转位置编码作为定位锚点；随后利用软后缀嵌入混合机制，将上一时间步在保留位置预测出的 Top-\(k\) 软分布混合入静态掩码向量；最后将精简后的稀疏输入送入扩散主干进行并行块去噪，并结合块级 EOS 早期终止机制阻断无意义填充。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入序列：提示词 P + 当前块 B_t + 全量后缀 S_t"] --> B["后缀异构区域划分<br/>局部区、中部区、尾部区"]
    B --> C["非对称结构化 Token 筛选<br/>全量邻近 + 块起始锚点 + 终末边界"]
    C --> D["软后缀嵌入混合<br/>上一时间步 Top-k 概率混合至 MASK 嵌入"]
    D --> E["稀疏输入张量构建<br/>拼接 Prompt、当前块与带原始位置编码的保留后缀"]
    E --> F["扩散语言模型并行块去噪与解码"]
    F -->|检测到 EOS| G["提前终止并截断后续计算"]
    F -->|未结束| H["更新预测分布用于下一时间步软嵌入"]

关键设计¶

1. 后缀异构区域划分：基于注意力动力学的功能区解耦 针对以往方法无视后缀空间异质性的缺陷，本文基于注意力得分的统计特性将后缀切分成三段不同职能的区域。在时间步 \(t\)，全量后缀块集合 \(S_t\) 被严格划分为局部区域 \(S_t^{\mathrm{local}}\)、中部区域 \(S_t^{\mathrm{middle}}\) 与尾部区域 \(S_t^{\mathrm{tail}}\)：

\[S_{t}=S_{t}^{\mathrm{local}}\cup S_{t}^{\mathrm{middle}}\cup S_{t}^{\mathrm{tail}}\]

其中由前 \(w\) 个后缀块构成局部区域（若剩余后缀块总数不足 \(w\) 则全部分配至局部区），其直接毗邻当前解码块，接收极高的注意力权重，承载最为关键的短程细粒度语义上下文；最后一个后缀块被划定为尾部区域，承担生成边界与终末序列规划职能；夹在局部区与尾部区之间的所有过渡块则构成中部区域，该区域注意力权重已处于低位平稳平台期，无需保留密集 token。

2. 非对称结构化 Token 筛选：保留骨架锚点与边界 确定功能分区后，模型根据各区域的职能分配非对称的保留预算。对于局部区域，保留全部后缀 token 以确保局部连贯性：\(\mathcal{R}_{t}^{\mathrm{local}}=S_{t}^{\mathrm{local}}\)。对于中部区域，若完全剔除会导致模型丧失对后文长度和粗粒度结构的感知，而全量保留又带来不必要的计算冗余；统计显示各块起始 token 的注意力普遍高于块内其余 token，因此仅提取中部每个块的起始 token 作为轻量结构骨架：

\[\mathcal{R}_{t}^{\mathrm{middle}}=\{\mathrm{start}(B_{j})\mid B_{j}\in S_{t}^{\mathrm{middle}}\}\]

对于尾部区域，除保留尾块首 token 外，特别强制保留尾块末端 token 以固化终末边界线索：\(\mathcal{R}_{t}^{\mathrm{tail}}=\{\mathrm{start}(B_{\mathrm{tail}}),\mathrm{end}(B_{\mathrm{tail}})\}\)。最终合并得到高度稀疏的保留集 \(\mathcal{R}_t = \mathcal{R}_{t}^{\mathrm{local}}\cup\mathcal{R}_{t}^{\mathrm{middle}}\cup\mathcal{R}_{t}^{\mathrm{tail}}\)。所有保留 token 严格继承其原始位置编码，使得模型能在零微调前提下精准感知跨度距离与整体篇章骨架。

3. 软后缀嵌入混合：跨时间步去噪演化状态传递 针对传统后缀丢弃中 token 每次被粗暴重置为无信息 [MASK] 嵌入导致历史去噪信号丢失的问题，本文提出了软嵌入混合机制。在当前时间步 \(t\)，对保留集中的任意位置 \(i \in \mathcal{R}_t\)，首先提取上一时间步 \(t+1\) 模型在该位置输出的预测概率分布 \(p_{t+1}^{i}\)，筛选置信度最高的 Top-\(k\) 候选候选集 \(\mathcal{T}_{t+1}^{i}\) 并重归一化：

\[\tilde{p}_{t+1}^{i}(v)=\frac{p_{t+1}^{i}(v)}{\sum_{u\in\mathcal{T}_{t+1}^{i}}p_{t+1}^{i}(u)},\quad v\in\mathcal{T}_{t+1}^{i}\]

随后在词表嵌入空间内进行凸组合加权，形成软解码嵌入 \(\tilde{\mathbf{e}}_{t+1}^{i}=\sum_{v\in\mathcal{T}_{t+1}^{i}}\tilde{p}_{t+1}^{i}(v)\cdot\mathbf{e}_{v}\)。最后通过混合因子 \(\alpha \in [0, 1]\) 将原始静态掩码向量 \(\mathbf{e}_{\mathrm{mask}}\) 与软解码嵌入插值：

\[\tilde{\mathbf{e}}_{t}^{i}=(1-\alpha)\cdot\mathbf{e}_{\mathrm{mask}}+\alpha\cdot\tilde{\mathbf{e}}_{t+1}^{i}\]

该机制使保留的后缀骨架 token 具备了类似连续隐状态的信息承载力，在逐步去噪迭代中不断吸收演化的语义置信度，显著稳定了后续去噪轨迹的收敛方向。

一个完整示例¶

假设序列块大小为 32，当前正在去噪第 1 块（index 0-31），后方有 7 个被掩码的后缀块（块 1 至块 7）。设定局部窗口 \(w=2\)： 1. 区域划分：块 1 与块 2 划为局部区（共 64 个 token）；块 3、4、5、6 为中部区；块 7 为尾部区。 2. 结构化保留：局部区保留全部 64 个 token；中部区保留块 3、4、5、6 的首个 token（即全局位置 96, 128, 160, 192 共 4 个 token）；尾部区保留块 7 的首尾 token（位置 224 和 255 共 2 个 token）。全量后缀从原先的 \(7 \times 32 = 224\) 个 token 骤降为 \(64 + 4 + 2 = 70\) 个 token，计算量压缩近 70%。 3. 软嵌入混入：对于保留的这 70 个位置，利用上一去噪步输出的 top-\(k\)（如 \(k=5\)）重归一化概率合成软向量，按 \(\alpha=0.4\) 的强度注入到 [MASK] 向量中输入模型，既大幅削减自注意力矩阵维度，又无缝维持了对全局上下文与生成终止态的敏锐感知。

实验关键数据¶

主实验¶

在 NVIDIA A800 40GB GPU 上评测 LLaDA-8B-Instruct 与 Dream-7B-Base 模型，涵盖数学推理（GSM8K、MATH）与代码生成（HumanEval、MBPP）基准。下表汇总了主实验核心对比数据（对比 Vanilla Top-1、并行解码 Par.、DPad+Par.、Streaming+Par. 及本文 SSM+Par.）：

基准数据集	模型	方法	端到端延迟(s)↓	加速比↑	吞吐(TPS)↑	长度比 \(\bar{\ell}/\ell_{\max}\)	任务精度(%) Strict/Pass@1↑
GSM8K (4-shot)	LLaDA-Instruct	Vanilla	28.52	1.00×	8.13	232 / 256	37.38
GSM8K (4-shot)	LLaDA-Instruct	Par. (并行解码)	8.65	3.30×	26.81	232 / 256	38.67
GSM8K (4-shot)	LLaDA-Instruct	DPad + Par.	6.68	4.27×	24.07	160 / 256	64.97
GSM8K (4-shot)	LLaDA-Instruct	Streaming + Par.	5.25	5.43×	21.86	114 / 256	74.00
GSM8K (4-shot)	LLaDA-Instruct	SSM + Par. (本文)	4.87	5.86×	22.63	110 / 256	76.95
MATH (4-shot)	LLaDA-Instruct	Par. (并行解码)	10.00	2.63×	24.84	249 / 256	8.76
MATH (4-shot)	LLaDA-Instruct	Streaming + Par.	7.35	3.57×	21.41	157 / 256	30.68
MATH (4-shot)	LLaDA-Instruct	SSM + Par. (本文)	6.11	4.30×	21.64	132 / 256	31.96
HumanEval (0-shot)	LLaDA-Instruct	Vanilla	35.09	1.00×	13.47	473 / 512	43.90
HumanEval (0-shot)	LLaDA-Instruct	Par. (并行解码)	11.59	3.03×	40.99	475 / 512	43.29
HumanEval (0-shot)	LLaDA-Instruct	Streaming + Par.	5.41	6.49×	54.32	293 / 512	46.34
HumanEval (0-shot)	LLaDA-Instruct	SSM + Par. (本文)	3.35	10.47×	65.76	220 / 512	48.17
MBPP (3-shot)	LLaDA-Instruct	Par. (并行解码)	14.43	4.35×	20.74	299 / 512	15.00
MBPP (3-shot)	LLaDA-Instruct	Streaming + Par.	6.17	10.18×	16.34	100 / 512	42.40
MBPP (3-shot)	LLaDA-Instruct	SSM + Par. (本文)	2.68	23.44×	19.44	52 / 512	45.20
MBPP (3-shot)	Dream-Base	SSM + Par. (本文)	8.99	5.46×	56.89	512 / 512	55.60

在长序列评测中（GSM8K 结合 LLaDA-1.5，最大长度 1024），SSM 与前缀缓存（Prefix Caching, PC）和并行解码叠加时，延迟由 Vanilla 的 128.16s 降至 1.76s，实现了高达 72.81× 的端到端极限加速，Strict 准确率同时从 48.98% 提升至 75.13%。

消融实验¶

在 LLaDA-Instruct 上对各模块组件进行消融分析（基于 GSM8K 与 HumanEval）：

配置变体	GSM8K 延迟(s) / 加速比	GSM8K Strict(%)	HumanEval 延迟(s) / 加速比	HumanEval Pass@1(%)	机制分析与结论
SSM + Par. (完整框架)	4.87 / 5.86×	76.95	3.35 / 10.47×	48.17	完整模型兼具最优延迟与最高精度
移除局部区域 (w/o Local)	1.96 / 14.55×	25.32	0.51 / 68.80×	15.24	精度雪崩且生成极度截断（6/512），证明邻近密集上下文必不可少
中部去掉起始 Token (w/o Start)	4.95 / 5.76×	75.59	3.59 / 9.77×	43.29	丢失结构骨架导致收敛放缓，延迟反弹且精度显著下降
中部起始换为结尾 Token	4.79 / 5.95×	75.82	3.30 / 10.63×	46.34	块尾 token 难以承载前向骨架引导，精度逊于首 token
尾部仅保留起始 Token	4.95 / 5.76×	76.19	3.30 / 10.63×	43.29	缺失终末界标导致代码生成任务边界感知弱化，Pass@1 掉点近 5%
去掉软嵌入混合 (w/o Soft Embed)	5.06 / 5.64×	73.54	4.26 / 8.24×	48.17	丢失跨步去噪信号，推理延迟上升且 GSM8K 严格匹配准确率下降 3.41%
去掉早期终止 (w/o Early Term)	5.34 / 5.34×	76.95	3.96 / 8.86×	48.17	无法跳过无效后续块计算，增加约 10%-18% 的冗余推理时间

关键发现¶

异构区域功能不可替代：消融表明局部区域是模型生成连贯文本的根基，剔除局部区域会导致灾难性的生成过早截断（HumanEval 实际输出长度仅 6 tokens）；中部块起始 token 则构成了支撑长文本篇章框架的必要支柱；尾部结尾 token 则为模型规划终末界限提供了明确锚点。
后缀精简诱导更紧凑的生成：实验显示 SSM 显著缩短了实际生成长度（如 MBPP 上 LLaDA-Instruct 的实际生成长度由 299 骤减至 52），这并非单纯由 early termination 引起，而是由于大幅减少冗余后缀 token 抑制了模型漫无目的地产生冗余中间幻觉，进而大幅提升了 Strict-match 的标准输出格式遵循能力。
软嵌入权重 \(\alpha\) 存在显著阈值效应：当超参数 \(\alpha\) 在 0.2 至 0.4 区间时，模型性能稳健且在 \(\alpha=0.4\) 时达到最优；但一旦增大至 0.5 以上，早期未收敛的噪声预测过度污染输入表征，将导致严重的表征崩溃与精度暴跌。

亮点与洞察¶

从同质粗暴修剪升维至功能区结构化建模：敏锐捕捉到注意力得分“近高、中平、尾翘”的固有物理先验，把后缀修剪从“粗粒度窗口截断”升级为保留骨架与边界的非对称布设，以极小的 token 预算保全了关键感知能力。
软嵌入跨步桥接破除静态掩码壁垒：巧妙利用上一时间步未完全固化的 Top-\(k\) 预测概率，通过嵌入空间软融合将去噪历史演化信号沉淀进当前步输入，破解了传统离散扩散模型反复重新初始化掩码嵌入的割裂缺陷。
完全免训练与高度正交兼容：纯推理期算法，零模型参数改动，能够无缝与并行解码门限策略、Prefix/KV Caching、EOS 早期截断正交叠加，在千 token 长文本生成场景下解锁了超 70 倍的惊人加速。

局限与展望¶

作者指出的局限：加速收益与后缀长度强相关，当面对极长上下文 Prompt 或极短目标生成长度时，后缀计算占比小，加速幅度有限；此外在少数特定下游任务上仍存在微弱掉点，且尚未在超过 10B 参数量的超大模型上充分验证。
深入洞察与潜在缺陷：固定超参数（如静态窗口 \(w\) 和固定 \(\alpha\)）未能根据去噪时间步的进度自适应退火；在扩散前期的强噪声阶段，Top-\(k\) 软嵌入可能注入较多错误偏差，而在后期接近真实语义时其置信度本应更大。
未来改进方向：可探索随去噪时间步动态自适应衰减/递增的动态软混合调度器（Dynamic \(\alpha_t\)），以及基于注意力熵动态自适应确定局部与中部保留边界的弹性机制。

评分¶

新颖性: ⭐⭐⭐⭐☆ 突破传统局部窗口同质修剪思维，提出功能区非对称保留与软嵌入演化机制，视角清晰独到。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖 3 个先进离散扩散大模型、4 个基准任务，覆盖消融、超参敏感性、长序列及与主流加速方案的全面正交实验。
写作质量: ⭐⭐⭐⭐⭐ 逻辑推导层次分明，动机分析依托坚实的注意力统计实验，公式与机制叙述凝练透彻。
价值: ⭐⭐⭐⭐⭐ 彻底免训练且能与主流算子/系统级优化无缝集成，长文本下超 70 倍的加速表现极大推动了扩散大语言模型的实际落地可行性。