Zyphra发布ZAYA1-8B-Diffusion-Preview：首个从自回归LLM转换的MoE扩散模型，速度提升达7.7倍

Zyphra，这家总部位于旧金山的AI实验室，发布了ZAYA1-8B-Diffusion-Preview，这是其扩散语言模型早期工作的预览版。该发布表明，现有的自回归语言模型可以转换为离散扩散模型，而不会造成评估性能的系统性损失，同时在AMD硬件上实现了显著的推理加速。

自回归解码的问题 要理解这一突破的重要性，首先需要了解大多数语言模型目前生成文本的方式。标准大型语言模型是自回归的：它们按顺序逐个解码token。对于每个新token，注意力机制必须回顾所有之前生成的token，并从GPU内存加载它们的存储表示（即KV缓存）。关键问题是，由于批次中每个用户的token历史不同，每个用户的KV缓存必须单独加载，无法在请求之间共享。这造成了瓶颈。当GPU花费更多时间从内存移动数据而非执行实际计算时，系统便受限于内存带宽而非计算能力。这限制了现代GPU硬件（其计算FLOPs增长速度超过内存带宽）在推理过程中的高效利用。

扩散模型提供了另一种方案。扩散模型不是逐个生成token，而是同时生成N个token的多个草稿，并多次迭代这一草稿过程。由于块中的所有N个token共享相同的KV缓存，操作从内存带宽受限转变为计算受限，这意味着GPU可以更高效地被利用。在ZAYA1-8B-Diffusion-Preview中，模型对块中的每个token执行从掩码到token的单步变换——即直接预测未被掩码的token，而非迭代去噪。

从自回归到扩散的转换：无需从头训练 从头训练扩散语言模型在技术上很困难，且已知的成熟方法很少。Zyphra团队给出了选择转换而非从头训练的两个理由：首先，从头训练本身就很难，已知配方极少；其次，以扩散模式训练并无优势，因为训练本身已是计算受限——扩散解决的带宽瓶颈仅出现在推理阶段。这意味着扩散的所有优势都是推理时的优势，现有的预训练堆栈可以原样复用。

基于TiDAR配方，Zyphra以ZAYA1-8B-base检查点为起点，额外进行了6000亿token的扩散转换中训练（上下文长度32k），随后进行了5000亿token的原生上下文扩展至128k，最后进行了扩散监督微调（SFT）阶段。

ZAYA1-8B-Diffusion-Preview是首个从自回归LLM转换的MoE扩散模型，也是首个在AMD GPU上训练的扩散语言模型。Zyphra报告称，与基础自回归检查点相比，评估性能几乎没有下降，在某些基准（如LCB-v6）上甚至有所提升。他们将此部分归因于改进的中训练数据集，以及扩散式块内非因果推理相比因果自回归具有更强的表达能力。

扩散采样器的工作原理 推理过程中，ZAYA1-8B-Diffusion-Preview同时生成16个token的草稿。基于从推测解码借鉴的采样标准，接受其中一部分token。关键在于，同一个模型在前向传播中同时扮演推测器和验证器的角色，这消除了传统方法（如EAGLE或dFlash）中运行两个独立模型的开销。在高度内存带宽受限的情况下，几乎所有被接受的token都代表了对自回归解码的免费加速——GPU已经加载完毕，额外token的计算成本极低。

Zyphra团队报告了两种具有不同速度-质量权衡的采样器：

• 无损扩散采样器：使用标准推测解码接受标准min(1, p(x)/q(x))，其中p是自回归模型的logit分布，q是扩散模型的分布。拒绝时，下一个token从p(x)-q(x)的残差分布中采样。该采样器实现4.6倍加速，且没有系统性的评估性能下降。
• Logit-mixing采样器：首先混合扩散推测器和自回归模型的logit，然后使用平均分布进行验证。这提高了接受率，因为验证logit更接近扩散logit，但对质量有一定影响。该采样器实现7.7倍加速。速度与质量的权衡可在运行时选择。

需要注意的是，由于ZAYA1-8B-Diffusion-Preview是未经RL训练的基础中训练检查点，Zyphra使用pass@评估而非标准准确率基准，以更好地代表模型在RL训练后的最终潜力。读者在将这些数字与其他模型报告的基准进行比较时应牢记这一点。

Zyphra团队还指出，扩散带来的加速高于多token预测（MTP）和EAGLE3等推测解码策略的替代方法。由于TiDAR风格的扩散模型仅使用单次前向传播，即使接受率与dFlash相当，也能实现显著的加速。

架构细节 ZAYA1-8B-Diffusion-Preview是一种单步推测扩散模型，采用顺序约束生成，即扩散模型仅能生成从前缀开始的连续子序列。与无约束掩码扩散目标或集合块解码相比，这一约束极大地提高了训练稳定性，这也是Zyphra基于TiDAR配方的主要原因。

该模型使用Zyphra的CCA注意力变体。CCA大幅减少了注意力中的预填充FLOPs，这对扩散直接有利，因为扩散将解码转变为类似预填充的操作。这意味着CCA允许模型在达到计算限制之前并行处理更多token。

具体来说，架构采用CCGQA，查询头与键头的比例为4:1。设计选择之一是故意避免MLA（多头潜在注意力），因为其高算术强度与CCGQA不匹配。由于块扩散访问相同的缓存，算术强度随块大小和每次前向传播的块数量而扩展。在AMD MI300x硬件上（bf16），系统每次前向传播大约支持三个块大小的提议；在MI355x上，这一数字上升到约五个。CCGQA还以2倍压缩运行，这使得Zyphra能够负担与TiDAR中训练相关的额外训练FLOPs。AMD GPU硬件更大的VRAM容量进一步使得扩散训练整体上更高效。

在实践中，实现理论加速更具挑战性，因为扩散带来了额外的操作开销，且扩散模型的推理堆栈远不如自回归推理的成熟工具那样优化。