[论文解读] Invertible Residual Network with Regularization for Effective Medical Image Segmentation

Kashu Yamazaki, Vidhiwar Singh Rathour+1 more|arXiv (Cornell University)|Mar 16, 2021

Advanced Neural Network Applications参考文献 18被引用 4

一句话总结

本文提出了两种可逆残差网络变体——部分可逆残差网络（Partially-InvRes）和完全可逆残差网络（Fully-InvRes），用于实现内存高效的3D医学图像分割。通过利用加法耦合和可学习的挤压操作实现可逆下采样/上采样，模型在训练过程中显著降低了内存使用量，同时在iSeg 2019和BraTS 2020数据集上保持了具有竞争力的分割性能，其中Fully-InvRes的内存使用量仅为2.12GB，相比基线3D Unet的4.18GB大幅减少。

ABSTRACT

Deep Convolutional Neural Networks (CNNs) i.e. Residual Networks (ResNets) have been used successfully for many computer vision tasks, but are difficult to scale to 3D volumetric medical data. Memory is increasingly often the bottleneck when training 3D Convolutional Neural Networks (CNNs). Recently, invertible neural networks have been applied to significantly reduce activation memory footprint when training neural networks with backpropagation thanks to the invertible functions that allow retrieving its input from its output without storing intermediate activations in memory to perform the backpropagation. Among many successful network architectures, 3D Unet has been established as a standard architecture for volumetric medical segmentation. Thus, we choose 3D Unet as a baseline for a non-invertible network and we then extend it with the invertible residual network. In this paper, we proposed two versions of the invertible Residual Network, namely Partially Invertible Residual Network (Partially-InvRes) and Fully Invertible Residual Network (Fully-InvRes). In Partially-InvRes, the invertible residual layer is defined by a technique called additive coupling whereas in Fully-InvRes, both invertible upsampling and downsampling operations are learned based on squeezing (known as pixel shuffle). Furthermore, to avoid the overfitting problem because of less training data, a variational auto-encoder (VAE) branch is added to reconstruct the input volumetric data itself. Our results indicate that by using partially/fully invertible networks as the central workhorse in volumetric segmentation, we not only reduce memory overhead but also achieve compatible segmentation performance compared against the non-invertible 3D Unet. We have demonstrated the proposed networks on various volumetric datasets such as iSeg 2019 and BraTS 2020.

研究动机与目标

• 解决3D卷积神经网络（CNN）在体素医学影像数据上训练时高内存占用的问题。
• 通过引入正则化方法，缓解医学影像中标签数据有限带来的过拟合问题。
• 通过利用可逆残差块，在GPU内存受限条件下实现更深、更大的3D网络训练。
• 通过整合可逆性与变分自编码器（VAE）正则化，改进标准3D Unet基线模型，以提升泛化能力与内存效率。

提出的方法

• 通过加法耦合实现部分可逆残差块（Partially-InvRes），将输入分割为两部分，通过可学习函数与残差连接进行变换，实现可逆性。
• 设计完全可逆残差块（Fully-InvRes），采用可学习的像素重排（挤压）实现可逆下采样与上采样，确保整个网络的完全可逆性。
• 引入VAE正则化分支以重建输入体素数据，提升小样本医学数据集上的泛化能力并减少过拟合。
• 基于3D Unet架构构建网络，将标准残差块替换为可逆版本，并增加VAE分支以实现自监督预训练。
• 使用ADAM优化器进行训练，基础初始学习率为2e-4，所有层均采用实例归一化和Leaky ReLU激活函数。
• 在PyTorch中利用MemCNN库，通过可逆函数反演实现内存高效的反向传播，无需存储中间激活值。

实验结果

研究问题

• RQ1 可逆残差网络是否能在不损失分割精度的前提下，降低3D医学图像分割中的内存使用量？
• RQ2 在体素数据集上，部分可逆与完全可逆残差网络在内存效率与性能方面如何比较？
• RQ3 在低数据量医学影像场景下，添加VAE正则化分支在多大程度上能提升泛化能力并减少过拟合？
• RQ4 可逆网络是否能有效集成到3D Unet架构中，从而在GPU内存受限条件下实现更深网络的训练？
• RQ5 与非可逆3D Unet基线相比，所提出的可逆网络在iSeg 2019和BraTS 2020等标准基准数据集上的表现如何？

主要发现

• 在iSeg 2019数据集上，Fully-InvRes训练时内存使用量最低，仅为2.12GB，显著低于非可逆3D Unet基线所需的4.18GB。
• 在iSeg 2019数据集上，Fully-InvRes对CSF、GM和WM的Dice分数分别达到95.09%、91.32%和90.31%，在所有结构上均略优于基线模型。
• 在BraTS 2020数据集上，Fully-InvRes对WT、ET和TC的Dice分数分别为88.20%、72.01%和77.74%，在所有指标上均优于基线3D Unet。
• 两组数据集的Hausdorff距离均有所降低：例如在BraTS 2020上，Fully-InvRes对WT的Hausdorff距离为10.46，低于基线的11.45，表明边界对齐能力更优。
• 可逆模块与VAE正则化的结合，使Partially-InvRes和Fully-InvRes均能有效避免过拟合，并在训练数据有限的情况下实现高精度分割。
• 表4证实，两种所提模型均具备内存效率高、避免过拟合且精度高的特点，而基线3D Unet则在内存效率与过拟合正则化方面均表现不足。