MaskAlign：用于高效扩散训练的词元子集表示对齐

与预训练视觉模型的表示对齐最近显示出加速扩散 Transformer 训练的强大潜力。通过将中间扩散特征与来自自监督视觉编码器的干净图像表示对齐，现有方法提高了收敛和生成质量。然而，这种对齐也引入了一个重要的约束：扩散模型对噪声输入进行操作，其可用信息随时间步长而变化，而参考特征是从干净的图像中提取的。在本文中，我们从 词元级 的角度重新审视这种不匹配。我们发现，在全词元表示对齐下，具有大对齐梯度范数的词元表现出稳定的空间偏好，这表明对齐目标不会均匀地影响所有词元，并且可能会鼓励模型依赖于完整的干净图像词元集。为了解决这个问题，我们提出了 MaskAlign，一种词元子集表示对齐方法，该方法在训练期间将对齐应用于随机采样的词元子集。通过在迭代中将模型暴露给不同的标记子集，MaskAlign 减少了表示对齐对完整标记集的依赖性，并鼓励在标记子集扰动下更稳定的对齐行为。为了减轻直接丢弃词元造成的信息丢失，我们进一步引入了一个轻量级的预屏蔽词元混合块，该块在屏蔽之前在词元之间共享信息。