只需对齐 $\bm{x}$：对齐预测，而不是表示

表示对齐已成为加速扩散训练的有效方法，但其优点并不能可靠地转移到像素空间干净图像预测。在 JiT 中，我们发现辅助特征对齐可以改善对语义特征的访问，同时减少对干净图像预测所需的图像变化的访问，从而在辅助目标和去噪任务之间造成不匹配。这提出了一个不同的原则：辅助监督应该改进预测目标本身，而不是强加一个单独的表示目标。我们引入了 JAx（Just Align x），这是一种预测监督方法，可以跨噪声级别对齐干净图像预测。 JAx 通过保留原始 JiT 输入分布的马尔可夫降级，将嘈杂的学生观察与更清晰的观察结合起来。在这种耦合下，来自更清洁状态的预言机预测具有与最优 JiT 目标相同的条件均值，而其条件目标协方差并不更大。因此，预言机预测对齐将群体 JiT 目标保持在一个常数，同时提供较低方差的训练目标。为了使这种构造在不完美的 EMA 教师中实用，JAx 将地面实况监督与可靠性门控耦合带结合起来，该耦合带根据预测风险选择附近的教师状态。在 ImageNet 256x256 上，JAx 持续改进 FID 并加速 JiT-B/16、L/16 和 H/16 的收敛，无需外部编码器或更改架构或采样程序。梯度诊断进一步显示小批量梯度方差减少，而消融则表明增益不能仅通过时间重新加权来解释。这些结果表明，预测空间监督为像素空间生成模型的表示对齐提供了一种简单且有原则的替代方案。