Universal Loss Reweighting to Balance Lesion Size Inequality in 3D Medical Image Segmentation

TL;DR

This paper introduces a loss reweighting method that enhances the detection of small lesions in 3D medical image segmentation by inversely weighting voxel contributions based on lesion volume, improving detection without sacrificing delineation quality.

Contribution

The authors propose a novel inverse lesion size reweighting technique applicable to various loss functions, addressing lesion size imbalance in medical image segmentation.

Findings

Inverse weighting improves small lesion detection.

Method maintains state-of-the-art delineation quality.

Effective across multiple datasets and imaging modalities.

Abstract

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.