Orthrus: Memory-Efficient Parallel Token Generation via Dual-View Diffusion¶

会议: NeurIPS 2026
arXiv: 2605.12825
OpenReview: https://openreview.net/forum?id=fwaVRSuVPd
代码: https://github.com/chiennv2000/orthrus
领域: LLM 效率
关键词: 并行解码, 扩散语言模型, 投机解码, 键值缓存复用, 无损加速

一句话总结¶

Orthrus 在冻结的自回归大模型中注入轻量扩散注意力头，首创完全共享单套 KV Cache 的双视角架构，通过扩散单步并行投影与自回归共识验证，以 \(O(1)\) 显存增量实现高达 7.8 倍的严格无损推理加速。

研究背景与动机¶

自回归（Autoregressive, AR）大语言模型因其严格的因果依赖建模，在复杂推理与代码生成等长程任务中展现出卓越的输出质量。然而在实际部署中，自回归推理面临严重的计算瓶颈：预填充（Prefill）阶段能够利用自注意力对提示词进行高度并行的上下文编码，但生成（Decoding）阶段却必须逐个生成 token。这种序列化循环迫使硬件频繁从高带宽显存中重复加载庞大的键值缓存（KV Cache），导致严峻的显存带宽瓶颈与算力利用率低下。

为打破序列化束缚，以块级并行生成见长的扩散语言模型（Diffusion Language Models, DLMs）成为研究热点。然而现存扩散语言模型受制于条件独立性假设，在多步去噪中极易出现严重的条件漂移（Conditional Drift），导致复杂推理精度断崖式下跌；即使采用连续预训练将成熟 AR 模型转变为扩散模型，也需要消耗数百亿乃至数千亿 token 的高昂算力，依然无法避免偏离基线模型的真实目标分布。另一方面，传统投机解码（Speculative Decoding）依赖独立的小草稿模型，推理时必须额外维护一份相互割裂的草稿 KV Cache，带来了随序列长度线性膨胀的显存冗余负担。

面对生成保真度与并行吞吐量之间的根本矛盾，本文提出关键洞察：自回归机制的性能瓶颈仅仅存在于解码阶段，其自注意力对历史上下文的编码表征是近乎最优且不可替代的。核心 idea：在完全冻结的自回归大模型上引入仅占 16% 参数的轻量扩散注意力模块，双视角完全共用单套高保真历史 KV Cache，由扩散头执行单步并行候选块投影，再由冻结自回归头在单次前向中完成序列级分布一致性验证与纠偏，以常数级 \(O(1)\) 显存增量实现严格无损且高达 7.8 倍的并行解码加速。

方法详解¶

整体框架¶

Orthrus 将并行的块级生成机制与因果自回归模型内生统一于同一个 Transformer 结构中，构建了兼顾高保真表征与并行吞吐的双视角（Dual-View）架构。系统整体流程包含三大核心阶段：首先在上下文预填充阶段，完全冻结的原生自回归主干执行标准因果前向，生成高质量因果历史 KV Cache；在后续的解码循环中，每个周期包含两次前向传播——第一步由轻量可训练的扩散头基于共享 KV Cache 和当前锚点 token，在单次前向内并行预测包含 \(K\) 个候选 token 的完整块；第二步将该候选块整体输入冻结的自回归头进行并行前向验证，计算真实的因果条件概率分布。最后通过模型内共识机制逐位置比对，若在前 \(j\) 个位置保持一致，则无损提交这 \(j-1\) 个投影 token 及自回归给出的第 \(j\) 个纠偏 token，并将共享 KV Cache 无缝截断同步至最新位置，进入下一轮迭代。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入 Prompt 序列"] --> B["上下文预填充<br/>冻结 AR 主干构建高保真 KV Cache"]
    B --> C["共享统一 KV Cache 的双视角注意力机制<br/>AR 头与扩散头无缝复用历史缓存"]
    C --> D["结构化双通道掩码与单步并行投影<br/>扩散头单次前向预测 K 个候选 Token"]
    D --> E["模型内分布匹配与无损共识验证<br/>冻结 AR 头单次并行校验并自回归纠偏"]
    E --> F["更新共享 KV Cache 并输出无损 Token 序列"]

