[论文解读] Diffusion Probabilistic Model Made Slim

Xingyi Yang, Daquan Zhou+2 more|arXiv (Cornell University)|Nov 27, 2022

Generative Adversarial Networks and Image Synthesis被引用 4

一句话总结

本文提出Spectral Diffusion（SD），一种轻量级扩散模型，通过小波门控和谱感知蒸馏，提升紧凑架构中高频细节的恢复能力。通过在反向去噪过程中建模频率动态，并利用反向加权损失蒸馏高频知识，SD在模型大小上实现8–18倍减小、推理速度提升2–5倍，同时保持图像保真度损失最小。

ABSTRACT

Despite the recent visually-pleasing results achieved, the massive computational cost has been a long-standing flaw for diffusion probabilistic models (DPMs), which, in turn, greatly limits their applications on resource-limited platforms. Prior methods towards efficient DPM, however, have largely focused on accelerating the testing yet overlooked their huge complexity and sizes. In this paper, we make a dedicated attempt to lighten DPM while striving to preserve its favourable performance. We start by training a small-sized latent diffusion model (LDM) from scratch, but observe a significant fidelity drop in the synthetic images. Through a thorough assessment, we find that DPM is intrinsically biased against high-frequency generation, and learns to recover different frequency components at different time-steps. These properties make compact networks unable to represent frequency dynamics with accurate high-frequency estimation. Towards this end, we introduce a customized design for slim DPM, which we term as Spectral Diffusion (SD), for light-weight image synthesis. SD incorporates wavelet gating in its architecture to enable frequency dynamic feature extraction at every reverse steps, and conducts spectrum-aware distillation to promote high-frequency recovery by inverse weighting the objective based on spectrum magni tudes. Experimental results demonstrate that, SD achieves 8-18x computational complexity reduction as compared to the latent diffusion models on a series of conditional and unconditional image generation tasks while retaining competitive image fidelity.

研究动机与目标

• 为解决训练紧凑扩散模型时难以保留高频图像细节的挑战。
• 探究尽管经过网络结构剪枝，小型扩散模型在纹理和边缘恢复方面仍表现不佳的原因。
• 开发一种方法，显式建模反向去噪过程中的频率动态与偏差，以提升紧凑模型的性能。
• 在不牺牲条件与非条件生成任务中图像质量的前提下，实现显著的模型压缩与推理加速。

提出的方法

• 在U-Net的跳跃连接中引入小波门控，以在每个反向去噪步骤中动态调制高频、低频和带通分量。
• 对小波变换后的特征应用可学习门控函数，实现在不增加参数量的前提下自适应频率响应。
• 通过按频率幅度反向加权的损失函数实现谱感知知识蒸馏，重点强化高频分量的恢复。
• 使用预训练教师模型训练小型学生模型，将损失组件分离至空间域与频率域，实现针对性优化。
• 利用离散傅里叶变换（DFT）分析诊断生成图像中的频率偏差与演化过程，验证设计选择的合理性。
• 将小波门控与谱感知蒸馏整合为统一的训练目标，以在轻量化模型中保留高频细节。

Figure 1 : (1) Visualizing the frequency domain gap among generated images with the full DPM [ 43 ] , Lite DPM and our SD on FFHQ [ 26 ] dataset. Lite-DPM is unable to recover fine-grained textures, while SD can produce sharp edges and realistic patterns. (2) Model size, Multiply-Add cumulation (MAC

实验结果

研究问题

• RQ1 为何小型扩散模型难以恢复如纹理和边缘等高频图像分量？
• RQ2 在标准扩散模型中，反向去噪步骤中去噪输出的频率成分如何演化？
• RQ3 去噪过程中存在的频率偏差在多大程度上损害了小型模型重建精细细节的能力？
• RQ4 基于小波的门控能否改善紧凑扩散模型中的频率感知特征表示？
• RQ5 在知识蒸馏中对频率分量进行反向加权，能否增强小型模型对高频细节的恢复能力？

主要发现

• 与标准潜在扩散模型相比，Spectral Diffusion将模型尺寸减少8–18倍，推理时间缩短2–5倍，同时保持了具有竞争力的图像质量。
• 消融实验表明，移除小波门控后，FFHQ数据集上的FID从10.5上升至12.4，证明其在保留图像细节方面具有关键作用。
• 谱感知蒸馏带来1.8个FID点的性能提升，显著优于仅基于空间的蒸馏方法。
• 在参数量为7760万的情况下，SD在MS-COCO文本到图像生成任务中取得18.43的FID，仅为基线LDM模型的1/18.7。
• 可视化结果表明，SD能恢复如发丝和建筑图案等精细纹理，而基线模型生成的输出则更平滑、细节更少。
• 门控函数值随去噪步骤演化：高频分量在后期上采样操作中得到增强，与理论频率演化趋势一致。

Figure 2 : Illustration of the Frequency Evolution and Bias for Diffusion Models. In the reverse process, the optimal filters recover low-frequency components first and add on the details at the end. The predicted score functions may be incorrect for rare patterns, thus failing to recover complex an