通过随机采样令牌子集对齐,提升扩散模型训练效率与稳定性。
MaskAlign: Token-Subset Representation Alignment for Efficient Diffusion Training

- 采用随机令牌子集进行表征对齐,缓解噪声与清晰图像间的不匹配问题。
- 在ImageNet上训练时,收敛速度提升30%,生成质量保持不变。
- 适合追求高效扩散模型训练的科研与工程人员使用。
利用预训练视觉模型进行表征对齐,近年来在加速扩散变换器训练方面展现出强大潜力。通过将扩散模型中间特征与自监督视觉编码器提取的清晰图像表征对齐,现有方法提升了收敛速度与生成质量。然而,这种对齐也带来显著约束:扩散模型处理的是随时间步变化的噪声输入,而参考特征来自清晰图像。本文从令牌级别重新审视这一错配问题。我们发现,在全令牌对齐下,梯度范数较大的令牌表现出稳定的空间偏好,表明对齐目标并非均匀作用于所有令牌,可能促使模型依赖完整的清晰图像令牌集。为解决此问题,我们提出MaskAlign,一种在训练中随机采样令牌子集进行对齐的方法。通过迭代暴露模型于不同子集,减少对完整令牌集的依赖,并增强在令牌子集扰动下的对齐稳定性。为进一步缓解直接丢弃令牌带来的信息损失,我们引入轻量级预掩码令牌混合模块,在掩码前实现令牌间信息共享。
原文摘要 · Abstract (English)
Representation alignment with pretrained vision models has recently shown strong potential for accelerating diffusion transformer training. By aligning intermediate diffusion features with clean-image representations from self-supervised vision encoders, existing methods improve convergence and generation quality. However, such alignment also introduces a non-trivial constraint: diffusion models operate on noisy inputs whose usable information varies across timesteps, while the reference features are extracted from clean images. In this paper, we revisit this mismatch from a token-level perspective. We find that, under full-token representation alignment, tokens with large alignment-gradient norms exhibit a stable spatial preference, suggesting that the alignment objective does not affect all tokens uniformly and may encourage the model to rely on the complete set of clean-image tokens. To address this issue, we propose MaskAlign, a token-subset representation alignment method that applies alignment to randomly sampled token subsets during training. By exposing the model to different token subsets across iterations, MaskAlign reduces the dependence of representation alignment on the complete token set and encourages alignment behavior that is more stable under token-subset perturbations. To mitigate the information loss caused by directly dropping tokens, we further introduce a lightweight pre-mask token mixing block that shares information across tokens before masking.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。