Impact of Soft Segmentation Training on Medical Image Segmentation and Uncertainty Representation

Avec l'essor de l'apprentissage profond, une quantité croissante de modèles sont développés pour le domaine médical afin d'automatiser les tâches fastidieuses et de réduire les erreurs médicales causées par l'homme. Cependant, en raison de leur impact potentiel sur la vie et la santé humaine, des préoccupations éthiques grandissantes font jour concernant la fiabilité et la transparence des modèles d'apprentissage profond. Une solution potentielle à ces problèmes est d'entraîner des modèles générant des prédictions calibrées avec une représentation fidèle de l'incertitude. De cette manière, les prédictions les plus susceptibles d'être incorrectes ou de donner lieu à des désaccords entre les experts peuvent être isolées et corrigées. Cependant, les réseaux neuronaux de segmentation modernes sont généralement trop confiants, c'est-à-dire qu'ils expriment une grande certitude même pour les prédictions erronées, et ne tiennent pas compte de considérations importantes en imagerie médicale telles que l'e˙et de volume partiel, la variabilité inter-expert ou la représentation de l'incertitude en raison de la faible qualité des images ou du manque de données. Ceci est dû à la nature binaire de la segmentation qui est considérée comme une tâche de classification où chaque voxel se voit attribuer une valeur de 0 ou 1. Dans la première partie de ce travail, nous proposons une méthode appelée SoftSeg qui traite la segmentation comme une tâche de régression afin d'encourager la représentation des informations sur les volumes partiels, la variabilité inter-expert et la représentation de l'incertitude. L'approche de segmentation non-binaire vise à réduire la confiance excessive du modèle. Trois caractéristiques principales définissent SoftSeg par rapport aux modèles de seg-mentation conventionnels : (i) la préservation du caractère non-binaire, i.e., entre 0 et 1, des segmentations utilisées pour l'entraînement après le traitement et l'augmentation des don-nées, (ii) une fonction d'activation finale linéaire normalisée pour éviter la perte d'information contrairement aux fonctions sigmoïde ou softmax non linéaires, et (iii) l'utilisation d'une fonction de perte de régression plutôt que de classification comme Dice ou d'entropie croisée. Nous avons exploré ces nouvelles fonctionnalités et évalué l'impact de chacune d'entre elles lors d'une étude d'ablation. La combinaison de ces trois nouvelles caractéristiques a permis d'obtenir de meilleures performances de segmentation sur trois ensembles de données de seg-mentation publiquement accessibles : matière grise de la moelle épinière, lésions de sclérose en plaques du cerveau et tumeur du cerveau. Dans un deuxième article, trois méthodes de fusion d'annotations d'expert, soit STAPLE, moyennage et l'échantillonnage aléatoire (c'est-à-dire sans fusion), pairées à SoftSeg ou à un entraînement conventionnel, ont été comparées. Les approches ont été étudiées sur deux ensembles de données avec respectivement quatre et sept annotations d'évaluateurs pour chaque image : segmentation de la matière grise et blanche de la moelle épinière et lésions de la sclérose en plaques du cerveau. La préservation de l'incertitude due au désaccord entre les évaluateurs, la calibration des prédictions, la qualité visuelle et les performances de segmentation ont été évaluées. Bien qu'il n'y ait pas eu de consensus entre les ensembles de données en ce qui concerne la meilleure méthode de fusion des annotations d'expert, les résultats étaient équivoques en ce qui concerne le type d'entraînement. Nos résultats indiquent que SoftSeg a produit une prédiction avec une meilleure calibration ainsi qu'une préservation de la variabilité inter-expert accrue, et ce, avec une performance de segmentation améliorée, ou minimalement équivalente. Toutes les approches étudiées dans ce travail ont été répétées 10 à 40 fois avec des séparations aléatoires des données pour éviter un biais au niveau des données de test et garantir des di˙érences statistiques (valeur p < 0,05). Toutes les recherches e˙ectuées dans ce projet ont été réalisées et rendues accessibles via le projet en libre accès ivadomed (https://ivadomed. org).

Abstract

With the rise of deep learning, an increasing amount of models are being developed and researched for the medical field to automate tedious tasks and mitigate medical errors. How-ever, due to the sensitive nature of medical tasks and their impact patient's health, increasing ethical concerns are arising regarding reliability and transparency of deep learning models. A potential avenue to address these issues is to have calibrated predictions with truthful uncer-tainty representation. Reliable uncertainty representation help identify predictions that are prone to model failure or inter-rater disagreement. However, modern segmentation neural networks are usually overconfident, i.e., express a high certainty even for incorrect predic-tions, and disregard important considerations in medical imaging such as partial volume e˙ect, inter-rater variability, or uncertainty representation due to low image quality or lack of data. This is partly due to the inherent binary nature of segmentation that is considered a classification task where each voxel is attributed a value of 0 or 1. In the first part of this work, we propose a method named SoftSeg that treats segmen-tation as a regression task to encourage the representation of partial volume information, inter-rater variability, and uncertainty. The soft segmentation approach aims at mitigating overconfidence. Three main features characterize SoftSeg compared with the conventional segmentation models: (i) preservation of soft input labels following data processing and aug-mentation, (ii) a normalized linear final activation to avoid information loss instead of the non-linear sigmoid or softmax, and (iii) the use of a regression loss function rather than the classification Dice or cross-entropy loss. We explored these new features and evaluated the impact of each feature through an ablation study. The combination of these three new fea-tures resulted in better segmentation performance on three publicly available segmentation datasets: spinal cord gray matter, brain multiple sclerosis lesions, and brain tumor. In a second article, three label fusion methods, STAPLE, average, and random sampling (no fusion), paired with SoftSeg or a conventional training framework, were compared. The approaches were studied on two datasets with respectively four and seven rater annotations for each image: spinal cord gray and white matter segmentation and brain multiple sclerosis lesions. The uncertainty preservation due to inter-rater disagreement, the calibration of predictions, the visual predictions, and the segmentation performance were evaluated. While there was no consensus between datasets in terms of the best label fusion method, results were equivocal regarding the training framework. Our results indicate that SoftSeg yielded prediction with better calibration and inter-rater variability preservation with higher, or minimally equivalent, segmentation performance. All the approaches studied in this work were repeated 10 to 40 times with varying random seeds to avoid data splitting biases and ensure statistical di˙erences (p-value < 0.05). All the research done in this project was developed and made accessible via the open-source project ivadomed (https://ivadomed.org).