关键设计¶

1. 共享统一 KV Cache 的双视角注意力机制：消除传统投机解码的双重显存冗余

传统投机解码为了提速通常引入独立的草稿模型，这导致系统必须在显存中同时维护验证模型和草稿模型的两套 KV Cache，显存开销随长上下文呈线性倍增。Orthrus 在每个 Transformer 层内部直接并联一个轻量化的可训练扩散注意力模块，其投影参数 \((W_{Q_{\text{diff}}}, W_{K_{\text{diff}}}, W_{V_{\text{diff}}})\) 直接由对应的冻结 AR 权重初始化。在生成并行 token 块时，扩散查询不仅在块内 mask token 之间做双向自注意力交互，更重要的是直接接入冻结 AR 主干所生成的历史键值状态 \((K_{\text{AR}}, V_{\text{AR}})\)：

\[O_{\text{diff}} = \text{Softmax}\left(\frac{Q_{\text{diff}} [K_{\text{AR}} \parallel K_{\text{diff}}]^\top}{\sqrt{d_{\text{head}}}}\right) [V_{\text{AR}} \parallel V_{\text{diff}}]\]

由于历史上下文的 KV Cache 是就地直接复用，扩散视角完全无需分配任何额外的历史 KV 存储空间，全模型仅需常数级 \(O(1)\) 的瞬态显存开销（在并行块长 \(K=32\) 时仅约为 4.5 MiB），并且微调参数量仅占总模型参数的约 16%，彻底根除了外置草稿模型带来的带宽冗余与显存膨胀难题。

2. 结构化双通道掩码与单步并行投影：严防上下文穿越并杜绝多步去噪开销

现有扩散语言模型通常需要 2 到 16 步迭代去噪才能生成较为连贯的文本，这大幅摊薄了并行计算的吞吐优势。Orthrus 抛弃了低效的多步去噪，采用单步确定性投影机制：在当前生成步，系统以当前 anchor token \(x_t\) 为起点，拼接 \(K-1\) 个 <mask> token 构造出待预测块。为了在训练和推理中统一信息流并防止未来数据穿越，设计了基于 FlexAttention 实现的结构化掩码矩阵 \(M_{\text{diff}}\)：

\[M_{\text{diff}}[q, k] = \mathbf{1}[k < L] \cdot \mathbf{1}[k \leq a_b - 1] \mid \mathbf{1}[k \geq L] \cdot \mathbf{1}\left[\lfloor q/K \rfloor = \lfloor (k - L)/K \rfloor\right]\]

该掩码严密划定了两套互不干扰的注意力通道：对于处于并行预测块内部的查询，其向前可以因果式关注 anchor 位置之前的所有干净 AR 历史上下文，严禁窥探后续真实 token；而在此块内部的 \(K\) 个位置之间则允许完全双向互视，使得模型能在单个前向步骤中充分聚合块内上下文，单步直接映射出 \(K\) 个候选 token 的联合先验分布。

3. 模型内分布匹配与无损共识验证：严格保障数学等价并消除条件漂移

纯扩散生成模型在长文本或数学推理任务中往往因累积误差而出现严重的幻觉与性能退化。Orthrus 将扩散生成的候选序列 \(\hat{\mathbf{y}} = (\hat{y}_1, \dots, \hat{y}_K)\) 直接路由回冻结的自回归注意力头中。由于输入序列的所有 \(K\) 个位置已被完全显式填充，自回归头利用其因果掩码，可以在单次前向传播中同时算出所有位置关于历史和前序候选的精确后验概率 \(p_{\text{AR}}(v \mid \mathbf{x}_{\leq t}, \hat{\mathbf{y}}_{1:k-1})\)。在贪婪解码模式下，系统执行从左至右的严格同构判定：

\[\hat{y}_k = \arg\max_{v \in \mathcal{V}} p_{\text{AR}}(v \mid \mathbf{x}_{\leq t}, \hat{\mathbf{y}}_{1:k-1})\]

