【论文阅读笔记】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只需调整相对较少的参数即可达到类似的性能。
4. 广泛的实验验证：文章通过在两个医学成像数据集上的实验，验证了PUNet架构的有效性。实验结果表明，即使在标注数据稀缺的情况下，这种方法也能达到令人满意的性能。
5. 对不同训练策略和网络变体的探讨：文章不仅提出了PUNet架构，还探讨了不同的训练阶段、提示令牌的使用方式（如二元预测或多类预测）以及不同的网络架构变体对模型性能的影响。

【模型结构】

1.提示可调的UNet (PUNet) 架构：PUNet是一种为医学图像分割设计的网络架构，其核心特点是在网络中引入了可学习的提示令牌（prompt tokens），以实现对预训练模型的高效微调。以下是对这一部分的详细解读：

• 架构设计：PUNet基于流行的UNet架构，但进行了重要的修改，使其能够整合额外的类依赖提示令牌。这些提示令牌被视为一组可学习的指令，它们聚集了所有任务依赖的信息，以实现参数高效的微调。在PUNet中，预训练后的主干网络（backbone model）被冻结，而提示令牌则在整个网络中进行调整。

  • Prompt-able SWin (PSWin) blocks:
    • 结合移位窗口和提示令牌：PSWin块是一种特殊的网络块，它结合了移位窗口（SWin）的概念和可调节的提示令牌。这种设计允许网络在处理图像内容的同时，考虑到与特定任务相关的提示信息。
    • 窗口化内容的处理：在PSWin块中，输入图像被分割成较小的区块（窗口化的内容），这些内容随后与提示令牌一起被送入变压器块中进行处理。
    • 提示令牌的作用：提示令牌为注意力层提供了关于目标任务的额外信息。在PSWin块中，这些令牌与每个窗口化内容一起被处理，从而使注意力层能够根据任务相关的信息调整其对图像内容的处理方式。
    • 注意力机制的调整：PSWin块通过计算窗口化内容和提示令牌之间的相似性分数来调整注意力权重，确保网络能够根据提示令牌中的信息调整其对图像内容的处理。
    • 灵活性和适应性：PSWin块的设计提供了极大的灵活性和适应性，使得PUNet能够有效地适应不同的医学图像分割任务。通过在网络的不同层级中插入PSWin块，可以实现对编码图像内容的深度和细粒度调整。
    • 记忆效率：PSWin块在设计时也考虑到了记忆效率。通过限制自注意力到局部非重叠窗口，并在后续的块中进行移位，PSWin块能够在保持线性复杂度的同时，有效地应用注意力机制。
  • Heterogeneous Bias Scores:
    • 异质偏差得分的概念：异质偏差得分是一种机制，用于在网络的不同部分引入额外的偏差，以增强模型处理不同类型数据的能力。这些偏差得分允许网络更好地区分和处理来自不同源的信息，例如图像内容和提示令牌。
    • 增强注意力机制：通过在注意力层中引入异质偏差得分，网络能够更有效地结合来自图像和提示令牌的信息。这种机制有助于提高模型对于不同任务特征的敏感性和适应性。
    • 提高模型的灵活性：异质偏差得分的引入增加了模型的灵活性，使其能够更好地适应不同的任务和数据类型。这对于处理复杂的医学图像分割任务尤其重要，因为这些任务通常涉及到多种类型的图像特征和标签。
    • 优化任务适应性：异质偏差得分使得PUNet在适应新任务时更加高效，尤其是在处理具有不同特征和标签的医学图像时。这种优化的适应性对于提高分割精度和减少训练时间至关重要。
    • 提升分割性能：通过更好地理解和处理图像内容与提示令牌之间的关系，异质偏差得分有助于提升医学图像分割的整体性能，特别是在精确度和鲁棒性方面。
  • Cosine Similarity Aggregation:
    • 余弦相似度聚合的目的：余弦相似度聚合是一种机制，用于评估和聚合网络中不同部分的特征相似度。这种方法特别用于处理提示令牌和图像特征之间的关系，以提高医学图像分割的准确性。
    • 工作原理：在PUNet中，余弦相似度聚合通过计算提示令牌和图像特征之间的余弦相似度来工作。这种方法允许模型评估不同类别的特征与提示令牌之间的相似性，并据此进行有效的分类。
    • 提高分割精度：通过使用余弦相似度聚合，PUNet能够更准确地将图像区域分配给正确的类别。这种精细的相似度评估对于提高医学图像分割的精度至关重要。
    • 增强模型适应性：余弦相似度聚合增强了模型对不同任务特征的适应性。这意味着PUNet可以更有效地处理多样化的医学图像数据，适应不同的分割任务。
    • 优化特征表示：通过聚合相似度信息，余弦相似度聚合有助于优化网络中的特征表示。这种优化的特征表示对于提高分割任务的整体性能非常重要。

