SDXS：知识蒸馏在高效图像生成中的应用

人工智能咨询培训老师叶梓 转载标明出处

扩散模型虽然在图像生成方面表现出色，但其迭代采样过程导致在低功耗设备上部署面临挑战，同时在云端高性能GPU平台上的能耗也不容忽视。为了解决这一问题，小米公司的Yuda Song、Zehao Sun、Xuanwu Yin等人提出了一种新的方法——SDXS，通过知识蒸馏简化了U-Net和图像解码器架构，并引入了一种创新的一步式DM训练技术，使用特征匹配和得分蒸馏，从而在单GPU上实现了大约100 FPS（比SD v1.5快30倍）和30 FPS（比SDXL快60倍）的推理速度。

图1为在图像生成时间限制为1秒的情况下，不同模型的性能对比。SDXL模型在这种情况下只能使用16次函数评估（NFEs）来生成稍微模糊的图像，而提出的SDXS-1024模型却能够生成30张清晰的图像。这表明SDXS-1024在保持图像质量的同时显著提高了生成速度。本方法还能够训练ControlNet，这是一种能够嵌入空间引导的网络，用于图像到图像的任务，如草图到图像的转换、修复和超分辨率等。证明了SDXS方法的灵活性和应用潜力。

方法

LDM框架由三个关键要素组成：文本编码器、图像解码器以及一个需要多次迭代以生成清晰图像的去噪模型。由于文本编码器的开销相对较低，因此优化其大小并不是研究的重点。

VAE优化：LDM框架通过将样本投影到计算效率更高的低维潜在空间，显著提高了高分辨率图像扩散模型的训练效率。这一过程通过使用预训练模型，如变分自编码器（Variational AutoEncoder, VAE）或向量量化变分自编码器（Vector Quantised-Variational AutoEncoder, VQVAE）来实现高比例图像压缩。VAE包含一个将图像映射到潜在空间的编码器，以及一个重建图像的解码器。其训练通过平衡重建损失、Kullback-Leibler (KL) 散度和GAN损失来优化。然而，训练中对所有样本同等对待引入了冗余。研究者们提出了一种VAE蒸馏（VD）损失，用于训练一个小型的图像解码器G： 其中，D是GAN判别器，用于平衡两个损失项，表示在8倍下采样图像上的L1损失。图2(a)展示了蒸馏小型图像解码器的训练策略。倡使用简化的CNN架构，不包含注意力机制和归一化层等复杂组件，只关注基本的残差块和上采样层。

U-Net优化： LDMs采用U-Net架构作为核心去噪模型，该架构结合了残差块和Transformer块。为了利用预训练的U-Nets的能力，同时减少计算需求和参数数量，研究者们采用了知识蒸馏策略，这一策略受到BK-SDM的块移除训练策略启发。这涉及从U-Net中选择性地移除残差和Tr