Staged Voxel-Level Deep Reinforcement Learning for 3D Medical Image Segmentation with Noisy Annotations

TL;DR

This work tackles the challenge of 3D medical image segmentation under noisy annotations. It introduces SVL-DRL, an end-to-end framework that treats each voxel as an autonomous agent within a staged reinforcement learning paradigm, augmented by a voxel-level A3C module (vA3C) and a tailored action space with a Dice-plus-spatial-continuity reward. The method achieves state-of-the-art performance across LA, Pancreas-CT, and BraTS 2021 under synthetic noise patterns (SFDA-induced and morphological transformations), with average improvements exceeding 3% in Dice and IoU. By directly refining input states and labels through reinforcement learning rather than filtering data, SVL-DRL enhances robustness to annotation noise and holds promise for more reliable clinical segmentation pipelines.

Abstract

Deep learning has achieved significant advancements in medical image segmentation. Currently, obtaining accurate segmentation outcomes is critically reliant on large-scale datasets with high-quality annotations. However, noisy annotations are frequently encountered owing to the complex morphological structures of organs in medical images and variations among different annotators, which can substantially limit the efficacy of segmentation models. Motivated by the fact that medical imaging annotator can correct labeling errors during segmentation based on prior knowledge, we propose an end-to-end Staged Voxel-Level Deep Reinforcement Learning (SVL-DRL) framework for robust medical image segmentation under noisy annotations. This framework employs a dynamic iterative update strategy to automatically mitigate the impact of erroneous labels without requiring manual intervention. The key advancements of SVL-DRL over existing works include: i) formulating noisy annotations as a voxel-dependent problem and addressing it through a novel staged reinforcement learning framework which guarantees robust model convergence; ii) incorporating a voxel-level asynchronous advantage actor-critic (vA3C) module that conceptualizes each voxel as an autonomous agent, which allows each agent to dynamically refine its own state representation during training, thereby directly mitigating the influence of erroneous labels; iii) designing a novel action space for the agents, along with a composite reward function that strategically combines the Dice value and a spatial continuity metric to significantly boost segmentation accuracy while maintain semantic integrity. Experiments on three public medical image datasets demonstrates State-of-The-Art (SoTA) performance under various experimental settings, with an average improvement of over 3\% in both Dice and IoU scores.

Figures (6)

• Figure 1: The overall framework of SVL-DRL.The SVL-DRL architecture is based on PA3C combining swin-unetr. $X^{(t)}$ is the temporary input in step $t$, $GT$ is the label. $a^{(t)}$ is sampled from the policy $\pi: a^{(t)} \sim \pi(a^{(t)} \mid s^{(t)})$.
• Figure 2: Visual comparison of different noisy labels on different datasets. Column 1: original images; Column 2: images with clean segmentation labels; Column 3: images with TC-Noise labels; Column 4: images with SFDA-Noise labels
• Figure 3: Visual comparison of segmentation results on the LA dataset under MT-noise label setting across different methods.
• Figure 4: Visual comparison of segmentation results on the Pancreas dataset under MT-noise label setting across different methods.
• Figure 5: Visual comparison of segmentation results on the BraTS 2021 dataset under MT-noise label setting across different methods.