[论文解读] ProSFDA: Prompt Learning based Source-free Domain Adaptation for Medical Image Segmentation

Shishuai Hu, Zehui Liao+1 more|arXiv (Cornell University)|Nov 21, 2022

Domain Adaptation and Few-Shot Learning被引用 11

一句话总结

该论文提出ProSFDA，一种基于提示学习的无源域自适应方法，用于医学图像分割，通过学习一种领域感知的视觉提示来优化目标域图像，从而减少领域差异。通过对齐批量归一化统计量并利用风格增强特征对齐来强化特征紧凑性，ProSFDA在多领域基准测试中实现了与无监督域自适应方法相当的性能。

ABSTRACT

The domain discrepancy existed between medical images acquired in different situations renders a major hurdle in deploying pre-trained medical image segmentation models for clinical use. Since it is less possible to distribute training data with the pre-trained model due to the huge data size and privacy concern, source-free unsupervised domain adaptation (SFDA) has recently been increasingly studied based on either pseudo labels or prior knowledge. However, the image features and probability maps used by pseudo label-based SFDA and the consistent prior assumption and the prior prediction network used by prior-guided SFDA may become less reliable when the domain discrepancy is large. In this paper, we propose a extbf{Pro}mpt learning based extbf{SFDA} ( extbf{ProSFDA}) method for medical image segmentation, which aims to improve the quality of domain adaption by minimizing explicitly the domain discrepancy. Specifically, in the prompt learning stage, we estimate source-domain images via adding a domain-aware prompt to target-domain images, then optimize the prompt via minimizing the statistic alignment loss, and thereby prompt the source model to generate reliable predictions on (altered) target-domain images. In the feature alignment stage, we also align the features of target-domain images and their styles-augmented counterparts to optimize the source model, and hence push the model to extract compact features. We evaluate our ProSFDA on two multi-domain medical image segmentation benchmarks. Our results indicate that the proposed ProSFDA outperforms substantially other SFDA methods and is even comparable to UDA methods. Code will be available at \url{https://github.com/ShishuaiHu/ProSFDA}.

研究动机与目标

• 解决在无法访问源数据的情况下，将预训练模型部署到新领域时医学图像分割中的领域差异问题。
• 克服基于伪标签和先验引导的无源域自适应方法在大领域偏移下预测不可靠的局限性。
• 开发一种仅利用预训练源模型和无标签目标数据来提升适应性能的方法。
• 通过引入可学习的领域感知提示和风格增强特征对齐，提升特征紧凑性和领域对齐性。
• 在无源数据和隐私限制下，实现与无监督域自适应（UDA）方法相当的性能。

提出的方法

• 引入一种领域感知的视觉提示，作为可学习的图像嵌入，添加到目标域图像中，以模拟源域特征。
• 通过最小化源模型与在提示增强目标图像上微调后的模型之间批量归一化（BN）统计量差异的统计对齐损失，来优化该提示。
• 利用优化后的提示生成在统计上更接近源域分布的优化目标图像。
• 在优化后的目标图像与其风格增强版本之间应用特征对齐，以增强特征紧凑性并减少实例级差异。
• 在特征对齐阶段联合优化模型泛化能力和领域不变性，结合分割损失与特征对齐损失。
• 利用冻结的源模型在提示增强的目标图像上生成预测，实现在不重新训练源数据情况下的知识蒸馏。

实验结果

研究问题

• RQ1 提示学习能否有效估计并缓解无源域自适应中医学图像分割的领域差异？
• RQ2 在处理大领域偏移时，可学习的领域感知提示相较于伪标签法或先验引导方法有何优势？
• RQ3 与风格增强图像进行特征对齐在多大程度上能提升特征紧凑性和分割性能？
• RQ4 在统计对齐损失中，BN层和超参数（如α）的最佳配置是什么，以实现有效的领域差异估计？
• RQ5 尽管缺乏源数据，该方法能否实现与无监督域自适应（UDA）方法相当的性能？

主要发现

• ProSFDA在两个多领域医学图像分割基准上优于现有无源域自适应方法，在无源设置下达到最先进性能。
• 在'BASE2'基准中，该方法在OD分割任务上达到94.71 ± 0.01的平均Dice分数，在OC分割任务上达到85.33 ± 0.08，优于基线无源域自适应方法。
• 消融研究证实，在特征对齐阶段同时使用分割损失与特征对齐损失可获得最佳性能，相比仅使用分割损失提升1.36%。
• t-SNE可视化显示，提示增强目标图像的特征更接近源域特征，验证了领域差异估计的有效性。
• 统计对齐损失中的最优超参数α为0.01，当在损失计算中使用最多5个BN块时，性能趋于稳定。
• 该方法仅在具有批量归一化层的网络上有效，因其依赖BN统计量进行领域差异估计，表明当前存在架构限制。