Zarya：一种混合自回归-掩码扩散语言模型，支持灵活训练与双模推理

# Zarya：一种具备灵活训练与双模式推理的混合自回归-掩码扩散语言模型
来源：https://arxiv.org/html/2609.19868
###### 摘要

自回归语言模型（\(ARMs\)）受限于顺序的、从左到右的生成方式，而掩码扩散模型（\(MDMs\)）虽能实现并行解码，却因无法复用键值（KV）缓存而导致计算开销巨大，且因在难以处理的词元组合空间上学习依赖关系而产生不连贯的生成。我们推出Zarya，一种混合语言模型家族，它在单一架构内联合优化自回归（AR）目标与掩码-扩散目标。Zarya将训练数据结构化为可变大小的*槽位*，并采用一种逐步增加槽位粒度的课程学习策略，实现了从细粒度AR学习到粗粒度扩散学习的平滑过渡。在推理阶段，Zarya通过统一接口提供两种不同的解码范式：（i）具有首击去噪的MDM采样，以及（ii）*槽位投机解码*，它将基于槽位间的扩散选择与槽位内的自回归填充交替进行，实现完全的KV缓存复用。训练与推理机制完全解耦，允许使用任意配置训练的模型以任一模式部署。广泛的可配置性——包括分组噪声模式（前缀补全、前缀填充、中间填充）、有序采样调度以及噪声级别置换策略——使得研究探索更加灵活。我们公开发布了0.6B、1.7B和4B参数量的Zarya模型，在标准基准上展示了其性能，同时提供了自回归与扩散范式的原则性整合方案。

## 1引言

