Joint Class-Affinity Loss Correction for Robust Medical Image Segmentation with Noisy Labels

TL;DR

This paper introduces a novel joint class-affinity loss correction framework that leverages pixel-wise and pair-wise affinity relations to improve robustness in medical image segmentation with noisy labels.

Contribution

It proposes the JCAS framework combining pixel-wise and pair-wise supervision, along with a differentiated affinity reasoning module and a class-affinity loss correction strategy.

Findings

Significant improvement over existing methods on noisy medical image segmentation tasks.

Effective noise mitigation demonstrated on both synthetic and real-world noisy labels.

Achieves near upper bound performance with robust label correction.

Abstract

Noisy labels collected with limited annotation cost prevent medical image segmentation algorithms from learning precise semantic correlations. Previous segmentation arts of learning with noisy labels merely perform a pixel-wise manner to preserve semantics, such as pixel-wise label correction, but neglect the pair-wise manner. In fact, we observe that the pair-wise manner capturing affinity relations between pixels can greatly reduce the label noise rate. Motivated by this observation, we present a novel perspective for noisy mitigation by incorporating both pixel-wise and pair-wise manners, where supervisions are derived from noisy class and affinity labels, respectively. Unifying the pixel-wise and pair-wise manners, we propose a robust Joint Class-Affinity Segmentation (JCAS) framework to combat label noise issues in medical image segmentation. Considering the affinity in pair-wise manner incorporates contextual dependencies, a differentiated affinity reasoning (DAR) module is devised to rectify the pixel-wise segmentation prediction by reasoning about intra-class and inter-class affinity relations. To further enhance the noise resistance, a class-affinity loss correction (CALC) strategy is designed to correct supervision signals via the modeled noise label distributions in class and affinity labels. Meanwhile, CALC strategy interacts the pixel-wise and pair-wise manners through the theoretically derived consistency regularization. Extensive experiments under both synthetic and real-world noisy labels corroborate the efficacy of the proposed JCAS framework with a minimum gap towards the upper bound performance. The source code is available at \url{https://github.com/CityU-AIM-Group/JCAS}.