【论文阅读笔记】Prompt Tuning for Parameter-efficient Medical Image Segmentation

Fischer M, Bartler A, Yang B. Prompt tuning for parameter-efficient medical image segmentation[J]. Medical Image Analysis, 2024, 91: 103024. 【开源】

【核心思想】

本文的核心思想是提出了一种用于医学图像分割的参数高效的提示调整（Prompt Tuning）方法。这种方法基于预训练的神经网络，通过插入可学习的提示令牌（prompt tokens）来适应新的下游任务，例如在语义分割中识别新的类别。这种方法的关键在于，它允许在不改变预训练模型主体（backbone）的情况下，通过少量参数的调整来适应新任务。文章中提出了一个名为PUNet（Prompt-able UNet）的架构，这是一个在预训练后固定不变，但通过类依赖的可学习提示令牌在整个网络中可调整的结构。PUNet通过一种基于在线生成原型的密集自监督方案进行预训练，这种方案结合了学生-教师模型和同时进行的分割损失。研究表明，这种方法能够在CT成像数据集上有效地缩小全面微调模型与参数高效适应模型之间的性能差距。此外，文章还探讨了不同的训练阶段、提示令牌的使用方式（如二元预测或多类预测）以及不同的网络架构变体对模型性能的影响。实验结果表明，这种提示调整方法在医学图像分割任务中，即使在标注数据稀缺的情况下，也能达到令人满意的性能，同时大大减少了所需调整的参数数量。

主要贡献如下：

1. 提示可调的UNet（PUNet）架构的提出：文章提出了一种新的神经网络架构，即提示可调的UNet（PUNet）。这种架构在预训练后固定不变，但可以通过类依赖的可学习提示令牌在整个网络中进行调整。这种设计使得模型能够在保持预训练网络主体不变的情况下，通过调整相对较少的参数来适应新的下游任务。
2. 密集自监督预训练方案：文章提出了一种基于在线生成原型的密集自监督预训练方案。这种方案结合了学生-教师模型和同时进行的分割损失，旨在建立解剖学表示，同时避免对比学习中吸引和排斥的硬性分离。
3. 参数高效的适应性：研究表明，通过使用提示令牌，PUNet能够以参数高效的方式适应新的分割任务。这意味着，相比于传统的完全微调方法，PUNet只需调整相对较少的参数即可达到类似的性能。