# Bringing together invertible UNets with invertible attention modules for memory-efficient diffusion models

arXiv: https://arxiv.org/abs/2504.10883 · HTML: https://mlsanity.com/paper/2504.10883/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Medical imaging and segmentation, one of 14 directions in Diffusion models.
- Builds on 43 of its 55 references that are in the map.
- First posted on arXiv 15 Apr 2025.

## Placed in

- Diffusion models / [Medical imaging and segmentation](https://mlsanity.com/diffusion/medical-imaging-and-segmentation/)

## Builds on (most cited first, 30 of 43)

- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [Deep Unsupervised Learning using Nonequilibrium Thermodynamics](https://mlsanity.com/paper/1503.03585/) (2015)
- [DreamFusion: Text-to-3D using 2D Diffusion](https://mlsanity.com/paper/2209.14988/) (2022)
- [CBAM: Convolutional Block Attention Module](https://mlsanity.com/paper/1807.06521/) (2018)
- [Glow: Generative Flow with Invertible 1x1 Convolutions](https://mlsanity.com/paper/1807.03039/) (2018)
- [Zero-1-to-3: Zero-shot One Image to 3D Object](https://mlsanity.com/paper/2303.11328/) (2023)
- [Attention U-Net: Learning Where to Look for the Pancreas](https://mlsanity.com/paper/1804.03999/) (2018)
- [Magic3D: High-Resolution Text-to-3D Content Creation](https://mlsanity.com/paper/2211.10440/) (2022)
- [Swin-Unet: Unet-like Pure Transformer for Medical Image Segmentation](https://mlsanity.com/paper/2105.05537/) (2021)
- [Diffusion Models: A Comprehensive Survey of Methods and Applications](https://mlsanity.com/paper/2209.00796/) (2022)
- [MedMNIST v2 -- A large-scale lightweight benchmark for 2D and 3D biomedical image classification](https://mlsanity.com/paper/2110.14795/) (2021)
- [Identifying the Best Machine Learning Algorithms for Brain Tumor Segmentation, Progression Assessment, and Overall Survival Prediction in the BRATS Challenge](https://mlsanity.com/paper/1811.02629/) (2018)
- [Flow Network based Generative Models for Non-Iterative Diverse Candidate Generation](https://mlsanity.com/paper/2106.04399/) (2021)
- [Brain Imaging Generation with Latent Diffusion Models](https://mlsanity.com/paper/2209.07162/) (2022)
- [ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning](https://mlsanity.com/paper/2104.07857/) (2021)
- [Diffusion Models for Implicit Image Segmentation Ensembles](https://mlsanity.com/paper/2112.03145/) (2021)
- [BAM: Bottleneck Attention Module](https://mlsanity.com/paper/1807.06514/) (2018)
- [Medical Image Segmentation Using Deep Learning: A Survey](https://mlsanity.com/paper/2009.13120/) (2020)
- [Patch Diffusion: Faster and More Data-Efficient Training of Diffusion Models](https://mlsanity.com/paper/2304.12526/) (2023)
- [An Overview of Diffusion Models: Applications, Guided Generation, Statistical Rates and Optimization](https://mlsanity.com/paper/2404.07771/) (2024)
- [Generation of 3D Brain MRI Using Auto-Encoding Generative Adversarial Networks](https://mlsanity.com/paper/1908.02498/) (2019)
- [A General Survey on Attention Mechanisms in Deep Learning](https://mlsanity.com/paper/2203.14263/) (2022)
- [Diffusion Adversarial Representation Learning for Self-supervised Vessel Segmentation](https://mlsanity.com/paper/2209.14566/) (2022)
- [Medical Diffusion: Denoising Diffusion Probabilistic Models for 3D Medical Image Generation](https://mlsanity.com/paper/2211.03364/) (2022)
- [Efficient Diffusion Models for Vision: A Survey](https://mlsanity.com/paper/2210.09292/) (2022)
- [Can segmentation models be trained with fully synthetically generated data?](https://mlsanity.com/paper/2209.08256/) (2022)
- [Spot the fake lungs: Generating Synthetic Medical Images using Neural Diffusion Models](https://mlsanity.com/paper/2211.00902/) (2022)
- [A Unified Framework for U-Net Design and Analysis](https://mlsanity.com/paper/2305.19638/) (2023)

## Abstract (arXiv)

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.