[论文解读] Improving Diffusion Model Efficiency Through Patching

Troy Luhman, Eric Luhman|arXiv (Cornell University)|Jul 9, 2022

Generative Adversarial Networks and Image Synthesis被引用 5

一句话总结

该论文提出了一种补丁扩散模型（Patched Diffusion Models, PDM），通过采用ViT风格的补丁化处理，显著降低了扩散模型中的计算和内存开销。通过将空间补丁转换为通道维度，PDM在保持生成质量的同时，大幅减少了推理时间与内存使用量，尤其在高分辨率图像（如ImageNet 256²和FFHQ 1024²）上表现显著。

ABSTRACT

Diffusion models are a powerful class of generative models that iteratively denoise samples to produce data. While many works have focused on the number of iterations in this sampling procedure, few have focused on the cost of each iteration. We find that adding a simple ViT-style patching transformation can considerably reduce a diffusion model's sampling time and memory usage. We justify our approach both through an analysis of the diffusion model objective, and through empirical experiments on LSUN Church, ImageNet 256, and FFHQ 1024. We provide implementations in Tensorflow and Pytorch.

研究动机与目标

• 解决扩散模型采样过程中计算与内存成本过高的问题，尤其是在高分辨率图像上的问题。
• 识别出在大多数时间步长下，高频率卷积层在扩散模型中是冗余的，原因在于最优重建过程中的模糊性。
• 提出一种简单、即插即用的方法，通过补丁化将空间特征重新组织为通道维度，以提升效率。
• 在降低计算开销的同时，保持扩散过程的统计等价性。
• 通过可扩展的基于补丁的架构，在ImageNet 256²等具有挑战性的数据集上实现高效且高保真度的图像生成。

提出的方法

• 应用ViT风格的补丁化处理，将空间图像补丁转换为通道维度，降低空间分辨率并增加通道数量。
• 重新设计扩散模型，使其在补丁化后的潜在表示上运行，同时保持相同的损失函数和生成过程。
• 采用数据预测头（x_θ）而非噪声预测头，以提升训练稳定性并降低模型脆弱性。
• 引入多阶段网络拆分策略，使PDM能够扩展至ImageNet 256²等复杂数据集。
• 通过调整损失函数中的超参数γₜ = √(αₜ / (1−αₜ))，优先学习低层次特征而非高频细节。
• 在TensorFlow和PyTorch中均实现该方法，以确保广泛适用性与可复现性。

实验结果

研究问题

• RQ1 在不降低样本质量的前提下，补丁化处理能在多大程度上降低扩散模型采样的计算与内存开销？
• RQ2 高频率卷积层的冗余性如何影响扩散模型中的最优重建？
• RQ3 在补丁化设置下，数据预测（x_θ）是否在训练稳定性和样本保真度方面优于噪声预测（ε_θ）？
• RQ4 补丁扩散模型是否能在ImageNet 256²等具有挑战性的数据集上实现高保真度图像合成？
• RQ5 补丁化变换如何影响扩散过程的统计等价性与似然目标？

主要发现

• 补丁扩散模型显著降低了采样时间与内存使用量，尤其在高分辨率图像（如FFHQ 1024²和ImageNet 256²）上表现突出。
• 最优重建x*(zₜ,t)位于低维流形上，且呈现模糊特性，表明在大多数时间步长下，高频分量并非关键。
• 高频率卷积层在扩散过程中是冗余的，因为其对最优重建的贡献极小，尤其是在后期时间步长。
• 数据预测（x_θ）比噪声预测（ε_θ）更具鲁棒性，从而在补丁化设置下实现更稳定的训练与更优的样本质量。
• 所提出的方法通过拆分网络架构，在ImageNet 256²上实现了高保真度图像合成，证明了其在复杂数据集上的可扩展性。
• 补丁化操作保持了扩散过程的统计结构，使得该方法可无缝集成至任意现有扩散模型中，实现即插即用。