# Memory Efficient 3D U-Net with Reversible Mobile Inverted Bottlenecks for Brain Tumor Segmentation

arXiv: https://arxiv.org/abs/2104.09648 · HTML: https://mlsanity.com/paper/2104.09648/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 11 of its 22 references that are in the map.
- First posted on arXiv 19 Apr 2021.

## Builds on (most cited first, 11 of 11)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks](https://mlsanity.com/paper/1905.11946/) (2019)
- [Fully Convolutional Networks for Semantic Segmentation](https://mlsanity.com/paper/1411.4038/) (2014)
- [MobileNetV2: Inverted Residuals and Linear Bottlenecks](https://mlsanity.com/paper/1801.04381/) (2018)
- [Group Normalization](https://mlsanity.com/paper/1803.08494/) (2018)
- [MnasNet: Platform-Aware Neural Architecture Search for Mobile](https://mlsanity.com/paper/1807.11626/) (2018)
- [Identifying the Best Machine Learning Algorithms for Brain Tumor Segmentation, Progression Assessment, and Overall Survival Prediction in the BRATS Challenge](https://mlsanity.com/paper/1811.02629/) (2018)
- [The Reversible Residual Network: Backpropagation Without Storing Activations](https://mlsanity.com/paper/1707.04585/) (2017)
- [TernausNet: U-Net with VGG11 Encoder Pre-Trained on ImageNet for Image Segmentation](https://mlsanity.com/paper/1801.05746/) (2018)
- [A Partially Reversible U-Net for Memory-Efficient Volumetric Image Segmentation](https://mlsanity.com/paper/1906.06148/) (2019)

## Built on by (most cited first, 3 of 3)

- [Reversible Vision Transformers](https://mlsanity.com/paper/2302.04869/) (2023)
- [RevBiFPN: The Fully Reversible Bidirectional Feature Pyramid Network](https://mlsanity.com/paper/2206.14098/) (2022)
- [Towards Zero Memory Footprint Spiking Neural Network Training](https://mlsanity.com/paper/2308.08649/) (2023)

## Abstract (arXiv)

We propose combining memory saving techniques with traditional U-Net architectures to increase the complexity of the models on the Brain Tumor Segmentation (BraTS) challenge. The BraTS challenge consists of a 3D segmentation of a 240x240x155x4 input image into a set of tumor classes. Because of the large volume and need for 3D convolutional layers, this task is very memory intensive. To address this, prior approaches use smaller cropped images while constraining the model's depth and width. Our 3D U-Net uses a reversible version of the mobile inverted bottleneck block defined in MobileNetV2, MnasNet and the more recent EfficientNet architectures to save activation memory during training. Using reversible layers enables the model to recompute input activations given the outputs of that layer, saving memory by eliminating the need to store activations during the forward pass. The inverted residual bottleneck block uses lightweight depthwise separable convolutions to reduce computation by decomposing convolutions into a pointwise convolution and a depthwise convolution. Further, this block inverts traditional bottleneck blocks by placing an intermediate expansion layer between the input and output linear 1x1 convolution, reducing the total number of channels. Given a fixed memory budget, with these memory saving techniques, we are able to train image volumes up to 3x larger, models with 25% more depth, or models with up to 2x the number of channels than a corresponding non-reversible network.
