SpeeDiff 用 TPR 损失打破扩散模型端到端训练禁区

我找到了一篇关于潜在扩散模型端到端训练的论文。和 @XiaoyuW9401147 讨论后，我们都认为，虽然现在在 ImageNet 上训练 DiT 感觉像个玩具实验，但这篇论文中的端到端训练仍然值得一看，因为它简单而优雅。

自从 REPA 发布以来，扩散社区的焦点已经从训练大规模 DiT 转向更高效地训练 DiT。由于 VAE 空间是为重建而设计的，缺乏生成能力，REPA-E 提出对 DiT 和 VAE 进行端到端训练，以获得更好的输入空间，从而加快收敛速度。其关键发现是，扩散损失不能反向传播到 VAE，因为这会导致潜在空间模式崩溃。RAE 革新了训练范式，证明可以直接在 Dinov2 或 SigLIP2 等表示空间上训练 DiT。虽然语义空间本身有利于生成，但它无法消除对两阶段训练（视觉编码器、生成器）的需求。

本文提出的 SpeeDiff 揭开了端到端训练潜在扩散的秘密。我们所需要的只是一个 TPR 损失，它使用 Tweedie 公式通过从中间噪声状态解码预测的干净潜在来计算像素级重建损失。有了 TPR 损失，我们可以使用扩散损失从头开始更新 VAE 编码器，推翻了 REPA-E 中的主张。

通过端到端训练的简洁公式，我们现在可以让生成模型自己决定什么是最好的生成输入空间。期待看到这是否适用于大规模文本到图像实验或视频生成。但是，从更广泛的角度思考，为什么社区需要潜在扩散模型的端到端训练？视觉编码器的存在意味着生成总是有损的。也许最终的解决方案是将噪声移到编码器之前：扩散发生在像素空间，VAE 成为去噪器的一部分，端到端训练决定在每个噪声级别压缩多少。

令人惊讶的是，这是去年的 CVPR 论文，四天前才在 arxiv 上发布🤯