知道何时持有它们：统一状态扩散语言模型中的正确标记保留

统一状态扩散模型（USDM）可以在任何去噪步骤修改任何标记，这让它们能够纠正自己的错误，这是相对于掩模扩散的一个关键优势。然而，自我纠正需要修改不正确的标记并保留正确的标记，我们表明当前的 USDM 缺乏后者。即使在贪婪尾解码下，最先进的 USDM（DUO、UDLM 和均匀噪声 SEDD）也会在每一步不断修改 512 个位置中的 173--270 个位置，而这些大型、不协调的编辑会破坏样本多样性。随机词元损坏实验将这种缺陷追溯到模型本身：它们以几乎相同的精度重建干净和损坏的词元，尽管干净的词元更容易成为目标。验证 NELBO 的分解表明，训练几乎不会奖励保留：错误的预测在损坏的位置上会受到严重惩罚，但在干净的位置上几乎不受惩罚。我们提出了正确词元保留正则化（CTR-Reg），这是一种简单但有效的辅助损失，可以训练模型保留不受前向过程干扰的词元，并且不需要更改采样器。 CTR-Reg 在六个基准测试中平均将干净词元的准确性提高了 26.5 个百分点，同时损坏词元的准确性几乎保持不变，并且其每步修正仅收敛到 3--11 个位置。只需五个贪心尾步骤，在 CTR-Reg 下，所有三个模型的生成困惑度都减少了一半以上，同时保留了多样性，并且这些收益在整个抽样预算中都保持不变。我们的结果表明，正确的标记保留是自我纠正扩散语言模型的一个关键缺失成分，并证明了一种有效的修复方法。