MaskAlign：用于高效扩散训练的令牌子集表示对齐 MaskAlign: Token-Subset Representation Alignment for Efficient Diffusion Training

MaskAlign的完整步骤包括：首先，从干净图像$z_*$得到干净潜在$z_0$，在时间步$t$构造噪声潜在$z_t = \alpha_t z_0 + \sigma_t \epsilon_z$。$z_t$被patch化和投影为patch令牌序列$x_{0,t} = \{x_{0,t,1}, ..., x_{0,t,N}\}$，拼接类别令牌$c_{0,t}$形成$H_{0,t} = [c_{0,t}, x_{0,t}]$。第二步，应用轻量级预掩码令牌混合块$M_\psi(\cdot)$共享令牌信息：$\bar{H}_{0,t} = M_\psi(H_{0,t}, t, y)$。第三步，对patch令牌采样二元保持掩码$m \in \{0, 1\}^N$，类别令牌总是保留：$H_{0,t}^{(m)} = [\bar{c}_{0,t}, \{\bar{x}_{0,t,i}\}_{i \in S(m)}]$。第四步，掩码序列输入SiT块，在层$\ell$产生$H^{[\ell]}_t(m)$。第五步，对保留的类别令牌和可见patch令牌计算预测损失和对齐损失，最终目标为$\mathcal{L}_{total} = \mathcal{L}_{pred} + \lambda\mathcal{L}_{REPA}$。

MaskAlign的技术新颖性体现在三个方面：一是首次从令牌级别分析表示对齐的训练行为，发现了梯度分布的空间偏好问题；二是提出随机令牌掩码策略进行表示对齐，与现有方法的对齐所有令牌有本质区别；三是引入预掩码令牌混合块来减少信息丢失，这种设计与仅使用掩码重建的方法（如MDT、MaskDiT）不同，MaskAlign使用随机掩码构造令牌子集进行表示对齐，预测和对齐损失在保留的类别令牌和可见patch令牌上计算。

独立分析的弱点包括：首先，固定掩码比例（25%）可能不是最优的，不同训练阶段或不同架构可能需要自适应掩码策略。其次，预掩码令牌混合块的设计相对简单，可以探索更复杂的混合机制如注意力或图结构。第三，实验仅在ImageNet上验证，缺乏在其他数据集（如COCO、LAION）和任务（如文本到图像、超分辨率）上的泛化性测试。第四，方法增加了参数量和计算复杂度，虽然减少了令牌数量，但混合块的引入可能抵消部分收益。改进方向包括设计自适应掩码策略（根据训练阶段或梯度动态调整掩码比例）、探索更高效的令牌混合机制、在更多数据集和任务上验证泛化性、优化混合块设计以减少参数和计算开销。

作者提出的未来工作方向包括探索自适应掩码策略和更广泛的模型族，以更好地理解令牌子集表示对齐的范围和鲁棒性。基于本文成果，可以延伸的研究方向包括：将MaskAlign应用于文本到图像生成（如Stable Diffusion架构）、更高分辨率生成（如512×512或1024×1024）、探索其他预训练视觉编码器（如CLIP、MAE）作为教师模型、研究掩码策略的调度（训练早期使用低掩码比例，后期增加掩码比例）、将令牌子集对齐思想应用于其他生成任务（如视频生成、3D生成）、分析MaskAlign在不同噪声水平下的行为差异、探索与 Curriculum Learning 结合的渐进式掩码策略。

复现评估方面，论文未提供代码仓库链接，但详细描述了实现细节，包括使用SiT-B/2和SiT-XL/2架构、固定批大小256、与REG相同的学习率和EMA设置。评估协议使用50K生成图像，SDE Euler-Maruyama求解器和250采样步骤。算力需求未明确说明，但从训练7M迭代和参数量（677M-732M）推测需要大量GPU资源。ImageNet数据集需要预先下载和预处理，包括中心裁剪和256×256调整大小，使用Stable Diffusion VAE编码为潜在表示。复现难度中等偏上，主要挑战在于计算资源和训练时间，但实现细节充分，算法描述清晰，预训练模型和评估代码可能需要自行实现。