条件掩码语言模型的并行解码，这个方向我在做文本生成加速的时候踩过不少坑。最早接触自回归解码那套东西，生成一段两百字的文本要等好几秒，用户端体验很差。后来看到Mask-Predict这篇工作，思路确实让人眼前一亮——它把生成任务从"从左到右一个字一个字蹦"变成了"先铺满整句再反复填空修正"，推理步数直接从O(n)降到O(log n)级别。这篇博文我就把这个方法的来龙去脉、核心机制、实操细节和踩坑经验完整梳理一遍，适合做文本生成加速、机器翻译、文本摘要的同行参考，也适合对非自回归生成感兴趣的朋友入门。

1. 为什么需要并行解码：自回归的瓶颈与非自回归的困境

1.1 自回归解码到底慢在哪里

先说说我们为什么要折腾并行解码这件事。标准的自回归（Autoregressive, AR）生成，核心逻辑是P(y_t | y_<t, x)，也就是第t个词的概率依赖于前面所有已经生成的词。这个设计天然串行——你必须先算出第1个词，才能算第2个，再算第3个，一路推到结束。生成长度为n的序列，就需要n次前向传播。

我实测过一个典型的Transformer翻译模型，在单张消费级显卡上，生成一条30词的译文大概需要300到500毫秒。如果做批量推理，batch size拉大能摊薄一些，但延迟的绝对值还是摆在那里。对于实时字幕、同声传译、对话系统这类场景，几百毫秒的延迟已经很要命了。更麻烦的是，序列越长，延迟线性增长，长文本生成基本没法做实时。

自回归的另一个隐性成本是显存带宽。每次前向传播都要把整个模型的权重从显存读一遍，而实际计算量很小，GPU的算力根本吃不满，瓶颈卡在访存上。这就是为什么AR解码的GPU利用率经常只有百分之十几。

1.2 非自回归的诱惑与硬伤

既然串行是瓶颈，那最直接的想法就是：能不能一次性把所有位置的词都预测出来？这就是非自回归（Non-Autoregressive, NAR）生成的出发点。一次前向传播，所有位置的输出同时出来，推理步数从n降到1，理论上能快几十倍。

但天下没有免费的午餐。NAR最大的问题是条件独立性假设——它假设各个位置的词互相独立，P(y|x) = ∏P(y_t|x)。这个假设在真实语言里根本不成立。"我吃饭"和"我饭吃"词都一样，但语序错了意思就变了。独立假设下，模型没法建模词与词之间的依赖，生成结果经常出现重复词、漏词、语序混乱。

早期的NAR工作靠引入隐变量、迭代精炼、CTC损失等手段来缓解，但效果始终和AR有差距。核心矛盾在于：要么快但质量差，要么质量好但慢。Mask-Predict的价值就在于，它在这条光谱上找到了一个很聪明的中间点。

1.3 Mask-Predict的定位：迭代精炼的折中路线

Mask-Predict的本质是迭代精炼（Iterative Refinement）。它不追求一步到位，也不接受逐词串行，而是用固定的T次迭代，每次迭代并行地预测所有被mask的位置，然后保留置信度高的词，把置信度低的重新mask掉，下一轮再猜。

这个思路的关键洞察是：生成一个词比判断一个词对不对要难。第一轮模型可能猜得不准，但它对"哪些位置猜得靠谱"是有判断力的。把靠谱的固定下来，不靠谱的交给下一轮，随着迭代进行，确定的位置越来越多，上下文越来越完整，剩下的词也就越来越好猜。

从计算复杂度看，AR是O(n)步，NAR是O(1)步，Mask-Predict是O(T)步，T通常取log2(n)左右。生成100个词，AR要100步，Mask-Predict大概7到10步就够了，加速比在10倍量级。而且质量能逼近甚至在某些任务上追平AR，这是它最吸引人的地方。

2. Mask-Predict的核心机制拆解

2.1 条件掩码语言模型的训练目标

Mask-Predict的底座是一个条件掩码语言模型（Conditional Masked Language Model, CMLM）。训练方式和BERT的掩码预测很像，但多了条件输入。

具体来说，给定源序列x和目标序列y，训练时随机mask掉y中一部分词，让模型根据x和未被mask的y来预测被mask的词。这里有个关键细节：mask的比例是随机采样的，而不是固定的。论文里用的是均匀采样，每个词被mask的概率独立，这样模型能见到各种mask密度的样本，推理时无论mask多少都能适应。

损失函数就是标准的交叉熵，只计算被mask位置的损失：

L = -Σ log P(y_t | x, y_observed)

