一句话总结

来自阿姆斯特丹大学和Meta FAIR的研究人员提出了一种上下文加权的离散流匹配方法，该方法利用局部上下文信息增强连续时间马尔可夫链（CTMC），在OpenWebText上将生成困惑度降低高达63%，可与强大的半自回归块扩散基线模型相媲美，同时保持任意顺序生成，并证明简单的上下文感知修改显著提升了采样和训练效率。

核心贡献

• 本文通过实验证明，离散流匹配中token级别的预测熵与未掩码邻居的密度密切相关，上下文越密集，不确定性越低。
• 论文引入了一种对CTMC转移动态的上下文加权修改，该修改根据局部邻域密度调整更新概率，同时保持每个token的边缘去掩码概率，将中间步骤引导至条件更好的token。
• 推导出两种实用的上下文感知机制：一种无需微调即可改善生成质量的推理时采样器，以及一种重新加权训练token的缩放交叉熵目标函数，在OpenWebText上将生成困惑度降低高达63%，MAUVE提升24%，有效和新型分子样本分别增加最多2.8倍和1.9倍，同时匹配半自回归基线并保持任意顺序生成。

引言

针对文本和分子等离散数据的生成模型越来越多地采用离散流匹配和扩散方法，这些方法学习连续时间马尔可夫过程（CTMC）以并行和任意顺序生成token。与以固定前缀为条件的自回归方法不同，这些方法对所有可能的token顺序进行边缘化，因此模型必须根据各种部分观测来预测token。这导致预测难度高度可变：一些token受到周围上下文的强烈约束，而另一些则保持模糊。先前的工作试图通过强制结构化生成顺序或在推理时调整采样器来处理这种不平衡，但没有将预测不确定性直接集成到生成动态中。

作者观察到，预测难度与token局部邻域中已去掩码的token数量强相关。他们修改CTMC转移概率，使其依赖于局部上下文，同时保持沿概率路径每个token的边际去掩码概率。这产生了两种实用机制：一种可应用于预训练模型而无需微调的上下文加权采样器，以及一种基于局部上下文重新加权token的缩放交叉熵训练目标。该方法在多个领域提升了生成质量，将文本生成分数提高最多24%，并大幅提高分子生成的有效性和新颖性，同时匹配块扩散基线的性能，且不牺牲任意顺序生成。

方法

作者基于离散流匹配（DFM）构建方法，该方法使用连续时间马尔可夫链（CTMC）对有限离散序列进行建模。在此框架中，在源噪声样本和数据样本之间定义条件概率路径。转移核依赖于分解速度，即每次转换改变一个坐标。这种分解实现了高效的逐坐标模拟。标准训练目标是分解的条件匹配目标，它分解为每个坐标的预测问题，其中所有token权重均等。

然而，作者观察到均匀加权是次优的，因为可降低的学习信号因token而异。他们通过实验发现，局部上下文可作为预测难度的简单代理。具体而言，token周围的未掩码邻居数量与预测熵和真实token的负对数似然都强相关。

如下图所示：

随着更多局部上下文可用，两个指标都持续下降，表明附近观察到的token使预测更确定、更准确。这种变化很重要，因为模型预测器的预期损失分解为不可降低的不确定性项和可降低的模型误差项。对于高熵坐标，很大一部分损失来自内在模糊性，提供的训练信号信息量较少。低熵坐标对应更确定的目标，提供更清晰的监督。理想情况下，加权方案应强调损失的可降低部分，而不是均匀加权原始每个token的损失。

为此，作者提出对标准CTMC过程进行上下文感知修改，称为上下文加权离散流匹配。他们引入一个上下文权重函数 $\alpha$α，将当前状态映射为每个token位置的正权重。该函数通过定义一个新的分解速度集成到CTMC中，该速度将标准概率速度按上下文权重缩放。这个新速度生成一条上下文加权的条件概率路径。为了确保路径连接噪声分布和数据分布，并保持标准凸路径的边际分布，作者对上下文权重函数设定了特定条件。这些条件要求权重在当前掩码坐标上平均为1，并且在循环移位下具有等变性。

