[论文解读] Fast-DDPM: Fast Denoising Diffusion Probabilistic Models for Medical Image-to-Image Generation

Hongxu Jiang, Muhammad Imran+5 more|arXiv (Cornell University)|May 23, 2024

Image and Signal Denoising Methods被引用 7

一句话总结

Fast-DDPM 通过使用面向任务的 10 次时间步调度器，在训练和采样时间上同时减少 DDPM 的开销，在医疗影像到影像生成任务上取得了最先进的结果。

ABSTRACT

Denoising diffusion probabilistic models (DDPMs) have achieved unprecedented success in computer vision. However, they remain underutilized in medical imaging, a field crucial for disease diagnosis and treatment planning. This is primarily due to the high computational cost associated with (1) the use of large number of time steps (e.g., 1,000) in diffusion processes and (2) the increased dimensionality of medical images, which are often 3D or 4D. Training a diffusion model on medical images typically takes days to weeks, while sampling each image volume takes minutes to hours. To address this challenge, we introduce Fast-DDPM, a simple yet effective approach capable of improving training speed, sampling speed, and generation quality simultaneously. Unlike DDPM, which trains the image denoiser across 1,000 time steps, Fast-DDPM trains and samples using only 10 time steps. The key to our method lies in aligning the training and sampling procedures to optimize time-step utilization. Specifically, we introduced two efficient noise schedulers with 10 time steps: one with uniform time step sampling and another with non-uniform sampling. We evaluated Fast-DDPM across three medical image-to-image generation tasks: multi-image super-resolution, image denoising, and image-to-image translation. Fast-DDPM outperformed DDPM and current state-of-the-art methods based on convolutional networks and generative adversarial networks in all tasks. Additionally, Fast-DDPM reduced the training time to 0.2x and the sampling time to 0.01x compared to DDPM. Our code is publicly available at: https://github.com/mirthAI/Fast-DDPM.

研究动机与目标

• 推动医疗成像中的扩散模型发展，解决训练和采样中的高计算成本。
• 提出 Fast-DDPM，使训练和采样在 10 次时间步框架下对齐。
• 为 2D/3D 医学数据开发面向任务的噪声调度器（均匀与非均匀），实现高效扩散。
• 展示在多图像超分辨率、去噪和图像对图像翻译等任务中的速度提升与质量提升。

提出的方法

• 在训练和采样时使用 10 步噪声调度器，而非 1000 步。
• 定义两种调度器：10 步的均匀采样，以及在更高噪声处步距更密的非均匀采样。
• 用一个简单的在 10 个时间步上的损失，对条件去噪器 εθ(x(t),c,t) 进行训练。
• 应用 DDIM 逆时间 ODE 求解器实现高效采样。
• 使用 10 步过程对给定条件 c 的条件模型 x0 进行训练与采样。
• 在三个医学影像到影像任务上用 PSNR 与 SSIM 进行评估。

实验结果

研究问题

• RQ1 在医学影像到影像任务中，10 步时间扩散调度是否能达到或超过 1000 步 DDPM 的性能？
• RQ2 均匀与非均匀的 10 步调度器是否在不同任务上具有依赖于任务的优势？
• RQ3 Fast-DDPM 是否在多图像超分辨率、去噪和翻译任务中胜过 DDPM 及基于 CNN/GAN 的基线？
• RQ4 在医学影像中，10 步 Fast-DDPM 能达到怎样的训练与推理速度提升？
• RQ5 在各任务中，性能对时间步数量的敏感度如何？

主要发现

方法	PSNR	SSIM	训练时间	推理时间
miSRCNN	26.5	0.87	1 h	0.01 s
miSRGAN	26.8	0.88	40 h	0.04 s
DDPM	25.3	0.83	136 h	3.7 m
Fast-DDPM	27.1	0.89	26 h	2.3 s

• Fast-DDPM 在所测试的任务上（如超分辨率、去噪、翻译）实现了比 DDPM 更高的 PSNR 和 SSIM。
• 在超分辨率任务上，训练时间从 136 h（DDPM）降至 26 h；每个体积的推理时间从 3.7 分钟降至 2.3 秒。
• 与 DDPM 相比，在去噪任务中的推理速度提升约 100x，训练提升约 5x。
• 采用 DDIM 采样的 10 步 Fast-DDPM 在各任务上得到具有竞争力或更优的结果。
• 均匀与非均匀调度器在任务上显示出依任务而定的优势；非均匀在超分辨率更好，均匀在去噪/翻译更好。
• 扩展训练（Fast-DDPM*）进一步提高了翻译任务的 SSIM。