Source-linked AI summary
Tversky loss function for image segmentation using 3D fully convolutional deep networks
Seyed Sadegh Mohseni Salehi, Deniz Erdogmus, Ali Gholipour
TL;DR
Medical image segmentation is challenged by extreme class imbalance, which can favor high precision over recall when missed lesions are especially costly. This paper introduces a Tversky-index loss for 3D U-net-based fully convolutional networks, and MS lesion experiments report improved test metrics over Dice loss, including DSC and F2 scores.
Problem
Extreme lesion/non-lesion imbalance can produce high-precision, low-recall segmentations, despite the importance of sensitivity in lesion detection.
Method
The paper adds a Tversky-index loss layer to a 3D fully convolutional U-net and adjusts α and β to control false-positive versus false-negative emphasis.
Results
All test performance metrics improved with Tversky loss rather than Dice loss in multiple sclerosis lesion segmentation, with β = 0.7 producing the best combined results.
Takeaways & Limitations
Tversky loss provides a generalized framework for training deep networks on highly imbalanced segmentation data while emphasizing recall over precision.
Abstract
from arXiv · showhide
Fully convolutional deep neural networks carry out excellent potential for fast and accurate image segmentation. One of the main challenges in training these networks is data imbalance, which is particularly problematic in medical imaging applications such as lesion segmentation where the number of lesion voxels is often much lower than the number of non-lesion voxels. Training with unbalanced data can lead to predictions that are severely biased towards high precision but low recall (sensitivity), which is undesired especially in medical applications where false negatives are much less tolerable than false positives. Several methods have been proposed to deal with this problem including balanced sampling, two step training, sample re-weighting, and similarity loss functions. In this paper, we propose a generalized loss function based on the Tversky index to address the issue of data imbalance and achieve much better trade-off between precision and recall in training 3D fully convolutional deep neural networks. Experimental results in multiple sclerosis lesion segmentation on magnetic resonance images show improved F2 score, Dice coefficient, and the area under the precision-recall curve in test data. Based on these results we suggest Tversky loss function as a generalized framework to effectively train deep neural networks.
1 Introduction
Medical image segmentation faces severe lesion/non-lesion imbalance, which can produce high-precision, low-recall predictions. The paper motivates a Tversky-index loss for controlling the precision–recall trade-off in 3D fully convolutional networks.
- Lesion detection can contain more than 500 times as many non-lesion as lesion voxels.
- Without label balancing, training may bias predictions toward non-lesion tissue, yielding high precision but low recall.
- High sensitivity is especially important in computer-aided diagnosis and clinical decision support systems.
- Existing responses to imbalance include balanced sampling, hierarchical or retraining strategies, and sample-reweighted loss functions.
- The proposed Tversky loss generalizes similarity-based losses and aims to emphasize false negatives while improving precision–recall trade-offs in 3D FCNs.
2 Method
The method combines a 3D U-net-style fully convolutional network with a Tversky-index loss layer for imbalanced lesion segmentation. The loss adjusts false-positive and false-negative penalties, while evaluation uses overlap, recall-oriented, and precision–recall metrics.
- Network architecture: The 3D U-net-style network combines high-resolution local features from the contracting path with upsampled low-resolution global features in the expanding path.Features from corresponding paths are concatenated to integrate local and global information.
- Network architecture: The network uses padded 3×3×3 convolutions and ReLU layers, 2×2×2 max pooling for downsampling, transposed convolutions for upsampling, and a final 1×1×1 softmax convolution.The output layer contains one plane per class, with c = 2 for lesion detection.
- Tversky loss layer: The Tversky loss layer is introduced because the Dice loss weights false positives and false negatives equally, whereas lesion segmentation requires greater emphasis on false negatives.This design targets highly imbalanced data where detecting small lesions is crucial.
- Tversky loss layer: The Tversky index uses α and β to control penalties for false positives and false negatives, respectively, enabling adjustment of the precision–recall trade-off.The formulation is used with softmax probabilities and ground-truth voxel labels.
- Training and evaluation: The model is trained end-to-end on T1-weighted, T2-weighted, and FLAIR MRI from 15 subjects using two-fold cross-validation.Images were rigidly registered to 128×224×256, and lesion predictions used a 0.5 probability threshold.
- Training and evaluation: Performance is evaluated with DSC, specificity, sensitivity, F2 score, and the precision–recall curve with its area under the curve for the highly skewed dataset.The F2 score gives greater importance to recall than F1 or DSC, while PR-based evaluation is preferred for skewed test data.
3 Results
Varying Tversky-loss penalties controlled the sensitivity–specificity trade-off, with β = 0.7 producing the strongest combined test performance across lesion-density conditions.
- β = 0.7 achieved the best combined test performance, including APR, F2, and DSC, while Dice-loss training at β = 0.5 performed poorly.
- Higher β increased sensitivity (recall) but reduced specificity, demonstrating control over the sensitivity–specificity balance through the loss-function parameters.
- The β = 0.7 model consistently provided the best precision–recall trade-off across the full test set and subjects with very different lesion densities.
- For subjects with extremely high or extremely low lesion density, varying false-positive and false-negative penalties showed the best results at β = 0.7.
4 Discussion and conclusion
The paper introduces a Tversky-index loss for highly imbalanced segmentation and integrates it into a 3D U-net-style fully convolutional network. In MS lesion segmentation, Tversky training improved test metrics relative to Dice loss, particularly when β = 0.7 emphasizes recall.
- The proposed Tversky loss generalizes the Dice coefficient and Fβ scores to improve precision–recall trade-offs in highly imbalanced segmentation.
- The loss was integrated into a state-of-the-art 3D fully convolutional neural network based on the U-net architecture.
- MS lesion-segmentation experiments showed improved test performance across evaluation metrics, including DSC and F2 scores, compared with Dice loss.
- Using β = 0.7 deliberately weighted recall higher than precision, while consistent test improvements indicated better generalization through this training approach.