mlsanityData to 7 Oct 2026

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

arXiv 2504.10883 ↗First posted 15 Apr 2025cs.CV cs.AIcited by 0

On the map: Diffusion models → Medical imaging and segmentation

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.

Key facts. No citing papers in the map yet (data through 7 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.

Builds on

43 of 55 references.

Closest ancestors · 7

Same direction, newest first

  1. 2024CT Synthesis with Conditional Diffusion Models for Abdominal Lymph Node Segmentation26 Mar
  2. 2024An Ordinal Diffusion Model for Generating Medical Images with Different Severity Levels1 Mar
  3. 2023MAM-E: Mammographic synthetic image generation with diffusion models16 Nov
  4. 2022Medical Diffusion: Denoising Diffusion Probabilistic Models for 3D Medical Image Generation7 Nov
  5. 2022Spot the fake lungs: Generating Synthetic Medical Images using Neural Diffusion Models2 Nov
  6. 2022Brain Imaging Generation with Latent Diffusion Models15 Sep
  7. 2021Diffusion Models for Implicit Image Segmentation Ensembles6 Dec

Other references · 36

  1. 2020An Image is Worth 16x16 Words: Transformers for Image Recognition at Scalecited by 29,915
  2. 2020Denoising Diffusion Probabilistic Modelscited by 20,851
  3. 2015U-Net: Convolutional Networks for Biomedical Image Segmentationcited by 19,604
  4. 2015Deep Unsupervised Learning using Nonequilibrium Thermodynamicscited by 7,005
  5. 2022DreamFusion: Text-to-3D using 2D Diffusioncited by 3,006
  6. 2018CBAM: Convolutional Block Attention Modulecited by 2,463
  7. 2018Glow: Generative Flow with Invertible 1x1 Convolutionscited by 1,539
  8. 2023Zero-1-to-3: Zero-shot One Image to 3D Objectcited by 1,483
  9. 2018Attention U-Net: Learning Where to Look for the Pancreascited by 1,420
  10. 2022Magic3D: High-Resolution Text-to-3D Content Creationcited by 1,285
  11. 2021Swin-Unet: Unet-like Pure Transformer for Medical Image Segmentationcited by 1,218
  12. 2022Diffusion Models: A Comprehensive Survey of Methods and Applicationscited by 1,107
24 more
  1. 2021MedMNIST v2 -- A large-scale lightweight benchmark for 2D and 3D biomedical image classificationcited by 682
  2. 2018Identifying the Best Machine Learning Algorithms for Brain Tumor Segmentation, Progression Assessment, and Overall Survival Prediction in the BRATS Challengecited by 508
  3. 2021Flow Network based Generative Models for Non-Iterative Diverse Candidate Generationcited by 384
  4. 2021ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learningcited by 253
  5. 2018BAM: Bottleneck Attention Modulecited by 171
  6. 2020Medical Image Segmentation Using Deep Learning: A Surveycited by 163
  7. 2023Patch Diffusion: Faster and More Data-Efficient Training of Diffusion Modelscited by 131
  8. 2024An Overview of Diffusion Models: Applications, Guided Generation, Statistical Rates and Optimizationcited by 72
  9. 2019Generation of 3D Brain MRI Using Auto-Encoding Generative Adversarial Networkscited by 65
  10. 2022A General Survey on Attention Mechanisms in Deep Learningcited by 64
  11. 2022Diffusion Adversarial Representation Learning for Self-supervised Vessel Segmentationcited by 55
  12. 2022Efficient Diffusion Models for Vision: A Surveycited by 50
  13. 2022Can segmentation models be trained with fully synthetically generated data?cited by 37
  14. 2023A Unified Framework for U-Net Design and Analysiscited by 27
  15. 2020AttentionNAS: Spatiotemporal Attention Cell Search for Video Classificationcited by 23
  16. 2019Dataset Growth in Medical Image Analysis Researchcited by 16
  17. 2020Memory-efficient GAN-based Domain Translation of High Resolution 3D Medical Imagescited by 14
  18. 2022Multi-scale Transformer Network with Edge-aware Pre-training for Cross-Modality MR Image Synthesiscited by 13
  19. 2023Diffusion-Model-Assisted Supervised Learning of Generative Models for Density Estimationcited by 13
  20. 2021Energy Consumption of Deep Generative Audio Modelscited by 12
  21. 2020iUNets: Fully invertible U-Nets with Learnable Up- and Downsamplingcited by 7
  22. 2021Invertible Attentioncited by 5
  23. 2023The Missing U for Efficient Diffusion Modelscited by 5
  24. 2023Memory Efficient Diffusion Probabilistic Models via Patch-based Generationcited by 3

Built on by

Not cited yet.

Similar, not cited

Close in wording, no citation either way

  1. 2026Performance Analysis and Optimization of 3D Generative Diffusion Models across GPU Architectures
  2. 2023Memory-Efficient 3D Denoising Diffusion Models for Medical Image Processing
  3. 2026PPDM: Pixel Puzzling Diffusion Model for Speed and Memory Efficient Volumetric Medical Image Translation
  4. 2026VolDiT: Controllable Volumetric Medical Image Synthesis with Diffusion Transformers
  5. 2024WDM: 3D Wavelet Diffusion Models for High-Resolution Medical Image Synthesis
  6. 2023Investigating Data Memorization in 3D Latent Diffusion Models for Medical Image Synthesis
  7. 2024DiG: Scalable and Efficient Diffusion Models with Gated Linear Attention
  8. 2025MedLoRD: A Medical Low-Resource Diffusion Model for High-Resolution 3D CT Image Synthesis