[论文解读] 3D U$^2$-Net: A 3D Universal U-Net for Multi-Domain Medical Image Segmentation

Chao Huang, Hu Han+3 more|arXiv (Cornell University)|Sep 4, 2019

Advanced Neural Network Applications参考文献 18被引用 9

一句话总结

本文提出3D U²-Net，一种用于多领域医学图像分割的通用3D U-Net架构，通过可分离卷积共享跨领域特征，同时利用通道卷积学习领域特定模式。该模型仅使用参数量的1%即可达到与任务专用网络相当的分割精度，并且仅增加0.3%的参数即可有效适应新器官。

ABSTRACT

Fully convolutional neural networks like U-Net have been the state-of-the-art methods in medical image segmentation. Practically, a network is highly specialized and trained separately for each segmentation task. Instead of a collection of multiple models, it is highly desirable to learn a universal data representation for different tasks, ideally a single model with the addition of a minimal number of parameters steered to each task. Inspired by the recent success of multi-domain learning in image classification, for the first time we explore a promising universal architecture that handles multiple medical segmentation tasks and is extendable for new tasks, regardless of different organs and imaging modalities. Our 3D Universal U-Net (3D U$^2$-Net) is built upon separable convolution, assuming that {\it images from different domains have domain-specific spatial correlations which can be probed with channel-wise convolution while also share cross-channel correlations which can be modeled with pointwise convolution}. We evaluate the 3D U$^2$-Net on five organ segmentation datasets. Experimental results show that this universal network is capable of competing with traditional models in terms of segmentation accuracy, while requiring only about $1\%$ of the parameters. Additionally, we observe that the architecture can be easily and effectively adapted to a new domain without sacrificing performance in the domains used to learn the shared parameterization of the universal network. We put the code of 3D U$^2$-Net into public domain. \url{https://github.com/huangmozhilv/u2net_torch/}

研究动机与目标

• 解决为每个医学图像分割任务训练独立深度学习模型所带来的效率低下与难以扩展的问题。
• 使单一神经网络能够泛化于多种解剖结构和成像模态（如CT、MRI）之间，而无需从头开始重新训练。
• 开发一种参数高效的架构，在多个领域中保持高分割精度的同时最小化模型规模。
• 通过极少的额外参数和对现有任务无性能下降的方式，实现对新器官（如脾脏）的迁移学习。
• 通过将图像分类中的成功经验扩展到密集预测任务，弥合医学图像分割中多领域学习的差距。

提出的方法

• 设计一种基于3D U-Net的通用架构，采用可分离卷积：每个领域使用通道卷积捕捉领域特定的空间模式，共享逐点卷积用于建模跨通道的共享相关性。
• 在适应新领域时引入领域特定参数（如输入层、通道卷积），同时冻结共享组件（如逐点卷积）。
• 使用轮询采样策略在所有领域间进行训练，以确保优化过程中各领域贡献均衡。
• 应用训练损失的指数移动平均以调度学习率衰减，在损失停滞时将其降低5倍。
• 使用ADAM优化器配合权重衰减和初始为3×10⁻⁴的学习率，对五个器官分割数据集进行端到端训练。
• 通过冻结共享层并仅微调领域特定模块（如新输入层和通道卷积）实现对新领域的零样本适应。

实验结果

研究问题

• RQ1 单一深度神经网络是否能在不同器官和成像模态（如心脏、肝脏、前列腺）的多个医学图像领域中实现具有竞争力的分割性能？
• RQ2 具有共享与领域特定组件的通用架构是否在分割精度和参数效率方面优于独立模型和仅共享模型？
• RQ3 该通用模型在仅增加极少额外参数且不降低现有任务性能的前提下，对未见过的新器官（如脾脏）的泛化能力如何？
• RQ4 使用可分离卷积（通道卷积用于领域特定特征，逐点卷积用于共享特征）如何提升多领域分割中的模型效率与性能？
• RQ5 在将通用模型扩展至新医学影像任务时，模型规模、参数共享与分割精度之间的权衡关系如何？

主要发现

• 3D U²-Net在五个基础领域（心脏、肝脏、海马体、前列腺、胰腺）上实现了83.10%的平均Dice分数，与独立模型的84.28%非常接近，而参数量仅为1.7M对比126.7M。
• 该通用模型在所有基础领域上均优于共享模型（82.56%平均Dice），尤其在前列腺的外周带（PZ）和移行带（TZ）分割中表现更优（68.50% vs. 68.38%）。
• 在适配新领域（脾脏）时，通用模型实现了91.60%的Dice分数，仅增加0.1M参数，优于共享模型（90.67%），且仅需全独立模型参数量的0.3%。
• 该通用模型的参数量约为独立模型的1%（1.7M vs. 126.7M），展现出极高的参数效率。
• 在适配新器官后，模型在所有原始领域上仍保持强劲性能，证实共享表征稳定且具备泛化能力。
• 通过冻结共享组件并仅训练领域特定模块，该架构实现了有效的零样本适应，性能下降极小，且具备高度可扩展性。