SpeeDiff:可扩展的像素锚定端到端潜扩散模型
SpeeDiff: Scalable Pixel-Anchored End-to-End Latent Diffusion Model
把扩散损失梯度直接回传到 VAE 编码器做端到端联合训练,并用 Tweedie 公式从中间噪声态解码预测的干净潜变量来加像素级重建约束,避免潜空间坍缩;做视觉生成的人可以借此替换掉「预训练冻结 VAE + 两阶段」的常规流程,显著缩短训练时间并复用其 32x 压缩 VAE 与 Transformer 联合缩放策略。
原文依据 a simple yet effective Tweedie Pixel Reconstruction (TPR) loss, which provides additional pixel-level feedback by decoding a predicted clean latent from an intermediate noisy state using Tweedie's formula
摘要译文
由模型翻译 · 术语统一 · 请对照原文核对关键结论本文提出可扩展的像素锚定端到端扩散方法 SpeeDiff,这是一种从零开始联合训练变分自编码器与扩散模型的潜扩散方法。原则上,联合训练使扩散损失的梯度能够直接引导变分自编码器编码器,促使形成对生成友好的潜空间,并有望比使用预训练冻结变分自编码器的传统两阶段方法更快收敛。然而,朴素的端到端实现会严重降低性能,因为扩散损失的无约束反向传播会导致潜空间坍缩。本文的主要技术贡献是一个简单而有效的 Tweedie 像素重建(TPR)损失:它利用 Tweedie 公式从中间噪声状态解码出预测的干净潜变量,从而提供额外的像素级反馈,缓解坍缩。此外,该方法支持对完全基于 Transformer 的架构进行联合缩放,并在端到端框架内增强表征对齐。在 ImageNet 256×256 生成任务上,SpeeDiff-XL 模型在无引导条件下取得 1.50 的弗雷歇初始距离,训练速度分别比 Vanilla SiT 和 REPA 快 140 倍以上和 61 倍。在使用更高效的 32 倍压缩变分自编码器时,该模型在 ImageNet 512×512 生成任务上于无引导条件下进一步达到 1.53 的弗雷歇初始距离。
英文原文摘要 1351 字符
We present Scalable Pixel-anchored End-to-end Diffusion (SpeeDiff), a latent diffusion method that jointly trains the VAE and the diffusion model from scratch. In principle, joint training allows the diffusion loss gradient to directly guide the VAE encoder, encouraging the formation of a generation-friendly latent space and potentially yielding faster convergence than the conventional two-stage approach with a pretrained frozen VAE. However, a naive end-to-end implementation severely degrades performance, as unrestricted backpropagation of the diffusion loss leads to latent space collapse. Our main technical contribution is a simple yet effective Tweedie Pixel Reconstruction (TPR) loss, which provides additional pixel-level feedback by decoding a predicted clean latent from an intermediate noisy state using Tweedie's formula, thereby alleviating collapse. Furthermore, our method enables jointly scaling a fully transformer-based architecture and enhances representation alignment within the end-to-end framework. Our SpeeDiff-XL model achieves over 140x and 61x faster training compared to Vanilla SiT and REPA, respectively, while attaining an FID of 1.50 without guidance on ImageNet 256x256 generation. With a more efficient 32x compressed VAE, our model further reaches an FID of 1.53 without guidance on ImageNet 512x512 generation.