其中求和只遍历被mask的位置。这个设计让模型学会"根据可见上下文填空"，而不是"从头生成"。

注意：训练时的mask策略直接影响推理表现。如果训练时mask比例总是很低，模型没见过高mask率的场景，推理第一轮全mask就会崩。所以mask比例的随机化范围要覆盖推理时的实际需求。

2.2 推理阶段的迭代解码流程

推理阶段是Mask-Predict的精华所在。整个流程分三步循环：

第一步，初始预测。把所有目标位置都mask掉，输入全是[MASK]，让模型一次性预测所有位置的词。这一步相当于NAR，质量一般，但给了个初始草稿。

第二步，计算置信度并筛选。对每个位置，取模型输出的概率分布，用最大概率值作为该位置的置信度分数。然后按置信度排序，保留置信度最高的那部分词，其余位置重新mask。

第三步，重新预测被mask的位置。把上一步保留的词作为已知上下文，重新预测被mask的位置。重复第二步和第三步，直到达到预设的迭代次数。

迭代次数T和每轮保留的词数是两个关键超参数。论文里用的策略是：第t轮保留n * (T-t+1)/T个词，也就是保留数量随迭代递减，越往后保留越多。这样早期大胆探索，后期精细收敛。

2.3 置信度打分与mask策略的细节

置信度怎么算，这里有个容易忽略的点。直接用softmax最大概率值是可以的，但论文里还试过用负熵作为置信度，效果差不多。我自己的经验是，最大概率值简单够用，但在某些任务上，用概率的几何平均或者考虑top-k的累积概率会更稳。

mask策略也有讲究。最朴素的是按置信度全局排序，保留top-k。但这样可能导致某些位置连续多轮都被mask，而有些位置一轮就固定。另一种做法是按位置分块，保证每个区域都有词被保留，避免局部信息缺失。我在做长文本摘要时发现，全局排序容易让句子开头和结尾先固定，中间一大片反复猜，分块策略能缓解这个问题。

还有一个细节是长度预测。Mask-Predict需要预先知道目标序列长度，因为要初始化对应数量的[MASK]。长度预测本身是个小任务，可以单独训一个分类器，也可以用源序列长度加个偏移量估计。长度预测错了，整个生成就偏了，这是NAR类方法的通病。

3. 实操复现：从零搭建一个Mask-Predict解码器

3.1 环境准备与依赖选型

我用的技术栈是PyTorch + HuggingFace Transformers，这套组合最省事。PyTorch版本建议1.10以上，Transformers用4.x任意较新版本都行。硬件方面，单张16G显存的卡就能跑中等规模实验，做小规模验证8G也够。

核心依赖：

pip install torch transformers datasets tokenizers

如果你要自己训模型，还需要准备数据加载和训练循环的代码。我建议先用预训练的掩码语言模型（比如多语言BERT类）做初始化，再在目标任务上微调，比从头训收敛快很多。

提示：不要直接用标准的因果语言模型（如GPT类）来改，因为它们的注意力是单向的，和CMLM的双向注意力不兼容。要用双向编码器架构。

3.2 模型结构的关键改动

标准Transformer编码器-解码器结构需要改两处：

第一，解码器注意力改成双向。自回归解码器的自注意力是带因果mask的，只能看左边。CMLM需要看两边，所以要把因果mask去掉，变成全连接注意力。这个改动在HuggingFace里可以通过设置is_decoder=True但add_cross_attention=True，然后手动去掉causal mask来实现。

第二，位置编码要支持任意长度。因为推理时是并行预测所有位置，位置编码必须能覆盖整个目标长度。用可学习的位置嵌入或者正弦位置编码都行，但要注意长度预测的最大值要留够余量。

下面是一个简化的解码器层实现示意：

import torch import torch.nn as nn class CMLMDecoderLayer(nn.Module): def __init__(self, d_model, nhead, dim_feedforward): super().__init__() self.self_attn = nn.MultiheadAttention(d_model, nhead, batch_first=True) self.cross_attn = nn.MultiheadAttention(d_model, nhead, batch_first=True) self.ffn = nn.Sequential( nn.Linear(d_model, dim_feedforward), nn.ReLU(), nn.Linear(dim_feedforward, d_model) ) self.norm1 = nn.LayerNorm(d_model) self.norm2 = nn.LayerNorm(d_model) self.norm3 = nn.LayerNorm(d_model) def forward(self, tgt, memory, tgt_key_padding_mask=None): # 双向自注意力，不加causal mask tgt2, _ = self.self_attn(tgt, tgt, tgt, key_padding_mask=tgt_key_padding_mask) tgt = self.norm1(tgt + tgt2) # 交叉注意力，看源序列 tgt2, _ = self.cross_attn(tgt, memory, memory) tgt = self.norm2(tgt + tgt2) # 前馈 tgt2 = self.ffn(tgt) tgt = self.norm3(tgt + tgt2) return tgt

