Adaptive decoder-block selection with filter reweighting for medical image segmentation

Abstract

Image segmentation models are mostly designed for specific tasks and trained on large-scale datasets. Medical image segmentation tasks, however, are highly diverse and usually have a small training dataset, making it difficult for a segmentation model to generalize well on other tasks. This architecture bias motivates a lot of research on designing a task-adaptive segmentation network with the most suitable settings including the network’s depth and width and feature fusion strategy. In this paper, we propose the adaptive Decoder-block Selection with Filter Reweighting (DeSFiR) algorithm to generate an efficient and effective medical image segmentation model for each specific dataset. This algorithm is a two-step iterative process. In the adaptive decoder-block selection step, we develop an auxiliary indicator function layer to bypass the decoder block whose individual responses are lower than a pre-defined threshold, aiming to generate a lightweight model. In the filter reweighing step, we replace invalid layers in the segmentation network with the external information from a simultaneously trained twin network. To this end, we utilize an entropy-based criterion to measure the informativeness of each layer and an adaptive weighting strategy to balance the replaced information between two networks. We have evaluated our algorithm on three tasks, including the segmentation of neural structures, nuclei, and liver. Our results indicate that the medical image segmentation model generated by the proposed DeSFiR algorithm has a significantly reduced complexity and comparable accuracy.