掩码语言流模型

# 掩码语言流模型 来源：https://arxiv.org/html/2606.27617 ###### 摘要 *掩码扩散模型*（MDMs）承诺实现快速、并行的语言生成，但其反向转移在词元位置上是因子化的——这种近似在少步采样（并行生成本应带来最大效率提升的领域）中会失效。*流语言模型*（FLMs）通过学习一个连续流来规避这一限制，该流将噪声输送到以欧几里得空间表示的干净序列中，并诱导出一个可被蒸馏以实现单步生成的流映射。然而，这使得需要多步推理的复杂任务对FLMs来说变得困难，因为FLMs在生成过程中必须解码每个词元。为了解决这个问题，我们引入了**掩码语言流模型**（MLFMs），它通过使用连续随机插值法将掩码融入FLMs中，以连接部分掩码序列和干净序列。这种设计通过连续流实现了条件生成，并允许通过简单、轻量级的适配将预训练的MDMs转换为MLFMs。利用这种灵活性，我们提出了一种新颖的采样器，它在连续去噪与对置信词元进行离散去掩码之间交替进行，以更好地支持多步推理。我们在GSM8K和MT-Bench上评估了我们的方法，并首次发现基于流的语言模型可以扩展到解决下游推理和指令遵循任务。代码可在以下网址获取：github.com/imbirik/mlfm (https://github.com/imbirik/mlfm)。

## 1 引言

