重新思考 Vanilla Pixel Diffusion Transformers 的设计空间

像素空间扩散模型近期重新成为潜在扩散的有力替代方案，避免了有损的潜在空间瓶颈以及单独训练的自编码器。JiT 表明，通过将常用的速度预测（$\boldsymbol{u}$-prediction）替换为干净图像预测（$\boldsymbol{x_0}$-prediction），原生扩散 Transformer 可以直接在像素空间中生成图像。在本文中，我们重新审视了这一选择，并引入 MFlow：一种自适应参数化方法，通过数据的主方向学习在 $\boldsymbol{x_0}$-prediction 与 $\boldsymbol{u}$-prediction 之间进行插值，并通过扩散训练以端到端方式学习。我们发现，最优参数化方式会随数据表示和模型容量系统性地变化——尤其是，随着模型容量增加，模型会更倾向于速度预测。我们在 ImageNet、UCF-101 视频以及大规模文本到图像生成任务上评估了 MFlow，结果表明，在多种设置下，它都取得了具有竞争力或更优的生成性能。此外，我们发现，无论采用何种参数化方式，原生像素扩散 Transformer 相较于真实图像仍持续存在高频信息不足的问题。为解决这一问题，我们引入 ConvRefiner：一种轻量级卷积模块，在冻结扩散骨干网络的同时，使用对抗损失进行事后训练。仅需少量额外训练，ConvRefiner 即可恢复生成图像的高频统计特性，同时保持 FID 不变。项目网页：https://soumik-kanad.github.io/mflow-diffusion/。