[论文解读] The Swiss Army Knife for Image-to-Image Translation: Multi-Task Diffusion Models

Julia Wolleb, Robin Sandkühler+2 more|arXiv (Cornell University)|Apr 6, 2022

Radiomics and Machine Learning in Medical Imaging被引用 17

一句话总结

本文提出一种用于图像到图像翻译的多任务扩散模型，通过利用回归或分割模型的外部梯度来引导去噪扩散采样，而无需微调基础扩散模型。该方法实现了可控的图像翻译——如人脸老化或脑肿瘤生长模拟——同时保留非目标特征，在使用单一预训练DDPM的情况下，在人脸老化和MR图像肿瘤生成任务上均取得了最先进性能。

ABSTRACT

Recently, diffusion models were applied to a wide range of image analysis tasks. We build on a method for image-to-image translation using denoising diffusion implicit models and include a regression problem and a segmentation problem for guiding the image generation to the desired output. The main advantage of our approach is that the guidance during the denoising process is done by an external gradient. Consequently, the diffusion model does not need to be retrained for the different tasks on the same dataset. We apply our method to simulate the aging process on facial photos using a regression task, as well as on a brain magnetic resonance (MR) imaging dataset for the simulation of brain tumor growth. Furthermore, we use a segmentation model to inpaint tumors at the desired location in healthy slices of brain MR images. We achieve convincing results for all problems.

研究动机与目标

• 使用单一预训练的扩散模型实现灵活、任务特定的图像到图像翻译。
• 解决仅修改与任务相关的图像特征，同时保留背景和结构细节的挑战。
• 将分类引导从分类任务扩展到扩散图像翻译中的回归和分割任务。
• 为医学影像生成逼真的合成数据，例如在健康MR扫描中模拟脑肿瘤。
• 证明同一扩散模型可通过推理时梯度注入在多种任务中重复使用。

提出的方法

• 使用在源数据集上训练的预训练去噪扩散概率模型（DDPM）作为基础生成器，且未进行任务特定的微调。
• 在相同数据集上独立训练一个外部任务特定模型（回归或分割），以提供引导信号。
• 在DDIM采样过程中，通过在每一步计算任务特定模型输出相对于噪声输入的梯度，引导去噪过程。
• 引导强度由缩放因子 $ s_t $ 控制，可设置为正或负值，以增强或减弱目标属性（如年龄或肿瘤大小）。
• 对于回归任务，使用目标值 $ i $ 计算 $ s_t = i - R(x_t, t) $，其中 $ R $ 为在步骤 $ t $ 时回归模型的预测值。
• 对于分割任务，标签掩码定义了期望的肿瘤位置，通过使用分割模型损失的梯度，在采样过程中将类似肿瘤的特征注入图像。

实验结果

研究问题

• RQ1 是否可以使用单一预训练的扩散模型在不微调的情况下完成多种图像到图像翻译任务？
• RQ2 来自外部回归和分割模型的梯度引导是否能有效控制采样过程中的图像属性（如年龄或肿瘤大小）？
• RQ3 该方法在仅修改目标属性的同时，是否能保留非目标图像特征（如背景、衣物、发型）？
• RQ4 该方法是否能生成逼真的合成医学图像（如在健康MR切片中模拟脑肿瘤）以用于数据增强？
• RQ5 与任务特定的扩散模型相比，该方法在图像保真度和属性控制方面表现如何？

主要发现

• 该方法成功在人脸图像上实现年龄回归，从40岁输入生成目标年龄（如10、20、60、80岁）的逼真图像。
• 对于脑部MR图像，通过调整回归损失梯度，模型能准确模拟肿瘤的生长或缩小，肿瘤大小可通过目标值 $ i $ 进行控制。
• 基于分割的引导方法成功在健康脑部MR切片中用户指定位置处修复（内补）肿瘤，差异图显示了区域特异性变化，验证了有效性。
• 同一基础DDPM被重复用于所有任务——回归、肿瘤生长模拟和肿瘤内补——展示了该框架的多功能性。
• 结果表明，仅与目标属性相关的特征被修改；背景、面部特征和衣物与输入图像保持一致。
• 该方法在人脸老化和脑肿瘤模拟任务上均取得了令人信服的视觉效果，具有极少的伪影和高度的结构保真度。