Adaptive Optimization with Fewer Epochs Improves Across-Scanner Generalization of U-Net based Medical Image Segmentation

The U-Net architecture is widely used for medical image segmentation. However, accuracy has been observed to drop, sometimes dramatically, when U-Nets are trained on images that have been acquired with a specific scanner, and are applied to images from another scanner. This indicates an overfitting to image characteristics that are irrelevant to the semantic contents, and is usually mitigated with data augmentation. We argue that early stopping additionally improves across-scanner generalization, while greatly reducing training times. For this, we first observe that the widely used stochastic gradient descent (SGD) trains different U-Net layers at different speeds, and demonstrate that this problem is reduced by switching to AvaGrad, a recently proposed adaptive optimizer. On two different datasets, this allows us to match accuracies from nnUNets with default settings, 1000 epochs of SGD, by training for only 50 epochs with AvaGrad, and to exceed their results in the across-scanner setting. This benefit is specific to combining adaptive optimization and early stopping, since it can be matched neither by SGD with a low number of epochs, nor by Avagrad with many epochs. Finally, we demonstrate that the choice of optimizer can have important implications for domain adaptation. In particular, the SpotTUnet, which was recently proposed to automatically select layers for fine-tuning, arrives at very different policies depending on the optimizer.

Citation information

Sheikh, Rasha; Klasen, Morris; Schultz, Thomas: Adaptive Optimization with Fewer Epochs Improves Across-Scanner Generalization of U-Net based Medical Image Segmentation, Domain Adaptation and Representation Transfer. DART 2022, 2022, https://link.springer.com/chapter/10.1007/978-3-031-16852-9_12, Sheikh.etal.2022a,

Associated Lamarr Researchers

lamarr institute person Schultz Thomas scaled e1663922506873 - Lamarr Institute for Machine Learning (ML) and Artificial Intelligence (AI)

Prof. Dr. Thomas Schultz

Principal Investigator Life Sciences to the profile