[论文解读] iUNets: Fully invertible U-Nets with Learnable Up- and Downsampling
本文提出了iUNets，一种基于指数斜对称矩阵的可学习、可逆上采样与下采样操作的全可逆U-Net架构。通过可逆性实现内存高效的反向传播，iUNets在GPU内存限制下能够训练更深、更大的网络，在3D医学图像分割和CT后处理任务中优于标准U-Net，同时还能作为归一化流使用，并在似然性能上表现具有竞争力。
U-Nets have been established as a standard architecture for image-to-image learning problems such as segmentation and inverse problems in imaging. For large-scale data, as it for example appears in 3D medical imaging, the U-Net however has prohibitive memory requirements. Here, we present a new fully-invertible U-Net-based architecture called the iUNet, which employs novel learnable and invertible up- and downsampling operations, thereby making the use of memory-efficient backpropagation possible. This allows us to train deeper and larger networks in practice, under the same GPU memory restrictions. Due to its invertibility, the iUNet can furthermore be used for constructing normalizing flows.
研究动机与目标
- 为解决标准U-Net在大规模3D医学成像中高内存消耗的问题，通过实现内存高效的训练。
- 开发一种全可逆U-Net架构，支持无需存储中间激活值的反向传播。
- 在固定GPU内存限制下，实现更大、更深网络的训练。
- 证明该架构可作为归一化流用于生成建模的可行性。
- 在3D图像到图像任务（如体积分割和CT重建）中提升性能。
提出的方法
- 提出使用指数斜对称卷积核实现可学习的可逆下采样，其为双射且可微。
- 将可逆上采样定义为所学下采样操作的逆，确保完全可逆性。
- 使用源自斜对称矩阵的正交卷积算子，以在训练过程中保持可逆性和稳定性。
- 采用结构化核重排机制，将空间下采样映射为通道扩展，以保留信息。
- 通过利用激活重建实现内存高效的反向传播，将该架构应用于3D医学成像任务。
- 通过堆叠iUNet模块与仿射耦合层构建归一化流，实现基于似然的生成建模。
实验结果
研究问题
- RQ1能否设计出可学习的可逆上采样与下采样操作，以替代U-Net架构中非可逆操作（如最大池化和转置卷积）？
- RQ2所得到的全可逆U-Net（iUNet）是否能在固定GPU内存下实现更深、更大网络的内存高效训练？
- RQ3iUNets在3D医学图像分割和后处理任务中是否能实现优于标准U-Net的性能？
- RQ4iUNet架构是否适用于归一化流？其似然性能与最先进方法相比如何？
- RQ5所学习的可逆操作是否能在无需引入伪影的情况下泛化到多样化的图像到图像任务？
主要发现
- 在BraTS2018验证集上，通道数为64的iUNet取得了0.850的Dice分数，优于标准U-Net（0.833），也优于更大的iUNets（例如iUNet-64: 0.850 vs. U-Net: 0.833）。
- 在3D CT后处理中，iUNet显著优于基线U-Net，表明内存效率使模型容量得以提升。
- 使用iUNet实现的内存高效反向传播产生的损失曲线与标准反向传播相当，证实了训练的稳定性。
- 基于iUNet的归一化流在CIFAR10上达到3.60 bits/dim的负对数似然，略差于Real NVP（3.49 bits/dim），表明仍有改进空间。
- 使用截断矩阵指数级数实现可逆性可能在似然估计中引入潜在偏差，这可能是性能差距的原因。
- iUNet结构使在内存限制下训练更大网络成为可能，例如更高的通道数（如iUNet-64）带来了性能提升。
更好的研究，从现在开始
从阅读论文到最终审阅，大幅缩短您的研究时间。
无需绑定信用卡
本解读由 AI 生成，并经人工编辑审核。