将可逆UNet与可逆注意力模块结合以实现内存高效的扩散模型
计算机视觉与模式识别
2025-04-16 v1 人工智能
摘要
扩散模型最近在众多图像生成任务中取得了最先进性能。然而，大多数模型需要显著的计算资源才能实现这一点。这在应用医学图像合成时尤为明显，因为医学数据集（如CT扫描、MRI、电子显微镜等）具有三维特征。本文提出一种新型架构，用于高维医学数据集的单GPU内存高效训练扩散模型。该模型基于可逆UNet架构与可逆注意力模块构建，带来以下两个贡献：1. 使扩散模型和去噪过程的内存使用与数据集维度无关，2. 降低训练期间的能源消耗。虽然该新模型可应用于多种图像生成任务，但我们在3D BraTS2020 数据集上展示了其内存效率，训练期间峰值内存消耗可降低最高达15%，同时以相当于SOTA的结果维持图像质量。
引用
@article{arxiv.2504.10883,
title = {Bringing together invertible UNets with invertible attention modules for memory-efficient diffusion models},
author = {Karan Jain and Mohammad Nayeem Teli},
journal= {arXiv preprint arXiv:2504.10883},
year = {2025}
}