Universal Loss Reweighting to Balance Lesion Size Inequality in 3D Medical Image Segmentation
This work addresses a specific imbalance issue in medical image segmentation for tasks like detecting multiple lesions, offering an incremental improvement over existing methods.
The paper tackles the problem of lesion size imbalance in 3D medical image segmentation, where small lesions are overshadowed by larger ones, by proposing a loss reweighting method that assigns higher weights to smaller lesions based on inverse volume proportionality. Experiments on public CT datasets (LiTS and LUNA16) and a private MR dataset show that this approach considerably increases detection quality while maintaining state-of-the-art delineation.
Target imbalance affects the performance of recent deep learning methods in many medical image segmentation tasks. It is a twofold problem: class imbalance - positive class (lesion) size compared to negative class (non-lesion) size; lesion size imbalance - large lesions overshadows small ones (in the case of multiple lesions per image). While the former was addressed in multiple works, the latter lacks investigation. We propose a loss reweighting approach to increase the ability of the network to detect small lesions. During the learning process, we assign a weight to every image voxel. The assigned weights are inversely proportional to the lesion volume, thus smaller lesions get larger weights. We report the benefit from our method for well-known loss functions, including Dice Loss, Focal Loss, and Asymmetric Similarity Loss. Additionally, we compare our results with other reweighting techniques: Weighted Cross-Entropy and Generalized Dice Loss. Our experiments show that inverse weighting considerably increases the detection quality, while preserves the delineation quality on a state-of-the-art level. We publish a complete experimental pipeline for two publicly available datasets of CT images: LiTS and LUNA16 (https://github.com/neuro-ml/inverse_weighting). We also show results on a private database of MR images for the task of multiple brain metastases delineation.