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.

本研究解决了当前扩散模型在医学图像合成中对计算资源的高需求问题。通过提出一种结合可逆UNet架构和可逆注意力模块的新型架构，实现了在单个GPU上内存高效的训练，并使去噪扩散模型的内存使用与数据集的维度无关，同时降低了训练期间的能耗。本模型在3D BraTS2020数据集上展示了高达15%的内存消耗减少，并且在保持图像质量的情况下与最新技术水平相当。

将可逆UNet与可逆注意力模块结合以实现内存高效的扩散模型