[论文解读] Towards Generic Semi-Supervised Framework for Volumetric Medical Image Segmentation

Haonan Wang, Xiaomeng Li|arXiv (Cornell University)|Oct 17, 2023

Generative Adversarial Networks and Image Synthesis被引用 16

一句话总结

引入 Aggregating & Decoupling (A&D)，一个通用的半监督框架，整合 SSL、类别不平衡的 SSL、UDA 与 SemiDG，用于体积医学影像分割，使用基于扩散的聚合编码器和解耦解码器，以抑制过拟合并捕捉分布不变特征。

ABSTRACT

Volume-wise labeling in 3D medical images is a time-consuming task that requires expertise. As a result, there is growing interest in using semi-supervised learning (SSL) techniques to train models with limited labeled data. However, the challenges and practical applications extend beyond SSL to settings such as unsupervised domain adaptation (UDA) and semi-supervised domain generalization (SemiDG). This work aims to develop a generic SSL framework that can handle all three settings. We identify two main obstacles to achieving this goal in the existing SSL framework: 1) the weakness of capturing distribution-invariant features; and 2) the tendency for unlabeled data to be overwhelmed by labeled data, leading to over-fitting to the labeled data during training. To address these issues, we propose an Aggregating & Decoupling framework. The aggregating part consists of a Diffusion encoder that constructs a common knowledge set by extracting distribution-invariant features from aggregated information from multiple distributions/domains. The decoupling part consists of three decoders that decouple the training process with labeled and unlabeled data, thus avoiding over-fitting to labeled data, specific domains and classes. We evaluate our proposed framework on four benchmark datasets for SSL, Class-imbalanced SSL, UDA and SemiDG. The results showcase notable improvements compared to state-of-the-art methods across all four settings, indicating the potential of our framework to tackle more challenging SSL scenarios. Code and models are available at: https://github.com/xmed-lab/GenericSSL.

研究动机与目标

• 解决体积医学影像分割中 SSL 方法在域和分布转移下失败的空缺（SSL、UDA、SemiDG）的问题。
• 开发一个在多个域之间学习分布不变特征的统一框架。
• 通过解耦训练路径和基于扩散的聚合，减轻对标注数据的过拟合以及域/类别偏差。

提出的方法

• 提出 Aggregating & Decoupling (A&D) 框架，包含 Aggregating 阶段和 Decoupling 阶段。
• 使用 Diffusion V-Net (Diff-VNet) 编码器，通过对增强体积进行扩散去噪来从多域数据学习分布不变特征。
• 在编码前应用 Sampling-based Volumetric Data Augmentation (SVDA) 以扩大分布多样性。
• 将训练解耦为三种解码器以处理有标注数据：一个基于扩散的去噪器（域无偏）和一个难度感知的 V-Net（类别无偏）；以及一个由伪标签监督的无标注数据解码器。
• 通过使用 Reparameterize & Smooth (RS) 策略，将域无偏（扩散）与类别无偏（难度感知）的预测进行集成，生成高质量的伪标签。
• 通过损失项优化训练：L_deno（Dice+CE）、L_diff（难度感知加权）、以及用于带蒸馏机制的无监督训练的 L_u。

实验结果

研究问题

• RQ1 一个通用的 SSL 框架是否能够在体积医学影像分割中同时处理 SSL、类别不平衡的 SSL、UDA 与 SemiDG？
• RQ2 如何在跨域学习分布不变特征，以及如何避免对标注数据过拟合？
• RQ3 在域和类别漂移下，将有标注与无标注数据训练解耦是否能提升性能？
• RQ4 基于扩散的聚合和所提 RS 策略对伪标签质量和分割准确性的影响是什么？
• RQ5 在标准基准测试中，SSL、类别不平衡的 SSL、UDA 以及 SemiDG 的经验收益是多少？

主要发现

• 在代表性数据集上实现了 SSL、类别不平衡的 SSL、UDA 和 SemiDG 的最新研究结果。
• 显示显著的 Dice 提升，包括 Synapse 上 12.3 Dice 点（20% 有标记）以及 MMWHS MR→CT 上 8.5 Dice 点。
• A&D 在报道的实验中在所有四种设置下均优于竞争方法。
• 消融研究表明扩散和 SVDA 对 UDA 至关重要，而难度感知重新加权对类别失衡（IBSSL）特别有益。
• 统一的代码库通过调整参数即可在单一脚本中运行任何任务。