Sauron U-Net: Simple automated redundancy elimination in medical image segmentation via filter pruning

TL;DR

Medical image segmentation often demands compact, fast, and interpretable models suitable for on-device deployment. Sauron introduces a single-phase filter pruning approach that jointly optimizes a CNN with a $\delta_{opt}$ regularization to promote feature-map clustering, then prunes during training using layer-specific thresholds without predefining cluster counts. Across four medical segmentation datasets, Sauron achieves substantial pruning (often >90%) with Dice scores and HD95 comparable to or better than baselines and delivers faster inference, while producing highly interpretable feature maps. This method offers a practical path to privacy-preserving, on-device segmentation with minimal hyperparameter tuning and broad applicability beyond the medical domain.

Abstract

We introduce Sauron, a filter pruning method that eliminates redundant feature maps of convolutional neural networks (CNNs). Sauron optimizes, jointly with the loss function, a regularization term that promotes feature maps clustering at each convolutional layer by reducing the distance between feature maps. Sauron then eliminates the filters corresponding to the redundant feature maps by using automatically adjusted layer-specific thresholds. Unlike most filter pruning methods, Sauron requires minimal changes to typical neural network optimization because it prunes and optimizes CNNs jointly, which, in turn, accelerates the optimization over time. Moreover, unlike with other cluster-based approaches, the user does not need to specify the number of clusters in advance, a hyperparameter that is difficult to tune. We evaluated Sauron and five state-of-the-art filter pruning methods on four medical image segmentation tasks. This is an area where little attention has been paid to filter pruning, but where smaller CNN models are desirable for local deployment, mitigating privacy concerns associated with cloud-based solutions. Sauron was the only method that achieved a reduction in model size of over 90% without deteriorating substantially the performance. Sauron also achieved, overall, the fastest models at inference time in machines with and without GPUs. Finally, we show through experiments that the feature maps of models pruned with Sauron are highly interpretable, which is essential for medical image segmentation.

Figures (14)

• Figure 1: a-c) tSNE plot of the feature maps of the first block of the decoder at initialization (epoch 0), and after optimizing with and without $\delta_{opt}$. d) Corresponding dip-test values during the optimization. e-g) Summary of the trends across the three clusterability measures in all convolutional layers. h) Number of layers with an increasing trend in the three clusterability measures with higher values of $\lambda$ (dashed line: Sauron's default configuration).
• Figure 2: Image slice from Rats (top-left), KiTS (top-right), ACDC (middle), and ATLAS (bottom) datasets, its ground-truth segmentation, and all feature maps at the second-to-last convolutional block after pruning with Sauron.
• Figure 3: Validation Dice coefficients of baseline nnUNet, Sauron, and two other approaches to normalize $\delta_{opt}$.
• Figure 4: Time required by each epoch divided by the time required by the first epoch during the optimization of nnUNet while pruning with Sauron.
• Figure 5: Diagram of an archetypal nnUNet with five levels.