注意self_attn调用时没有传attn_mask，这就是双向的关键。如果传了上三角mask，就退化成自回归了。

3.3 迭代解码的完整实现

解码循环是整个方法的核心，我把它拆成几个函数来写，逻辑更清晰。

def mask_predict_decode(model, src, src_mask, max_len, num_iterations=10): batch_size = src.size(0) device = src.device # 初始：全部mask tgt = torch.full((batch_size, max_len), MASK_TOKEN_ID, device=device) # 记录每个位置是否已经确定 fixed = torch.zeros(batch_size, max_len, dtype=torch.bool, device=device) for it in range(num_iterations): # 前向传播 logits = model(src, src_mask, tgt) probs = torch.softmax(logits, dim=-1) max_probs, pred_tokens = probs.max(dim=-1) # 计算本轮要保留的数量 # 保留数随迭代递增：第it轮保留 max_len * (it+1) / num_iterations num_to_keep = int(max_len * (it + 1) / num_iterations) # 对每个样本，按置信度选top-k for b in range(batch_size): # 已经固定的位置不参与 scores = max_probs[b].clone() scores[fixed[b]] = float('inf') # 已固定的强制保留 # 选置信度最高的num_to_keep个位置 _, topk_idx = scores.topk(num_to_keep) new_fixed = torch.zeros(max_len, dtype=torch.bool, device=device) new_fixed[topk_idx] = True # 更新tgt：固定的位置用预测词，其余保持MASK tgt[b][new_fixed] = pred_tokens[b][new_fixed] tgt[b][~new_fixed] = MASK_TOKEN_ID fixed[b] = new_fixed # 如果所有位置都固定了，提前结束 if fixed.all(): break return tgt

这段代码有几个实操要点。scores[fixed[b]] = float('inf')这行保证已经确定的词不会被重新mask，这是单调收敛的关键。num_to_keep的递增策略决定了收敛速度，我试过线性递增和指数递增，线性更稳，指数在长序列上收敛快但容易早熟。

注意：topk操作在batch内是逐样本做的，如果batch很大，这个Python循环会成为瓶颈。实际部署时可以用向量化的方式，把所有样本的置信度拼在一起做批量topk，或者用torch.topk的batch版本。

3.4 长度预测模块的接入

长度预测我单独训了一个小的回归/分类头。输入是源序列的编码表示（取平均或CLS位置），输出是目标长度。训练数据直接用平行语料的真实长度。

