iUNets: Fully invertible U-Nets with Learnable Up- and Downsampling
This addresses memory limitations for researchers and practitioners in fields like medical imaging, enabling more efficient training of deep networks, though it is incremental as it builds on the established U-Net architecture.
The paper tackles the prohibitive memory requirements of U-Nets in large-scale data like 3D medical imaging by introducing iUNets, a fully-invertible U-Net architecture with learnable up- and downsampling, enabling memory-efficient backpropagation and allowing training of deeper and larger networks under the same GPU memory constraints.
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.