[Paper Review] 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 Applications18 references4 citations

TL;DR

This paper proposes two invertible residual network variants—Partially-InvRes and Fully-InvRes—for memory-efficient 3D medical image segmentation. By leveraging additive coupling and learnable squeezing operations for invertible downsampling/upsampling, the models reduce memory usage during training while maintaining competitive segmentation performance on iSeg 2019 and BraTS 2020, with Fully-InvRes achieving 2.12GB memory usage compared to 4.18GB for the baseline 3D Unet.

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.

Motivation & Objective

• Address the high memory footprint of 3D Convolutional Neural Networks (CNNs) when training on volumetric medical imaging data.
• Overcome the challenge of limited labeled data in medical imaging by incorporating regularization to reduce overfitting.
• Enable training of deeper and larger 3D networks under constrained GPU memory by leveraging invertible residual blocks.
• Improve upon the standard 3D Unet baseline by integrating invertibility and variational autoencoder (VAE) regularization for better generalization and memory efficiency.

Proposed method

• Implement Partially-InvRes using additive coupling to create invertible residual blocks, where input is split into two parts and transformed via learnable functions with residual connections.
• Design Fully-InvRes with invertible downsampling and upsampling using learnable pixel shuffling (squeezing), ensuring full invertibility of the entire network.
• Integrate a VAE regularization branch to reconstruct input volumetric data, improving generalization and reducing overfitting on small medical datasets.
• Construct the network architecture based on 3D Unet, replacing standard residual blocks with invertible counterparts and adding the VAE branch for self-supervised pretraining.
• Train models using ADAM optimizer with a base learning rate of 2e-4, employing instance normalization and Leaky ReLU in all layers.
• Utilize the MemCNN library in PyTorch to enforce memory-efficient backpropagation via invertible function inversion, eliminating the need to store intermediate activations.

Experimental results

Research questions

• RQ1 Can invertible residual networks reduce memory usage in 3D medical image segmentation without sacrificing segmentation accuracy?
• RQ2 How do partially invertible and fully invertible residual networks compare in terms of memory efficiency and performance on volumetric datasets?
• RQ3 To what extent does adding a VAE regularization branch improve generalization and reduce overfitting in low-data medical imaging scenarios?
• RQ4 Can invertible networks be effectively integrated into the 3D Unet architecture to enable training of deeper models under limited GPU memory?
• RQ5 How do the proposed invertible networks perform on standard benchmarks like iSeg 2019 and BraTS 2020 compared to the non-invertible 3D Unet baseline?

Key findings

• Fully-InvRes achieved the lowest memory usage at 2.12GB during training on the iSeg 2019 dataset, significantly reducing the 4.18GB required by the non-invertible 3D Unet baseline.
• On the iSeg 2019 dataset, Fully-InvRes achieved a Dice score of 95.09% for CSF, 91.32% for GM, and 90.31% for WM—slightly better than the baseline across all structures.
• On the BraTS 2020 dataset, Fully-InvRes achieved a Dice score of 88.20% for WT, 72.01% for ET, and 77.74% for TC, outperforming the baseline 3D Unet in all metrics.
• The Hausdorff distance was reduced in both datasets: for example, on BraTS 2020, Fully-InvRes achieved 10.46 for WT, compared to 11.45 for the baseline, indicating improved boundary alignment.
• The combination of invertible blocks and VAE regularization enabled both Partially-InvRes and Fully-InvRes to avoid overfitting and achieve high accuracy despite limited training data.
• Table 4 confirms that both proposed models are memory-efficient, avoid overfitting, and achieve high accuracy—unlike the baseline 3D Unet, which lacks both memory efficiency and overfitting regularization.