Improving MS lesion segmentation by alleviating the class imbalance and conformal risk control

1886 

Abstract Submission 

YANG SUN¹, Bernd Taschler², Laura Gaetano³, Thomas Nichols⁴, Dieter Haering³, Habib Ganjgahi⁴

¹University of Oxford, Oxford, Oxfordshire, ²University of Oxford, Goring-on-Thames, South Oxfordshire, ³Novartis, Basel, Switzerland, ⁴University of Oxford, Oxford, United Kingdom

YANG SUN  
University of Oxford
Oxford, Oxfordshire

Bernd Taschler  
University of Oxford
Goring-on-Thames, South Oxfordshire

Laura Gaetano  
Novartis
Basel, Switzerland

Thomas Nichols  
University of Oxford
Oxford, United Kingdom

Dieter Haering  
Novartis
Basel, Switzerland

Habib Ganjgahi  
University of Oxford
Oxford, United Kingdom

Multiple sclerosis (MS) is a chronic and ultimately debilitating disease of the central nervous system that affects approximately 2.5 million individuals worldwide[5]. The accumulated damage from past and ongoing neuroinflammation appears as hyperintense areas in T2-weighted MRI, known as T2 lesions. The count and volume of T2 lesions are an important prognostic factor for MS monitoring and progression[3]. The best detection approach is manual segmentation but it is costly, subject to intra/inter‐rater variability, hence there has been growing interest to develop automated tools. Deep learning has been shown to consistently outperform traditional image processing algorithms[4] however class imbalance, the unequal distribution of voxels among foreground (lesion) and background (non-lesion) classes, leads to poor generalizability and inflated false negatives (FN). In this work, we (a)evaluate the performance of different loss functions designed to mitigate the impact of class imbalance, and (b)assess uncertainty using conformal risk control.

Data resampling and loss function modifications are the common methods to handle class imbalance in image segmentation. It can be addressed by oversampling the minority class or undersampling the majority class during training. Besides, using loss functions that assign higher weights to the minority class and lower weights to the majority class can make the model more robust to class imbalance. The conventional loss functions for image segmentation are Binary Cross-Entropy (BCE), which quantifies the discrepancy between true lesion mask and its prediction, Dice loss, which is a harmonic mean of precision and recall; an alternative is Tversky loss, which is a generalised version of Dice loss that gives different penalty terms to false positives (FP) and FN.
In this work, we evaluate the impact of undersampling the majority class (only using lesion slices to train) and the use of combined loss functions to address class imbalance for lesion segmentations: BCE with and without class weights plus Dice; BCE with and without class weights plus Tversky loss with FP and FN penalties of 0.3 and 0.7, respectively; and WBCE plus Tversky loss with FP and FN penalties of 0.7 and 0.3. When used, the weights are proportional to the inverse of class frequency.
Conformal inference is an approach to get uncertainty quantification from an arbitrary predictive model. To obtain a conformal risk guarantee in our segmentation task, we manipulate the threshold value to regulate the fraction of misssegmentations [1]. Here, we use the trained model with weighted BCE (WBCE) and Tversky loss (α=0.7&β=0.3) and we target a false negative rate (FNR) less than 0.1 (recall over 0.9).
We use the NO.MS dataset [2] for evaluating the performance of different methods. It is currently the largest and most comprehensive clinical trial dataset in MS with more than 35,000 MS patients and over 200,000 MRI scans. We use 11,707 scans, randomly split into 10,571 for training and 1,136 for testing, where the ground truth lesion masks are obtained either manually or semi-automatically.

Our experiments show that resampling has the worst performance in terms of generalizability to test sets, F1 and precision. In addition, up-weighting lesion voxels and down-weighting background voxels in the BCE and Dice improve recall but suffer from inflated FPs. However, Tversky loss that penalizes FPs improves precision while WBCE keeps good recall (see Figure 1). Conformal risk control was used to guarantee FNR is controlled at 0.1 level with high probability.

We have shown WBCE and Tversky loss function improves T2 lesion segmentation by alleviating the class imbalance problem. Tuning the hyperparameters in Tversky can facilitate the efficient balance between recall and precision. Implementing conformal risk control ensures the FNR is managed and consequently improves the recall.

Methods Development ¹

Segmentation and Parcellation ²

MRI

White Matter

Other - Deep Learning, Lesion Segmentation