自回归模型（\(ARMs\)）在广泛自然语言任务中取得了显著成功Brown et al\., 2020 (https://arxiv.org/html/2609.19868#bib.bib2);OpenAI, 2023 (https://arxiv.org/html/2609.19868#bib.bib7);Touvron et al\., 2023 (https://arxiv.org/html/2609.19868#bib.bib9)\。然而，其顺序、从左到右的解码方式从根本上限制了推理吞吐量，阻碍了并行化Chen et al\., 2023 (https://arxiv.org/html/2609.19868#bib.bib4);Cai et al\., 2024 (https://arxiv.org/html/2609.19868#bib.bib3)\。掩码扩散模型（\(MDMs\)）通过迭代去噪过程提供了一种引人注目的替代方案，无需固定生成顺序即可实现并行生成Li et al\., 2026b (https://arxiv.org/html/2609.19868#bib.bib6)\。然而，大多数MDMs存在两个关键缺陷：（i）它们不包含键值（KV）缓存，导致推理时计算开销高昂；（ii）它们在难以处理的词元组合空间上学习依赖关系，导致生成结果不连贯Li et al\., 2026a (https://arxiv.org/html/2609.19868#bib.bib5)\。

近期的混合方法试图弥合这些范式间的鸿沟。块扩散（\(BD3\-LMs\)）Arriola et al\., 2025 (https://arxiv.org/html/2609.19868#bib.bib1)通过将词元分组为固定大小的块并从左到右解除遮蔽，在AR和MDM之间进行插值，但其块大小固定且缺乏灵活性。奥秘语言模型（\(Eso\-LMs\)）Sahoo et al\., 2026 (https://arxiv.org/html/2609.19868#bib.bib8)在损失层面融合AR和MDM范式，实现了MDMs的KV缓存，并比同时代的MDMs推理速度更快。ReFusionLi et al\., 2026a (https://arxiv.org/html/2609.19868#bib.bib5)将并行解码从词元层面提升到更高的*槽位层面*，将基于槽位的扩散选择与槽位内的自回归填充交替进行，在GSM8KCobbe et al\., 2021 (https://arxiv.org/html/2609.19868#bib.bib11)和MBPPAustin et al\., 2021 (https://arxiv.org/html/2609.19868#bib.bib13)上超越了Qwen3\-8BYang et al\., 2025 (https://arxiv.org/html/2609.19868#bib.bib10)，同时平均速度快了2\.33倍。

我们介绍了Zarya，一种具有以下关键贡献的新型混合架构：

1. 1\.具有逐渐增长槽位长度的训练：我们将预测序列划分为固定长度、连续的子序列，称为槽位。槽位的大小在训练期间逐渐增加。对于每个序列，我们随机遮蔽几个槽位，重新排列输入使干净槽位位于遮蔽槽位之前，并且同时打乱遮蔽和干净槽位的原始顺序。模型同时学习在可见槽位上进行自回归下一词元预测以实现顺序生成，并在遮蔽槽位上进行去噪损失以实现上下文感知的并行重构。
2. 2\.训练-推理解耦：与先前工作（训练配置决定推理模式）不同，Zarya允许任何训练好的模型通过单个推理标志部署为MDM采样或槽位投机解码模式，提供更高的灵活性。两种模式都完全利用带有因果注意力掩码的KV缓存

## 2相关工作

##### 块扩散（\(BD3\-LMs\)）\.

BD3\-LMs通过将词元分组为块并从左到右解除遮蔽，在AR和MDM之间进行插值Arriola et al\., 2025 (https://arxiv.org/html/2609.19868#bib.bib1)\。然而，块大小固定，限制了灵活性。

##### 奥秘语言模型（\(Eso\-LMs\)）\.

Eso\-LMs使用因果注意力融合AR和MDM范式，实现了精确似然计算和MDMs的KV缓存Sahoo et al\., 2026 (https://arxiv.org/html/2609.19868#bib.bib8)\。它们在十亿词数据集LM1B;Chelba et al\., 2014 (https://arxiv.org/html/2609.19868#bib.bib14)和OpenWebTextOWT;Gokaslan et al\., 2019 (https://arxiv.org/html/2609.19868#bib.bib15)上取得了扩散模型中的低困惑度。

##### ReFusion\.

ReFusion引入了槽位级并行解码，将生成从词元提升到固定长度的槽位Li et al\., 2026a (https://arxiv.org/html/2609.19868#bib.bib5)\。它将基于槽位的扩散选择与槽位内的自回归填充交替进行，在每次迭代后将新生成的槽位重排到剩余遮蔽之前。这种设计实现了完全的KV缓存复用，并将学习复杂度从难以处理的词元组合空间降低到可控的槽位级排列空间。然而，ReFusion推理仅适用于前缀补全任务。

## 3Zarya架构

### 3\.1模型骨干与配置

Zarya将一个Qwen3骨干Yang et al\., 2025 (https://arxiv.org/html/2609.19868#bib.bib10)封装在自定义的Zarya类中，并注册到Hugging Face的AutoModel和AutoConfig系统中。我们分别从Qwen3\-0\.6B、Qwen3\-1\.7B和Qwen3\-4B检查点初始化Zarya，并在一个涵盖数学、编码和通用指令遵循任务的多样化3200万样本指令SFT数据集（约377亿词元）上进行微调。我们发布了三种尺寸的Zarya（参见表1 (https://arxiv.org/html/2609.19868#S3.T1)）。

由于Qwen3骨干保留了因果注意力，遮蔽位置无法关注未来的遮蔽位置。因此，Zarya使用的扩散目标是一个因果掩码重构目标，而非掩码扩散语言模型中常用的完全双向掩码词元目标。可见和遮蔽位置的重排确保了所有遮蔽位置都可以关注可见前缀，同时保持了因果注意力模式和KV缓存兼容性。

表1：Zarya模型的模型架构该架构以扩散特定的超参数进行参数化，扩展了常见的AR模型参数：

- •α0\\alpha\_\{0\}\(默认0\.25\)和ε\\epsilon\(默认0\.001\)：线性噪声调度αt=α0\(1−t\)\\alpha\_\{t\}=\\alpha\_\{0\}\(1\-t\)的参数；
- •λ=diffusion\_loss\_proportion\\lambda=\\texttt\{diffusion\\\_loss\\\_proportion\}\(默认0\.5\)：MDM和AR损失之间的权重；
- •sequential\_shuffle/diffusion\_shuffle：控制每个阶段的槽位/词元打乱；
- •ordered\_sampling：从左到右单调增加pmaskp\_\{\\text\{mask\}\}；
- •grouped\_noise和max\_span\_length：启用前缀补全、前缀填充、中间填充和随机跨度掩码模式；
- •noise\_sorting：在向前传递前根据遮蔽/未遮蔽状态重新排列词元\.

## 4训练

### 4\.1槽位训练

当slotted\_training=True时，forward\_process\(\)按如下方式转换每个批次（此模式深受Li et al\. \(2026a\) (https://arxiv.org/html/2609.19868#bib.bib5)启发）：

1. 1\.槽位划分：每个答案被分割成大小为slot\_size（来自slot\_size\_set，例如发布的模型检查点为\[2,4,8,16,32,64\]）的槽位。
2. 2\.掩码采样：对于批次中的每个样本，均匀采样掩码概率（pmaskp\_\{\\text\{mask\}\}）。这决定了将有多少比例的槽位被遮蔽（作为扩散任务处理）。
3. 3\.槽位分配： - •*AR槽位*：词元保持可见；模型预测每个槽位内的下一个词元，产生Lseq\\mathcal\{L\}\_\{\\text\{seq\}\}。 - •*MDM槽位*：所有词元被特殊的⟨mdm\_mask⟩\\langle\\texttt\{mdm\\\_mask\}\\rangletoken替换；模型重构原始词元，产生Ldif\\mathcal\{L\}\_\{\\text\{dif\}\}。槽位内分配了pmaskp\_\{\\text\{mask\}\}值的每个词元，在每词元归一化中用作1pmask\\frac\{1\}\{p\_\{\\text\{mask\}\}\}权重。

最终损失是线性组合：

L=λ⋅Ldif\+\(1−λ\)⋅Lseq,\\mathcal\{L\}=\\lambda\\cdot\\mathcal\{L\}\_\{\\text\{dif\}\}\+\(1\-\\lambda\)\\cdot\\mathcal\{L\}\_\{\\text\{seq\}\},\(1\)其中λ\\lambd是扩散损失比例。

这种设计使模型能够在*相同*输入上同时学习下一词元预测（AR）和掩码词元重构（MDM），并以明确的槽位边界提供结构化的归纳偏置。

可选标志包括：

- •ordered\_sampling：在每个槽位内修改每词元的pmaskp\_\{\\text\{mask\}\}使其从左到右增加，由于Ldif\\mathcal\{L\}\_\{\\text\{dif\}\}按1pmask\\frac\{1\}\{p\_\{\\text\{mask\}\}\}值缩放，这使得模型学会正确预测槽位起始词元更为重要；

#### 4\.1\.1槽位大小课程

槽位大小通过slot\_step\_borders（周期或步长阈值）在训练期间演变。例如，使用slot\_size\_set=\[2,4,8,16,32,64\],\\displaystyle=\[2,4,8,16,32,64\],slot\_step\_borders=\[0\.06,0\.2,0\.4,0\.6,0\.8,1\.0\],\\displaystyle=\[0\.06,0\.2,0\.4,0\.6,0\.8,1\.0\]，槽位大小在训练过程中从2逐渐增加到64。这种课程学习将模型从细粒度的AR学习（小槽位，多次预测）引导至粗粒度的扩散学习（大槽位，整体重构），在两种范式间提供了平滑过渡。值得注意的是，槽位大小的增加也会增加计算的损失（参见图1 (https://arxiv.org/html/2609.19868#S4.F1)）。

参见图例图1：Zarya\-4B模型训练过程中的损失曲线。槽位大小逐步增加也导致损失增加

### 4\.2非槽位训练

当slotted\_training=False时，模型在未划分槽位的原始序列上运行目标（此模式深受Sahoo et al\. \(2026\) (https://arxiv.org/html/2609.19868#bib.bib8)启发）：

1. 1\.对于批次中的每个样本，均匀采样噪声级别\(tt\)。
2. 2\.构建带噪输入。噪声调度为αt=α0\(1−t\)\\alpha\_\{t\}=\\alpha\_\{0\}\(1\-t\)，每词元掩码概率为pmask=1−αtp\_\{\\text\{mask\}\}=1\-\\alpha\_\{t}\)。
3. 3\.运行两次独立的向前传递： - •*顺序阶段*：干净序列x0x\_\{0\}，仅预测遮蔽位置的词元（其他设为ignore\_index）。 - •*扩散阶段*：带噪序列xtx\_\{t\}，重构遮蔽词元。
4. 4\.最终损失是顺序阶段和扩散阶段损失的线性组合，计算方式与第4\.1节 (https://arxiv.org/html/2609.19868#S4.SS1)相同。

可选标志包括：

- •noise\_sorting：在两次向前传递前根据遮蔽/未遮蔽状态重新排列词元，之后logits再置换回去；
- •grouped\_noise：遮蔽连续跨度（前缀补全、前缀填充、中间填充、随机跨度），最大长度为max\_span\_length；
- •ordered\_sampling：按词元位置偏移tt，使pmaskp\_\{\\text\{mask\}\}从左到右增加。

## 5推理：双模式统一接口

Zarya通过单个model\.generate\(\)调用提供两种不同的解码范式，通过generation\_config\.slotted\_generation进行路由。关键是，训练与推理机制完全解耦：使用slotted\_training=False训练的模型仍可以slotted\_generation=True部署，反之亦然。

### 5\.1模式A：MDM采样

当slotted\_generation=False时，推理代码执行首击去噪过程Zheng et al\., 2025 (https://arxiv.org/html/2609.19868#bib.bib16);Sahoo et al\., 2026 (https://arxiv.org/html/2609.19868#bib.bib8)\。从用⟨mdm\_mask⟩\\langle\\texttt\{mdm\\\_mask\}\\rangletoken填充到max\_length的提示开始，模型迭代地揭示词元：

1. 1\.掩码预算规划：\_tokens\_unmasked\_per\_step\(\)确定每步揭示多少遮蔽。α0\\alpha\_\{0\}是扩散过程生成的遮蔽词元的预期比例。 - •如果离散化步数设置为T\>0T\>0：使用精确的TT扩散步。使用二项分布计算通过扩散过程去噪的遮蔽词元数量，以及剩余顺序去噪的词元（此模式深受Sahoo et al\. \(2026, Appendix B\.5\) (https://arxiv.org/html/2609.19868#bib.bib8)启发）。 - •如果离散化步数设置为T=0T=0（忽略噪声计算）：TT步数自动计算为遮蔽词元的14\\frac\{1\}\{4\}。
2. 2\.重排：输入序列被重排，使遮蔽词元始终位于未遮蔽词元之后。
3. 3\.从分类分布中逐步步采样：在每一步，模型接收逐步填充的序列（复用KV缓存）并为遮蔽位置生成logits。添加Gumbel噪声Gumbel, 1935 (https://arxiv.org/html/2609.19868#bib.bib17);Zheng et al\., 2025 (https://arxiv.org/html/2609.19868#bib.bib16), Appendix F)用于分类采样，并尊重标准采样参数（temperature,top\_p,repetition\_penalty）。
4. 4\.恢复：所有步骤后，序列恢复到原始词元顺序。

之所以能复用KV缓存，是因为序列被重排，使得在给定去噪步骤其值已固定的词元位于剩余遮蔽位置之前。在因果注意力下，当遮蔽位置逐步填充时，缓存的前缀状态因此保持不变。对于双向掩码-扩散注意力模式，此属性不成立，因为更改任何先前遮蔽的词元都可能影响其他遮蔽位置的表示。

### 5\.2模式B：槽位投机解码

当slotted\_generation=True时，推理代码通过将解码单元从词元提升到槽位来实现并行化，完全复用KV缓存以避免重新计算。此模式遵循ReFusion范式Li et al\., 2026a (https://arxiv.org/html/2609.19868#bib.bib5)：

1. 1\.重排：在第一次向前传递前，输入序列被重排，使遮蔽词元始终位于未遮蔽词元之后。
2. 2\.块构建：max\_new\_tokens数量的遮蔽词元被分成serial\_num\_blocks个长度为block\_size=⌊max\_new\_tokens/serial\_num\_blocks⌋\\texttt\{block\\\_size\}=\\lfloor\\texttt\{max\\\_new\\\_tokens\}/\\texttt\{serial\\\_num\\\_blocks\}\\rfloor的块。在每个块内，词元被分组为大小为slot\_size的槽位。如果max\_new\_tokens较小，serial\_num\_blocks会被回