在推理时，该公式允许对采样过程进行纯推理时修改。作者将乘法权重注入欧拉求解器的每个坐标，得到上下文加权更新规则。他们引入了两种具体的上下文加权求解器：邻居加权求解器，使用局部窗口内未掩码token的数量；以及熵加权求解器，使用测试时可用的token级预测熵。在这两种情况下，权重都经过归一化，使得局部上下文良好或低熵位置获得更大的权重。

对于训练，由于对于一般上下文权重，上下文加权路径没有解析形式，因此无法进行免模拟采样。作者没有依赖昂贵的时序模拟，而是利用已知的未掩码token数量分布，构建了一个具有线性复杂度的精确token级采样器。他们首先采样未掩码token的数量，然后顺序更新token，从而无需时间轴模拟，同时从上下文加权路径返回精确样本。

最后，基于预测不确定性与可用局部上下文之间的经验联系，作者引入缩放交叉熵作为标准目标的即插即用替代方案。该损失函数使用上下文权重函数重新加权每个坐标的贡献，提高条件更好坐标的权重，降低模糊坐标的权重。这种方法保持训练时路径采样不变，保留了标准训练和推理时任意顺序采样的效率，同时有效捕捉不同token更新的上下文加权重要性。

实验

该方法在OpenWebText文本生成和QM9分子生成上使用离散流匹配架构进行评估。推理时上下文加权采样，特别是邻居加权更新，提高了生成质量，在低数据情况下尤其有效，而训练时缩放交叉熵使模型接近强大的半自回归基线。使用上下文加权路径进行训练进一步有利于小数据设置，并且中等邻域半径提供了最佳的质量-多样性权衡。

带有缩放交叉熵的上下文加权离散流匹配将MAUVE从0.628提升至0.777，并将生成困惑度降低63%，使其接近强大的半自回归块扩散基线。它明显优于非自回归基线SEDD和MDLM，同时保留了任意顺序生成的灵活性。缩放交叉熵将MAUVE提高到0.777，接近L'=4的BD3LM（0.784）并超过L'=8（0.734），远远领先于SEDD（0.574）和MDLM（0.685）。生成困惑度从82.17降至30.20，降幅63%，介于BD3LM L'=8（29.47）和L'=16（32.95）之间，远优于SEDD（110.23）和MDLM（41.73）。

缩放交叉熵（SCE）相比标准交叉熵、NELBO和Bregman散度，大幅提升了离散流模型的生成质量。在均匀源条件下的OpenWebText基准上，SCE提高了MAUVE，并将生成困惑度降低约三分之二，同时熵保持接近基线水平。对于QM9上的分子生成，SCE在比较的损失函数中产生了最多的有效分子，并在新颖性上保持竞争力。在均匀源下，SCE在OpenWebText上将MAUVE从0.628提高到0.777，并将生成困惑度从82.2降至30.2。SCE下的熵（7.66）接近交叉熵水平（7.93），表明质量提升并非以多样性崩溃为代价。在QM9上，SCE产生了556个有效分子，远多于标准交叉熵的475个，而新颖性保持相似（297.6 vs. 287.0）。NELBO和Bregman损失相比CE仅提供中等改进，而SCE在MAUVE和困惑度上都有更大的跃升。

从混合路径切换到邻居加权路径，在OWT和QM9的掩码源设置下都带来了明显改进，而对均匀源在OWT上几乎没有益处。在QM9上，邻居路径持续提高有效性和新颖性，在掩码源下有效分子数几乎翻倍。对于OWT的均匀源，邻居路径没有改善MAUVE或生成困惑度，仅略微增加熵。对于OWT的掩码源，邻居路径将MAUVE从0.690提高到0.751，并降低生成困惑度。在QM9上，邻居路径在两种源类型下均提高了有效性和新颖性，掩码源有效性从177近乎翻倍到298。邻居路径在掩码源设置下影响最大，这与它在低数据设置下的增强有效性一致。

评估涵盖OpenWebText文本生成和QM9分子生成，考察了使用不同损失函数和源分布的离散流匹配模型。缩放交叉熵显著优于标准交叉熵、NELBO和Bregman损失，在生成质量上取得重大提升，同时不牺牲多样性。邻居加权掩码路径在掩码源条件下进一步提高了生成质量，特别是在低数据分子设置中，但没有增强均匀源性能。