从视觉基础模型中提取任务特定知识以提升医学图像分割性能

医疗图像分割中基于任务特定知识蒸馏的轻量化模型训练方法，通过LoRA高效适配大预训练模型并利用扩散生成合成数据增强知识迁移效果，在五组医学数据集上显著优于传统知识蒸馏和自监督方法，尤其在小样本标注场景下提升显著。

　　

梁鹏晨|黄海山|蒲斌|曾全红|陈建国

上海大学微电子学院，中国上海 201800

摘要

大规模预训练模型，如视觉基础模型（VFMs），通过迁移通用知识，在各种下游任务中展现了令人印象深刻的性能，尤其是在目标数据有限的情况下。然而，它们的高计算成本以及自然图像和医学图像之间的领域差异限制了它们在医学分割任务中的实际应用。基于此，我们提出了以下重要问题：“在训练数据有限的情况下，我们如何有效地利用大型预训练VFMs的知识来训练一个小型、特定于任务的模型以进行医学图像分割？”为了解决这个问题，我们提出了一个新颖且可泛化的特定于任务的知识蒸馏框架。我们的方法首先在目标分割任务上对VFMs进行微调，以捕获特定于任务的特征，然后通过低秩适应（LoRA）技术来降低微调的计算成本。此外，我们还结合了由扩散模型生成的合成数据来扩充转移集，从而在数据有限的场景中提升模型性能。在五个医学图像数据集上的实验结果表明，我们的方法始终优于任务无关的知识蒸馏和自监督预训练方法（如MoCo v3和Masked Autoencoders (MAE)）。例如，在KidneyUS数据集上，我们的方法使用80个标记样本进行微调时，其Dice分数比任务无关的KD方法高出28%；在CHAOS数据集上，使用100个标记样本时，其性能比MAE提高了11%。这些结果突显了将特定于任务的KD引入医学分割中的创新性，并证明了其在标注稀缺情况下构建准确高效模型的实际意义。

引言

医学图像分割对于临床诊断、治疗计划和手术指导至关重要，因为它能够精确地划分解剖结构和病理区域[1]、[2]、[3]、[4]、[5]、[6]、[7]。尽管深度学习取得了进展，但由于医学标注数据的稀缺，开发高性能的分割模型仍然具有挑战性。标注医学图像既昂贵又耗时，需要专门的专家来细致地标记复杂结构[8]。这种稀缺性阻碍了有效模型的训练，因为这些模型通常需要大量的标记数据才能很好地泛化。

在广泛多样的数据集上预训练的视觉基础模型（VFMs）已成为有前景的工具，因为它们能够捕获可迁移到各种下游任务的通用特征[9]、[10]。这些模型通过迁移学习在医学图像分割中显示出潜力[7]、[11]、[12]、[13]、[14]。然而，直接将VFMs应用于医学成像面临两个重大挑战。(1) VFMs通常在自然图像上训练，而医学图像与自然图像之间存在领域差异，这限制了它们的有效性，因为两者在特征和结构上存在差异[15]、[16]、[17]。(2) VFMs通常体积庞大且计算密集，使得它们在资源有限的临床环境中难以部署。

一个有效的解决方案是将VFMs的知识转移到更小、特定于任务的模型中，这些模型高效且专为医学图像分割设计。知识蒸馏（KD）为此提供了框架，其中大型“教师”模型将其知识传授给较小的“学生”模型[18]、[19]、[20]、[21]、[22]。在基线的任务无关KD流程中，首先通过将小型ViT学生的编码器特征与冻结的VFM教师的编码器特征在未标记的转移集上进行匹配来进行预训练（图1a），然后使用标记的目标数据进行微调（图1b）。任务无关的KD方法通常侧重于传输通用特征表示，而不针对目标任务的特定需求进行定制[23]。这些方法可能无法充分捕获准确分割所需的关键特征，如精细的解剖边界和微妙的病理变化，从而降低了较小模型的有效性。

虽然VFMs具有强大的泛化能力，但它们直接应用于医学分割仍受到领域不匹配、计算成本以及现有KD方法中缺乏特定于任务的适应性的限制。当前的研究很少探讨如何在统一框架下结合特定于任务的微调、双层知识迁移和基于扩散的增强技术，特别是在标签极度稀缺的情况下。

为了解决这些挑战，我们提出了一个新颖的特定于任务的知识蒸馏框架（图1（底部）），该框架能够有效地将大型预训练VFMs的知识转移到专为医学图像分割设计的小型高效模型中。我们的框架首先使用有限的标记医学数据在目标分割任务上对VFMs进行微调，使VFMs能够学习与医学图像相关的任务特征。然后，我们通过使微调后的VFMs和小型模型之间的中间表示和最终分割输出对齐，将这种特定于任务的知识蒸馏到小型模型中。这一过程确保了小型模型继承了高精度分割所需的关键特征，克服了传统KD方法的局限性。