若在索引 \(j\) 处首次检测到扩散预测与 AR 输出不一致，验证过程立即终止。系统当即提交已完全匹配的合法前缀 \(\hat{\mathbf{y}}_{1:j-1}\)，并无缝追加由自回归分布生成的标准纠偏 token \(y_j \sim p_{\text{AR}}\)。在非贪婪随机采样下则自动转为无偏拒绝采样。这一内生共识机制使得 Orthrus 在算法层面完全等价于原生 AR 基座的输出分布，达成了绝对无损的加速保证。

损失函数 / 训练策略¶

训练阶段完全冻结自回归主干权重，仅优化注入的扩散注意力参数。从训练语料中随机抽取 \(B\) 个起始锚点并构造长度为 \(K\) 的遮蔽块，训练目标采用前向 KL 散度（Forward KL Divergence），将冻结 AR 教师头的完整预测分布作为软蒸馏信号迁移给扩散视角：

\[\mathcal{L}_{\text{Orthrus}} = \mathbb{E}_{\mathbf{x}, \{a_b\}}\left[\sum_{b=1}^B \sum_{k=1}^K D_{\text{KL}}\left(p_{\text{AR}}(\cdot \mid \mathbf{x}_{\leq a_b + k - 1}) \parallel p_{\text{diff}}(\cdot \mid \mathbf{x}_{< a_b}, \tilde{\mathbf{y}}_b)\right)\right]\]

训练数据采样自开源高质量数据集 Nemotron-Post-Training-Dataset-v2，按 1:1:1 比例严格均衡分布在数学推理、代码生成与通用指令三大领域，打包为 2048 长度的序列。模型基于 PyTorch FSDP-2 与 FlashAttention-4/FlexAttention 在单台 8×H200 节点上采用 bfloat16 混合精度优化 2 个 epoch，总计消耗仅 0.96B token（训练耗时小于 24 小时），峰值学习率为 \(2 \times 10^{-4}\)，并配备 5% 的线性 warmup 与余弦学习率衰减。

实验关键数据¶

主实验¶

实验基于 Qwen3（1.7B、4B、8B）系列模型在数学推理（GSM8K、MATH-500、AIME24、AIME25）与代码算法（HumanEval、MBPP、Pseudo2code、LiveCodeBench-v5）等核心基准上展开评估。基线序列化自回归解码的单次前向生成效率 TPF（Tokens Per Forward Pass）为 1.0，加速比为 1.0×。

模型	基准任务	贪婪解码 TPF (T=0)	贪婪加速比 (T=0)	采样解码 TPF (T=1)	采样加速比 (T=1)
Qwen3-1.7B	GSM8K	4.20	4.37×	3.88	4.04×
Qwen3-1.7B	MATH-500	4.71	4.74×	4.47	4.50×
Qwen3-1.7B	Pseudo2code	4.60	4.90×	4.37	4.65×
Qwen3-1.7B	平均值 (Average)	3.89	4.25×	3.65	4.27×
Qwen3-4B	GSM8K	4.91	4.48×	4.97	4.53×
Qwen3-4B	MATH-500	6.02	5.76×	5.62	5.38×
Qwen3-4B	Pseudo2code	7.40	7.38×	6.95	6.93×
Qwen3-4B	平均值 (Average)	5.04	5.20×	4.72	5.13×
Qwen3-8B	GSM8K	5.30	5.22×	4.69	4.62×
Qwen3-8B	MATH-500	6.35	5.95×	5.80	5.43×
Qwen3-8B	AIME24	5.63	6.81×	4.70	5.68×
Qwen3-8B	Pseudo2code	7.51	7.83×	7.29	7.60×
Qwen3-8B	平均值 (Average)	5.39	5.36×	4.43	5.02×

在对比其他前沿扩散模型的生成质量方面，Orthrus 展现出压倒性优势：在 MATH-500 上，Dream-7B 仅取得 39.6% 准确率，Fast-dLLM-v2 仅取得 61.5% 准确率，同源自回归微调的 SDAR-Qwen3-8B 取得 78.6% 准确率，而 Orthrus-Qwen3-8B 达到 86.2%（与 Qwen3-8B AR 原始精度完全 100% 相同）；在 HumanEval 上，Fast-dLLM-v2 仅为 43.3%，而 Orthrus 达到 95.1%。