• 提示令牌的作用：提示令牌在网络中起到关键作用，它们使得模型能够根据不同的任务进行适应。每个任务都有一组新的提示令牌，其中每个类的子集代表二元或多类情况。这意味着，通过更换提示令牌，PUNet可以灵活地适应不同的分割任务。

• 深度集成：作者在PUNet中深度集成了提示令牌，这允许在网络中对编码的图像内容进行中间级别的调整。特别是，注意力层提供了一种结构化的方式来组合和处理异质编码的图像和提示信息。

  

  生成嵌入，利用对比学习，其中解剖学上相似的区域彼此靠近地表示。它结合了动量模型和 EMA 更新的教师和学生，将两名学生纳入其中，其中一名学生处理比教师输入 更小的输入，第二名学生 进行更严格的裁剪，以强化鲁棒嵌入，两个学生都共享权重。

  • Contrastive Prototype Assignments (CPA):
    • 原型的生成：CPA方法首先涉及生成一组原型，这些原型是数据特征的代表性集合。在医学图像分割的上下文中，这些原型可以被理解为代表不同解剖结构或图像模式的特征集合。
    • 特征与原型的对比：网络通过将输入图像的特征与这些原型进行对比来学习。这种对比不是简单的匹配，而是通过计算特征与原型之间的相似度来进行的。通常，这种相似度是通过余弦相似度或其他相关度量来计算的。
    • 自监督学习：在CPA中，网络被训练以最大化输入特征与相应原型之间的相似度。这种方法是自监督的，因为它不依赖于外部标注，而是依赖于数据本身的结构和内在模式。
    • 对比学习：CPA利用对比学习的原理，即通过比较不同的特征表示来学习区分不同的数据点。在这种情况下，网络被训练以区分不同的原型，并将图像特征正确地与这些原型对齐。
    • 促进特征区分性：通过这种方法，网络学习生成更加区分性和信息丰富的特征表示。这对于后续的图像分割任务至关重要，因为它提供了更精确的特征来指导分割过程。
    • 适应性和泛化：CPA通过强调数据内在结构的学习，提高了模型的适应性和泛化能力。这意味着模型能够更好地处理在训练数据中未见过的新图像或结构。
  • Online Prototype Generation：
    • 在线原型生成的目的：在线原型生成旨在实时创建和更新代表数据特征的原型。这些原型用于指导网络学习过程中的特征提取和表示。
    • 动态原型更新：与传统的静态原型不同，在线原型生成涉及到在训练过程中不断更新原型。这意味着原型能够适应训练数据中出现的新模式和变化，从而更准确地反映数据的当前状态。
    • 自监督学习的加强：通过在线更新原型，网络能够更有效地进行自监督学习。这种动态的学习过程有助于提高特征表示的质量和模型对新数据的适应能力。
    • 实时特征对齐：在线原型生成允许模型实时地将输入特征与最新的原型进行对齐和比较。这种实时对齐机制对于处理复杂和变化多端的医学图像尤其重要。
    • 提高模型的灵活性和准确性：由于原型能够动态更新，模型在处理多样化的医学图像时更加灵活和准确。这对于提高图像分割的性能至关重要。
    • 减少对大量标注数据的依赖：在线原型生成作为一种自监督方法，减少了对大量标注数据的依赖。这使得PUNet在标注数据有限的情况下仍能有效学习。

• 提示令牌的插入：在PUNet中，提示令牌被深入地插入到网络中，这允许对编码的图像内容在网络的不同层次上进行调整。这种深度集成的方法有助于更有效地适应下游任务。

• 记忆效率：为了提高记忆效率，PUNet架构中包括了记忆效率较高的移位窗口（SWin）注意力块。这些块在网络的编码器和解码器中被使用，以进一步节省内存。

• 适应性：PUNet的设计允许网络通过非冻结的（可学习的）提示令牌进行适应，这对于医学成像数据的分割任务来说是足够的。这种方法在保持主干网络任务不变的同时，实现了对特定任务的高效适应。