为了降低微调大型VFMs的计算成本，我们采用了低秩适应（LoRA）[24]、[25]、[26]、[27]，该方法允许在最小额外参数的情况下高效地调整大型模型。此外，考虑到医学数据的稀缺性，我们利用扩散模型[28]、[29]、[30]生成合成医学图像。这些合成图像增强了知识蒸馏过程中使用的未标记转移集（图1（顶部）的多样性，使小型模型在数据有限的场景中能够更有效地从微调后的VFMs中学习。

在我们的实验中，我们还将我们的特定于任务的KD框架与自监督学习（SSL）方法进行了有效性比较，后者作为一种替代方案，可以减少对大量标记数据的需求。基于SSL的方法，如Momentum Contrast (MoCo v3) [31]和Masked Autoencoders (MAE) [32]，允许模型通过解决伪装任务来学习有用的表示。尽管SSL在医学成像中显示出潜力[33]、[34]、[35]、[36]，但它通常需要大规模的未标记医学语料库进行有效预训练，而这在实践中往往难以实现。相比之下，我们的方法利用在自然图像上预训练的通用VFMs，并将其高效地适应于目标任务，且只需少量标签。正如我们的实验所证明的，当标记数据极其有限时，这种策略能够持续提升性能。

我们的主要贡献如下：

•

我们提出了一个特定于任务的KD框架，利用微调后的VFMs将领域特定知识转移到较小的视觉变换器（ViT）模型中。

•

我们结合了基于LoRA的微调，使得VFMs的任务适应在计算上可行，同时保留了它们的大规模预训练优势。

•

我们通过生成增强型转移集来改进KD过程，增加了训练数据的多样性和规模。我们对不同转移集大小（1,000；2,000；和3,000张合成图像）的分析表明，特定于任务的KD随着转移集的增大而有效扩展，从而在数据有限的分割任务中显著提升了性能。

•

通过五个数据集的全面实验，我们证明了我们的框架在标注有限的条件下，始终优于任务无关的KD和自监督学习基线。

部分内容摘录

通用且不依赖于模态的医学图像分割

越来越多的研究致力于构建能够处理多种成像领域的通用或不依赖于模态的医学图像分割模型。UniverSeg [37]提出了一个跨模态和解剖结构的统一分割框架。HA-U³Net [38]通过结合混合注意力和嵌套的U-Net架构扩展了这一方向，实现了通用3D分割。MA-SAM [11]将Segment Anything Model (SAM) 适应于多种医学模态，

方法

在这项工作中，我们提出了一个通用的特定于任务的知识蒸馏（TS-KD）框架，用于医学图像分割，以解决数据稀缺、模型高效适应和全面知识迁移的问题。(1) 基于Swin-transformer的扩散模型生成合成数据集，以扩充有限的医学成像数据。(2) 我们使用低秩适应（LoRA）对Segment Anything Model (SAM) 进行微调，有效地调整SAM的参数以满足特定于任务的分割需求。

实验设置

为了全面评估所提出的特定于任务的知识蒸馏（KD）策略在医学图像分割中的有效性，我们设计了涉及六个不同模型组的实验，评估了各种预训练和微调策略对轻量级模型的影响。本节描述了实验设置，包括数据集、模型配置、程序和评估指标。

局限性和失败情况

尽管提出的特定于任务的蒸馏框架在有限的监督下表现出强大的性能，但仍存在一些局限性。首先，我们目前的实验是在2D医学图像上进行的，将框架扩展到3D体积分割是未来工作的一个有意义的方向。其次，扩散生成的合成图像的有效性取决于它们的真实性和多样性；当合成分布与实际数据差异较大时，可能会出现失败情况

讨论与结论

我们的研究表明，在标注有限的条件下，特定于任务的知识蒸馏（KD）是一种有效的策略，用于训练轻量级的ViT模型进行医学图像分割。通过使用低秩适应（LoRA）对视觉基础模型（VFM）进行微调，并将编码器表示和解码器logits蒸馏到ViT-Tiny学生模型中，我们在具有多种成像模式的多个数据集上实现了持续的性能提升。值得注意的是，经过微调的SAM

CRediT作者贡献声明

梁鹏晨：正式分析，数据整理。黄海山：正式分析，数据整理。蒲斌：正式分析，数据整理，概念化。曾全红：调查，数据整理。陈建国：正式分析，数据整理。

利益冲突声明

作者声明与本文的创作和发表没有利益冲突。

致谢

本工作部分由国家自然科学基金（项目编号：62372486）和广东省自然科学基金（项目编号：2023A1515011179）资助。