消融实验¶

消融 1：去噪迭代步数对有效吞吐量的影响（MATH-500 基准）

去噪策略	前向传播次数 (Passes)	单步有效吞吐 TPF	说明
多步迭代去噪 (Multi-Step)	2	3.53	沿用 Fast-dLLM-v2 策略，额外前向显著拖慢整体吞吐
Orthrus (单步投影)	1	6.35	单次前向直接生成高质量块，TPF 提升 1.8 倍

消融 2：训练目标函数对验证阶段接受率的影响（MATH-500 基准）

训练目标 (Objective)	推理准确率 (Accuracy)	单步有效吞吐 TPF	说明
硬标签交叉熵 (Cross-Entropy)	86.2%	5.86	拟合表面语法而非基座因果轨迹，导致验证拒绝率上升
软标签 KL 散度 (KL Divergence)	86.2%	6.35	完全习得 AR 教师的概率分布空间，接受率与吞吐达到最优

关键发现¶

并行块长扩展效应良好：并行预测块长从 \(K=4\) 扩大至 \(K=32\) 时，由于扩散视角是单次前向与共享缓存交互，单次推理延迟保持平坦不变，而单步有效吞吐 TPF 从 1.76 跃升至 6.35，带来了 3.6 倍的纯净吞吐提升。
软分布蒸馏至关重要：使用软标签 KL 散度相比直接使用硬标签交叉熵，使 TPF 从 5.86 提升至 6.35。这表明扩散头只有精准复现自回归教师的连续概率流，才能最大化通过共识验证机制。
超长上下文吞吐极其稳健：传统投机解码（如 DFlash）在上下文长度扩展至 40K 时，因双 KV Cache 的内存与带宽开销暴涨导致生成速度腰斩；而 Orthrus 的 KV Cache 显存增量严格恒定在约 4.5 MiB（\(O(1)\)），在 40K 超长上下文中依然保持近乎满速的稳定吞吐。

亮点与洞察¶

零冗余 KV 缓存共享机制：打破了以往投机解码和扩散加速框架中主模型与辅助模块割裂管理上下文的惯性思维，通过在 Transformer 内并联扩散头直接读取原生自回归缓存，巧妙达成了零历史缓存冗余。
软蒸馏与结构化掩码的优雅闭环：借助 FlexAttention 构建兼具因果上下文防穿越与块内双向互视的结构化掩码，仅用不到 1B token 的微调就在 24 小时内为大模型内生注入了高质量的并行解码能力。
严格数学无损的推理底色：通过模型内一致性验证和拒绝采样，彻底根除了纯扩散模型在逻辑推理与长程依赖中无法克服的条件分布漂移弊端，为学术与工业界提供了一种即插即用的高质量加速范式。

局限与展望¶

生成能力严格受限于基座上限：由于自回归主干完全冻结且扩散头仅对齐 AR 概率分布，Orthrus 本质是一个无损推理加速器，基座模型固有的事实性幻觉或知识盲区无法通过该架构得到改善。
并行块长收益存在饱和边界：当并行块长 \(K\) 超过一定阈值后，后续位置的条件不确定性指数级上升，长尾 token 被验证拒绝的概率显著增大，限制了更大候选块的实用性。
未来方向：探索在推理解码过程中根据输入语义熵动态调整候选块长 \(K\)，以及将该双视角机制拓展到多模态大模型与长文本思维链推理的进一步加速中。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 巧妙提出 Transformer 内部并联扩散头与共享 KV Cache 的双视角机制，彻底破除了投机解码的冗余缓存弊端。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖 1.7B 到 8B 多尺寸模型及 8 项数理推理与代码基准，包含消融、吞吐测试、长上下文及显存分析。
写作质量: ⭐⭐⭐⭐⭐ 动机阐述深刻，数学建模严谨，注意力掩码与双视角设计表达极为清晰。
价值: ⭐⭐⭐⭐⭐ 训练极低成本（<1B tokens / 24h）、显存极省（O(1) 增量）、完全无损加速（最高 7.8×），极具工业落地与开源适配价值。