class LengthPredictor(nn.Module): def __init__(self, d_model, max_len): super().__init__() self.fc = nn.Sequential( nn.Linear(d_model, d_model // 2), nn.ReLU(), nn.Linear(d_model // 2, max_len) ) def forward(self, encoder_output, src_mask): # 对源序列做masked mean pooling mask = src_mask.unsqueeze(-1).float() pooled = (encoder_output * mask).sum(1) / mask.sum(1) logits = self.fc(pooled) return logits # 用交叉熵训练，预测长度类别

长度预测的准确率直接影响最终质量。我实测下来，长度误差在±2以内时，BLEU掉点不明显；误差超过5，质量断崖式下跌。所以长度预测器要单独调好，别指望主模型能兜底。

4. 调参经验与性能优化实录

4.1 迭代次数与保留策略的权衡

迭代次数T是最重要的超参数。T太小，模型没机会修正错误，质量差；T太大，加速比下降，而且后期迭代收益递减。

我的经验是T取ceil(log2(n))到ceil(log2(n)) + 2之间比较合适。n=50时T取7到8，n=100时T取8到9。再往上加，BLEU提升不到0.5，但延迟增加明显。

保留策略我试过三种：

线性递增是最稳的，我大部分实验都用这个。指数递增在长度小于30的短序列上能省1到2轮迭代，但长序列上容易把错误的词过早固定，后面改不回来。

4.2 置信度校准的坑

模型输出的softmax概率往往过度自信，这是神经网络的通病。直接拿最大概率当置信度，会导致一些实际错误的词因为概率高而被过早固定。

我试过两种校准方法。一种是温度缩放，在验证集上拟合一个温度参数T，把logits除以T再softmax，让概率分布更平滑。另一种是用多个轮次的预测一致性作为置信度——如果一个词连续两轮都被预测成同一个，说明模型比较确定。

温度缩放实现简单，效果稳定，我一般都会加上。代码就一行：

logits = logits / temperature # temperature > 1，通常1.5到2.5 probs = torch.softmax(logits, dim=-1)

一致性判断需要多存一份上一轮的预测，显存开销小，但逻辑稍复杂。两者可以结合用。

4.3 批处理与显存优化

Mask-Predict的批处理比AR复杂，因为不同样本的迭代进度可能不同。最朴素的做法是等所有样本都跑完T轮再返回，但这样短序列样本会陪跑。

我用的优化是动态提前退出：每个样本维护一个finished标志，所有位置都固定后就标记完成，后续迭代跳过该样本的计算。实现上可以用mask把已完成样本的loss和梯度置零，或者干脆把它们从batch里摘出来。

显存方面，因为要同时保存所有位置的logits和概率，显存占用比AR高。生成长度512时，batch size要相应调小。可以用梯度检查点和混合精度来省显存，我实测混合精度能省40%左右，速度还快一点。

提示：推理时不需要梯度，记得用torch.no_grad()包住整个解码循环，否则显存会爆得莫名其妙。

5. 常见问题排查与避坑指南

5.1 生成结果重复、漏词怎么办

这是NAR类方法的经典问题。重复通常是因为模型在某些位置反复预测同一个高频词，漏词是因为长度预测偏短或者某些位置一直没被固定。

排查思路：先看长度预测是否准确，如果长度偏短，漏词是必然的。再看置信度分布，如果大量位置置信度都很低，说明模型没学好，要检查训练mask策略。如果只是个别位置反复，可以调低保留阈值，让这些位置多迭代几轮。

我遇到过一次严重的重复问题，最后发现是训练时mask比例设得太低（平均只有15%），模型没见过高mask率场景。把mask比例范围调到10%到90%均匀采样后，问题基本消失。

5.2 迭代不收敛或震荡

正常情况下，每轮迭代固定的词应该单调增加，生成质量单调提升。如果出现震荡——某些位置这轮固定了下轮又被改——说明保留策略有问题。

最常见的原因是没有强制保留已固定的词。上面代码里scores[fixed[b]] = float('inf')这行如果漏了，已固定的词可能因为新一轮置信度低而被重新mask，导致来回震荡。另一个原因是学习率太大，模型输出不稳定，这个在微调阶段要注意。

5.3 加速比不达预期

理论上T轮迭代应该比AR快n/T倍，但实际测下来经常只有3到5倍。原因有几个：

• Python循环开销：逐样本的topk和更新操作在Python层做，batch大时开销显著。解决方法是向量化。
• 长度预测和预处理：这些额外步骤吃掉了部分加速收益。
• GPU利用率：并行解码的计算模式对GPU不如AR友好，kernel启动开销占比高。

我的优化顺序是：先向量化topk操作，再把长度预测和主模型合并到一次前向，最后考虑用CUDA graph减少kernel启动开销。做完这些，加速比能从3倍提到8倍左右。

5.4 常见问题速查表

6. 这套方法适合什么场景，不适合什么场景

Mask-Predict不是万能的，用对场景才能发挥价值。

适合的场景：对延迟敏感、对质量要求不是极致、序列长度中等的任务。比如实时翻译、语音识别的文本后处理、对话系统的回复生成。这些场景下，10倍加速带来的体验提升，远大于BLEU掉的那一两个点。

不太适合的场景：对质量要求极高的任务，比如法律文书翻译、医学报告生成，这些场景宁可慢也要准，AR还是首选。另外，超长序列（比如生成长度超过1000）时，迭代次数虽然是对数增长，但每轮的计算量还是O(n^2)的注意力，总开销依然很大，这时候要考虑稀疏注意力或者分块生成。

还有一个容易被忽略的点：Mask-Predict对模型容量的要求比AR高。因为要并行预测所有位置，模型需要更强的全局建模能力。小模型上Mask-Predict和AR的差距会拉大，大模型上差距缩小。所以如果你的模型参数量很小，可能得不偿失。

我在实际项目里的体会是，Mask-Predict最适合作为质量与速度的调节旋钮。通过调整迭代次数T，可以在推理时动态权衡。T=1就是纯NAR，最快但质量最差；T=log2(n)是标准配置；T再往上逼近AR质量。这种灵活性是纯AR或纯NAR给不了的。

最后分享一个小技巧：如果做的是多语言或者多任务场景，可以把任务标识拼在源序列前面，让同一个CMLM处理多个任务。我试过翻译和摘要共用一个模型，只要任务标识区分清楚，效果不比单独训差，还省了部署成本。这个思路后续还可以扩展到更多生成任务上，值得一试。