*自回归模型*（ARMs）通过将语言生成建模为下一个词元预测（Brown 等人，2020 (https://arxiv.org/html/2606.27617#bib.bib6)），推动了语言建模领域的诸多最新进展。然而，它们对联合分布的从左到右的因子化使得这一过程本质上是顺序的，因为每个词元都必须以前面所有词元为条件。因此，推理成本随输出长度线性增长，为长序列创造了瓶颈。这一局限性激发了对*掩码扩散模型*（MDMs）（Austin 等人，2021 (https://arxiv.org/html/2606.27617#bib.bib4); Campbell 等人，2022 (https://arxiv.org/html/2606.27617#bib.bib7); Lou 等人，2024 (https://arxiv.org/html/2606.27617#bib.bib29); Shi 等人，2024 (https://arxiv.org/html/2606.27617#bib.bib39); Sahoo 等人，2024 (https://arxiv.org/html/2606.27617#bib.bib38)）日益增长的兴趣，它们用掩码词元的并行解码取代了顺序生成。尽管MDMs在语言建模和下游任务中表现出了强大的性能（Nie 等人，2025a (https://arxiv.org/html/2606.27617#bib.bib30), b (https://arxiv.org/html/2606.27617#bib.bib31)），但它们通常依赖于跨掩码位置的因子化反向转移，使得在组合上巨大的离散状态空间中进行推理变得可行。这种近似在无穷小步长极限下是准确的，但在少步长机制（Deschenaux 和 Gulcehre, 2024 (https://arxiv.org/html/2606.27617#bib.bib16)）中变得高度受限，因为每一步必须同时独立地解码许多掩码词元——忽略了它们之间的依赖关系。因此，正是并行解码承诺带来最大加速的那些机制（Dieleman, 2023 (https://arxiv.org/html/2606.27617#bib.bib17); Zheng 等人，2024 (https://arxiv.org/html/2606.27617#bib.bib47)）恰恰是这种独立性假设最严重地损害生成质量的地方。

*流语言模型*（FLMs）（Roos 等人，2026 (https://arxiv.org/html/2606.27617#bib.bib37); Lee 等人，2026 (https://arxiv.org/html/2606.27617#bib.bib25); Potaptchik 等人，2026 (https://arxiv.org/html/2606.27617#bib.bib35); Davis 等人，2026 (https://arxiv.org/html/2606.27617#bib.bib13); Chen 等人，2026 (https://arxiv.org/html/2606.27617#bib.bib10); Hu 等人，2026 (https://arxiv.org/html/2606.27617#bib.bib23)）通过从离散状态空间转移到连续状态空间来解决这个问题。具体来说，FLMs学习一个流（Song 等人，2021 (https://arxiv.org/html/2606.27617#bib.bib41); Lipman 等人，2023 (https://arxiv.org/html/2606.27617#bib.bib27); Albergo 等人，2025 (https://arxiv.org/html/2606.27617#bib.bib3)），将噪声输送到词元序列的连续嵌入中，使得潜在状态能够在词元位置之间联合演化。FLM的一个显著优点是它们也诱导出一个流映射，该映射可以被蒸馏以支持少步甚至单步生成。然而，将生成过程坍缩为单个流映射对于需要多步推理的语言任务来说可能过于严格。事实上，许多语言任务受益于迭代生成，其中部分完成的中间状态为后续精炼提供了上下文（Ghazvininejad 等人，2019 (https://arxiv.org/html/2606.27617#bib.bib19); Nye 等人，2021 (https://arxiv.org/html/2606.27617#bib.bib32); Wei 等人，2022 (https://arxiv.org/html/2606.27617#bib.bib44)）。

为了解决这一局限性，我们提出了**掩码语言流模型**（MLFMs）。MLFMs将来自MDMs的掩码融入FLMs中，使用布朗桥作为随机插值法来连接嵌入的局部掩码序列与嵌入的干净序列。因此，MLFMs保留了MDMs的精确任意位置条件结构，同时学习一个耦合的连续流来生成掩码位置处的干净词元。此外，MLFMs的这种特殊结构为我们提供了训练和推理的自然策略。对于训练，我们提出将预训练的MDMs适配为MLFMs：这是可能的，因为在布朗桥的边界处，随机插值法简化为一个嵌入的局部掩码序列，与MDMs的设置相匹配。这允许我们从大型预训练MDMs启动MLFM训练，大大减少了训练模型所需的计算量。对于推理，我们利用MLFMs通过任意位置条件生成提供的灵活性。具体来说，我们引入了一种新颖的采样方案，该方案将一种新的无分类器指导方法与在线提升置信词元相结合。在这种指导方法下，干净词元在参考速度中被加噪以隔离其贡献。同时，如果特定词元位置的后验众数达到至少1−ε的概率，我们立即承诺使用该词元并替换为它的干净嵌入。通过提前提升这些已解决的词元，它们立即成为后续生成步骤的有用上下文，而不是在最后一步之前一直保持噪声状态。

通过实验，我们将一个来自Nie等人（2025a (https://arxiv.org/html/2606.27617#bib.bib30)）的具有1028M参数的预训练MDM适配为一个MLFM，并在GSM8K（Cobbe等人，2021 (https://arxiv.org/html/2606.27617#bib.bib12)）和MT-Bench（Zheng等人，2023 (https://arxiv.org/html/2606.27617#bib.bib48)）上评估了结果模型。MLFM在MT-Bench上优于SMDM和类似大小的AR基线，并在GSM8K上获得了令人鼓舞的性能，展示了其处理指令遵循和推理任务的能力。在这些基准测试中，我们还发现我们的新颖采样方案提供了显著的改进。据我们所知，这是基于流的语言模型首次被扩展到下游推理和指令遵循任务。

## 2 背景

在本节中，我们概述我们的方法所需的背景。第2.1节 (https://arxiv.org/html/2606.27617#S2.SS1)讨论掩码扩散模型，第2.2节 (https://arxiv.org/html/2606.27617#S2.SS2)讨论流语言模型，特别是LangFlow（Chen等人，2026 (https://arxiv.org/html/2606.27617#bib.bib10)），我们将其作为我们方法的基础。

##### 符号表示。** 在本文中，让V表示词元词汇表，|V|表示词汇表大小，L表示序列长度。我们将V中的元素与自然数对应：{1，…，|V|}。一个（干净的）序列X∼p，其中p表示我们的数据分布，记作X=(X1，…，XL)∈VL。我们用s∈[0，1]表示掩码概率，用t∈[0，1]表示连续加噪时间。

### 2.1 掩码扩散模型

*掩码扩散模型*（MDMs）通过向词汇表V引入一个特殊的*掩码*词元[MASK]并定义一个连续时间马尔可夫链（CTMC）（Del Moral 和 Penev, 2017 (https://arxiv.org/html/2606.27617#bib.bib14)）将p输送到完全掩码序列上的狄拉克测度，从而构建一个非自回归生成模型。在掩码概率s∈[0,1]和X∼p下，CTMC的前向过程Xs∼ps|0(Xs|X)由以下采样程序给出：
Xsl={ Xl，如果Bl=0，[MASK]，如果Bl=1，其中Bl∼i.i.d. Bernoulli(s)。 (1)
我们令Ms={l:Xsl=[MASK]}表示Xs中掩码词元的位置，令Us={l:Xsl≠[MASK]}表示Xs中干净词元的位置。我们可以通过模拟CTMC的反向过程来生成p的近似样本。这涉及逐渐去掩码词元以揭示干净词元，这些干净词元随后在后续步骤中被固定。这依赖于访问干净词元在给定未掩码词元情况下的真实因子化后验pdata l(Xl|XsUs)（Zheng等人，2024 (https://arxiv.org/html/2606.27617#bib.bib47); Ou等人，2025 (https://arxiv.org/html/2606.27617#bib.bib33)），其中XsUs表示Xs中由干净词元组成的子集合。这可以通过以下基于似然的目标从数据中学习：
LMDM(θ)=EX∼p，Xs∼ps|0(·|X)，s∼U(0，1)[1s∑l∈Ms−log pθ l(Xl∣XsUs)]，
其中我们用神经网络参数化pθ l以逼近pdata l，U(0,1)表示[0,1]上的均匀分布。在推理过程中，MDMs可以在一单步中并行去掩码多个词元。当给定当前掩码状态时，由于pθ l的因子化结构，要解掩码的词元接近条件独立时，这种方法最有效。然而，这通常只对一小部分词元成立，因为增加要解掩码的词元数量会增加词元之间强依赖的可能性。因此，尽管MDMs具有并行解码结构，但它们仍然需要许多精炼步骤才能获得高质量样本。

### 2.2 流语言模型

我们考虑*LangFlow*（Chen等人，2026 (https://arxiv.org/html/2606.27617#bib.bib10)）方法来实现*流语言模型*（FLMs），该方法在学习的嵌入空间而不是one-hot嵌入空间中操作。令H为嵌入维度。对于词汇表中的每个词元a∈V，我们将其对应的嵌入向量Ea∈RH表示该词元，并令E∈R|V|×H表示整体的嵌入矩阵。此外，对于序列X∼p，我们通过嵌入在欧几里得空间中表示X：
z1=(EX1，…，EXL)∈RL×H。
有了这种离散序列的连续表示z1，我们可以通过考虑基于高斯的随机插值法来构建p的生成模型：
zt=t z1+(1−t)ε，
其中ε∼N(0，IL×H)，t∈[0，1]。我们注意到z0是高斯噪声，z1是来自p的嵌入序列。为了从这个随机插值法诱导的流中采样，我们参数化去噪器pφ，它近似干净词元上的因子化后验：
pφ(·∣zt，t)∈Δ(V)L⊂RL×|V|。
这通过以下交叉熵目标进行训练：
LFLM(φ)=Et∼πt，X，ε[1L∑l=1L−log pφ l(Xl∣zt，t)]，
其中πt表示[0,1]上的某个分布。尽管监督目标Xl是离散的，但学习到的去噪器会诱导出一个连续的端点估计：
ẑ1，φ(zt，t):=Pφ(zt，t)E， 其中 Pφ(zt，t)=pφ(·∣zt，t)∈RL×|V|， (5)
并且ẑ1，φ(zt，t)估计后验均值嵌入E[z1∣zt，t]。此外，ẑ1，φ(zt，t)允许我们估计流的速度场：
vφ(zt，t):=ẑ1，φ(zt，t)−zt1−t， t<1。 (6)
因此，我们可以通过积分ODE来生成p的近似样本：
ddt zt=vφ(zt，t)。
在实践中，ODE在有限时间网格上离散化，最终词元序列通过从词汇空间中的pφ解码去噪器概率获得。连续流很有吸引力，因为它们定义了从噪声到干净序列的联合轨迹。这个轨迹可以被蒸馏成一个直接的流映射或一个一致性风格的求解器（Song等人，2023 (https://arxiv.org/html/2606.27617#bib.bib42); Boffi等人，2026 (https://arxiv.org/html/2606.27617#bib.bib5)），从而减少生成步骤的数量。然而，精确的一步映射对于语言模型来说是一个过于苛刻的对象——即某些问题（如数学或编程）可能自然需要迭代推理，其中中间承诺被用作后续决策的上下文。

## 3 掩码语言流模型

我们在第3.1节 (https://arxiv.org/html/2606.27617#S3.SS1)中介绍MLFM的核心框架。在第3.2节 (https://arxiv.org/html/2606.27617#S3.SS2)中，我们展示了预训练MDMs如何自然地适配到MLFMs，包括这种适配所需的结构变化。最后，第3.3节 (https://arxiv.org/html/2606.27617#S3.SS3)描述了如何将监督微调应用于指令-响应数据的MLFMs，据我们所知，这一方面在先前关于FLMs的工作中尚未被探索。

### 3.1 框架

我们保留来自MDMs和FLMs的掩码词元和词元嵌入E的使用。我们首先在欧几里得空间中定义一个前向过程zs，t∈RL×H。这个过程由(s，t)∈[0，1]2参数化。