用可逆网络和注意力模块，让医学图像生成模型更省显存
Bringing together invertible UNets with invertible attention modules for memory-efficient diffusion models
- 采用可逆UNet与可逆注意力模块构建新架构
- 3D脑肿瘤数据训练时显存峰值降低15%，效果接近顶尖模型
- 适合医疗影像等高维数据的高效生成任务
扩散模型在图像生成任务中表现优异，但通常需要大量计算资源。这一问题在医学图像合成中尤为突出，如CT、MRI、电子显微镜等3D数据集。本文提出一种新型架构，可在单张GPU上实现内存高效的扩散模型训练。该模型结合可逆UNet与可逆注意力模块，实现了两个关键贡献：1）去噪扩散模型，使内存使用与数据维度无关；2）降低训练能耗。该方法适用于多种图像生成任务，我们在3D BraTS2020数据集上验证其内存效率，训练时峰值显存减少最多达15%，生成质量与当前最先进模型相当。
原文摘要 · Abstract (English)
Diffusion models have recently gained state of the art performance on many image generation tasks. However, most models require significant computational resources to achieve this. This becomes apparent in the application of medical image synthesis due to the 3D nature of medical datasets like CT-scans, MRIs, electron microscope, etc. In this paper we propose a novel architecture for a single GPU memory-efficient training for diffusion models for high dimensional medical datasets. The proposed model is built by using an invertible UNet architecture with invertible attention modules. This leads to the following two contributions: 1. denoising diffusion models and thus enabling memory usage to be independent of the dimensionality of the dataset, and 2. reducing the energy usage during training. While this new model can be applied to a multitude of image generation tasks, we showcase its memory-efficiency on the 3D BraTS2020 dataset leading to up to 15\% decrease in peak memory consumption during training with comparable results to SOTA while maintaining the image quality.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。