[论文解读] MixSearch: Searching for Domain Generalized Medical Image Segmentation Architectures

Luyan Liu, Zhiwei Wen+7 more|arXiv (Cornell University)|Feb 26, 2021

Advanced Neural Network Applications参考文献 50被引用 5

一句话总结

MixSearch 提出了一种新颖的神经架构搜索（NAS）框架，通过混合多领域、多任务的医学影像构建复合医学数据集，以训练一种通用的U型分割架构。通过采用具有可学习特征融合的编织式编码器-解码器搜索空间，MixSearch 发现的架构在未见数据集上优于当前最先进的人工设计模型——包括注意力U-Net和U-Net变体，展现出显著提升的Dice分数和Jaccard分数，体现出强大的域泛化能力。

ABSTRACT

Considering the scarcity of medical data, most datasets in medical image analysis are an order of magnitude smaller than those of natural images. However, most Network Architecture Search (NAS) approaches in medical images focused on specific datasets and did not take into account the generalization ability of the learned architectures on unseen datasets as well as different domains. In this paper, we address this point by proposing to search for generalizable U-shape architectures on a composited dataset that mixes medical images from multiple segmentation tasks and domains creatively, which is named MixSearch. Specifically, we propose a novel approach to mix multiple small-scale datasets from multiple domains and segmentation tasks to produce a large-scale dataset. Then, a novel weaved encoder-decoder structure is designed to search for a generalized segmentation network in both cell-level and network-level. The network produced by the proposed MixSearch framework achieves state-of-the-art results compared with advanced encoder-decoder networks across various datasets.

研究动机与目标

• 为解决现有医学图像分割模型在单一、小规模数据集上训练时泛化能力不足的问题。
• 通过构建大规模、复合的训练分布，克服医学图像分析中的数据稀缺和域偏移挑战。
• 开发一个端到端的NAS框架，以发现适用于多领域医学图像分割的鲁棒、可泛化的U型架构。
• 评估所学架构在来自不同成像模态和解剖区域的完全未见数据集上的泛化能力。

提出的方法

• 通过混合多个医学图像分割任务和领域的样本（包括真实和增强的复合图像）构建复合数据集。
• 提出一种新颖的编织式编码器-解码器架构，包含可学习的水平和垂直特征融合模块，实现在单元级和网络级的联合搜索。
• 采用受DARTS启发的可微架构搜索策略，在复合数据集上优化架构。
• 使用数据混合策略，如MixUp，混合系数μ=0.5，以提升域多样性与分布鲁棒性。
• 通过在解码器多个层级设置深层监督头，实现多尺度监督，以增强特征学习和空间精度。
• 在结合ISIC、LiTS、KiTS和ETIS的复合数据集上执行架构搜索，以模拟跨域泛化。

实验结果

研究问题

• RQ1 通过混合多领域医学影像构建的复合数据集，能否提升所学分割架构的泛化能力？
• RQ2 具有可学习特征融合的编织式编码器-解码器搜索空间，是否能在多样化医学成像领域中实现更优的性能与鲁棒性？
• RQ3 在复合数据集上进行NAS，能否发现比人工设计模型在未见数据集上泛化能力更强的架构？
• RQ4 与专用型和联合型NAS策略相比，所提出的MixSearch框架在泛化能力和性能方面表现如何？

主要发现

• 在未见的ETIS数据集上，MixSearch的Dice分数比Attention U-Net高出1.0%，Jaccard相似度高出1.5%。
• 在LiTS数据集上，尽管MR U-Net相比U-Net++仅有0.01%的Dice优势，MixSearch仍以显著优势超越了该最佳人工设计模型。
• 在KiTS数据集上，MixSearch取得94.53%的Dice分数，与R2T U-Net的94.64%相比仅存在微小差距，展现出接近SOTA的性能。
• MixSearch所学架构的泛化能力优于UnionSearch，后者在各源域上出现过拟合，甚至表现不如基础U-Net。
• MixSearch中的分步特征聚合显著提升了边缘检测与分割掩膜质量，Jaccard分数从X^{0,5}提升至X^{0,9}。
• 最终输出X^{0,7}具备高精度，可作为剪枝后的推理头使用，在不损失性能的前提下降低计算量。