重掩码，而非替换：掩码扩散语言模型中的 Token-to-Mask 精修

# 掩码扩散语言模型中的 Token-to-Mask 精修  
代码已开源：https://github.com/synsis/remasked_DLM  
来源：https://arxiv.org/html/2604.18738  

## 重新掩码，而非替换：掩码扩散语言模型中的 Token-to-Mask 精修  
††代码已开源：https://github.com/synsis/remasked_DLM  

姚林¹²  
¹ 上海交通大学计算机科学学院，上海 200240  
² 中关村学院，北京 100097  
[email protected]  

###### 摘要  
掩码扩散语言模型（如 LLaDA2.1）依赖 Token-to-Token（T2T）编辑来自纠生成错误：一旦某个备选 token 的置信度超过阈值，当前 token 即被覆写。我们指出该规则存在三种结构性失效：  
1. 当没有任何单一备选足够置信时，触发器无法启动；  
2. 替换是在本身可能含错的上文下计算的；  
3. T2T 流在训练时使用的均匀扰动与推理时模型实际产生的语义合理错误分布不符。  

作为替代，我们提出 **Token-to-Mask（T2M）重掩码**：不再用新猜测覆写可疑 token，而是将该位置重置为掩码状态，让下一步去噪在分布内上下文重新预测。该方法无需训练，仅改动编辑规则，不引入新参数。我们配套三种检测启发式，并从理论上说明为何掩码是比错误 token 更好的条件信号。  

在 8 项基准上，T2M 在需要精确 token 级输出的任务中提升准确率；在 CMATH 上最大提升 **+5.92** 分，其中 79.9% 的基线错误属于“最后一公里损坏”（推理正确但答案 token 错乱），T2M 修复了 41.3% 的此类案例。  

## 1 引言  
离散掩码扩散语言模型（dLLM）从全掩码序列出发，通过迭代并行填词生成文本。LLaDA、Dream 等已证明其精度可与同规模自回归模型媲美。并行解码的代价是：同一步填充的 token 彼此独立，可能相互矛盾；一旦出错，后续步骤会在错误条件下继续扩散。  

LLaDA2.1 引入 T2T 编辑：每步 M2T 后，重新评估已提交 token，若任一备选概率超过阈值 τ_t2t 则覆写。该机制帮助 LLaDA2.1 追平自回归精度，但也带来三种失效：  

1. **检测-替换耦合**：当后验多峰（如 “sad”:0.12，“happy”:0.11…）时，无候选过阈值，明显错误被保留，称“修正惯性”。  
2. **上下文污染**：替换是在可能被污染的上文下计算的，错误会扩散并反噬自身。  
3. **训练-推理噪声分布失配**：训练时 T2T 流用均匀随机扰动，而推理错误语义合理、局部连贯。  

我们将检测与解耦：一旦检测器标记可疑 token，即重置为 [M]，让下一步 M2T 在净化后的上下文中联合重预测，称 **Token-to-Mask（T2M）重掩码**。掩码无语义偏向，可一次重掩多位置，推理时零参数、零再训练。  

## 2 相关工作  
#### 离散掩码扩散语言模型  
D3PM 引入吸收态离散扩散；MDLM、SEDD 给出高效训练目标；LLaDA 扩至 8B 参数；Dream 继续放大。LLaDA2.1 加入半自回归块生成与 T2T 编辑，架构与训练不变。  

#### 重掩码策略  
- **ReMDM**：对每一步已提交 token 以均匀概率 σ_t 随机重掩，好坏等比例抹除。  
- **CORE**：通过掩码扰动敏感度标记“上下文脆弱”token，需 O(k) 额外前向。  

本文 T2M 仅改“替换→重掩”，检测器不变，二者可组合。  

#### 基于训练的自纠正  
RemeDi、ProSeCo、PRISM、MDPO 等均需额外训练学习检测或纠正。T2M 与之正交，仅改动纠正动作。  

## 3 预备知识  
LLaDA2.1 采用半自回归块扩散：响应按块并行生成，块内迭代 M2T 去噪；每步后 T2T 编辑按  

x_i ← x_i^* if p_θ(x_i^*|z) > τ_t2t ∧ x_i^* ≠ x_i^old  

覆写，直至无掩码且无编辑。  

## 4 方法：Token-to-Mask 重掩码  
### 4.1 概述  
T2M 仅改编辑规则：对 T2T 会复查的位置 i，将动作由“替换”变为“重置”  

x_i ← [M] if ShouldRemask(i, z, θ)  

随后 M2T 在更新后的上下文中重预测。模型权重、M2T 填词规则、块调度、KV 缓存均不变。  

### 4.2 错误检测  
我们给出三种 ShouldRemask 实例：  

#### LowProb  
直接利用模型对当前 token 的概率：若 p_θ(x_i^old | z_{-i}) < τ_lp 则重掩，无需候选。  

#### T2T-Remask  
复用 T2T 的触发条件，仅把动作换成重掩，用于隔离“remask vs replace”效应。  

#### LogitDiff  
追踪连续两次迭代间模型对 x_i^old 的置信度变化：若 p_θ^(t-1) - p_θ^(t) > τ_ld 则重掩，捕捉轨迹级信号。  

### 4.3 安全上限  
防止重掩震荡：  
- 每位置最多重掩 C_max 次（默认 1）；  
- 每步重掩比例不超过 ρ_max（默认 0.25），超出时按置信度最低者优先。  

算法 1 给出完整 T2M 流程，直接替换 LLaDA2.1 的 T2T 编辑阶段，其余不变。