用单步扩散解码器提升图像编码效率,速度更快质量更高。
SSDD: Single-Step Diffusion Decoder for Efficient Image Tokenization
- 设计无对抗训练的单步扩散解码器,结合蒸馏技术优化性能。
- 重建FID降至0.46,吞吐量提升1.4倍,采样速度达DiT的3.8倍。
- 可直接替换现有VAE,适合追求高效生成的模型部署场景。
图像生成模型中的分词器是关键组件,负责从信号中提取重要特征,降低数据维度与冗余。当前主流分词器基于KL正则化变分自编码器(KL-VAE),采用重构、感知和对抗损失进行训练。扩散解码器作为更严谨的替代方案,能建模条件于潜在表示的图像分布,但要达到KL-VAE性能仍需对抗损失,且因迭代采样导致解码时间较长。为此,本文提出一种新型像素级扩散解码器架构,利用Transformer结构与无GAN训练方式实现更好扩展性与训练稳定性。通过知识蒸馏将扩散解码器性能迁移到高效单步解码器中。SSDD是首个无需对抗损失即可实现单步重建的扩散解码器,在保持生成质量的同时,使重建FID从0.87降至0.46,吞吐量提升1.4倍,采样速度达扩散Transformer(DiTs)的3.8倍。该方法可作为KL-VAE的即插即用替代品,助力构建更高性能、更快速的生成模型。
原文摘要 · Abstract (English)
Tokenizers are a key component of state-of-the-art generative image models, extracting the most important features from the signal while reducing data dimension and redundancy. Most current tokenizers are based on KL-regularized variational autoencoders (KL-VAE), trained with reconstruction, perceptual and adversarial losses. Diffusion decoders have been proposed as a more principled alternative to model the distribution over images conditioned on the latent. However, matching the performance of KL-VAE still requires adversarial losses, as well as a higher decoding time due to iterative sampling. To address these limitations, we introduce a new pixel diffusion decoder architecture for improved scaling and training stability, benefiting from transformer components and GAN-free training. We use distillation to replicate the performance of the diffusion decoder in an efficient single-step decoder. This makes SSDD the first diffusion decoder optimized for single-step reconstruction trained without adversarial losses, reaching higher reconstruction quality and faster sampling than KL-VAE. In particular, SSDD improves reconstruction FID from $0.87$ to $0.46$ with $1.4\times$ higher throughput and preserve generation quality of DiTs with $3.8\times$ faster sampling. As such, SSDD can be used as a drop-in replacement for KL-VAE, and for building higher-quality and faster generative models.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。