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

TL;DR

Sauron U-Net introduces a novel filter pruning method that reduces CNN model size by over 90% in medical image segmentation tasks, maintaining performance, accelerating inference, and enhancing interpretability.

Contribution

The paper presents Sauron, a joint optimization and pruning method that automatically eliminates redundant feature maps without needing pre-specified cluster numbers.

Findings

Achieved over 90% model size reduction without performance loss.

Produced the fastest inference times among compared methods.

Generated highly interpretable feature maps for medical segmentation.

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.