MoMA: 基于多头注意力的动量对比学习知识蒸馏，用于组织病理学图像分析|文献速递-视觉大模型医疗图像应用

Title

题目

MoMA: 基于动量机制的对比学习方法在多头注意力机制基础上实现 histopathology图像分析的知识蒸馏

MoMA: 利用多头注意力机制进行动量对比学习方法的知识蒸馏技术，在组织病理学图像分析中有应用。

01

文献速递介绍

计算病理学是一门新领域，在显著提升了准确性与可靠性方面展现了巨大潜力，并显著提升了患者护理、治疗与管理的质量（Cui and Zhang, 2021）。随着先进的人工智能技术（AI）以及机器学习（ML）的进步，并得益于高质量与高分辨率的数据资源的广泛获取，在计算病理学方法已被成功整合到传统病理学工作流程中的多个环节中去，如细胞核识别（Graham et al., 2019）、组织分类（Marini et al., 2021）、层次分析（Chunduru et al., 2022）以及生存分析（Huang et al., 2021；Li et al., 2023）。然而，在这一领域内仍存在工具泛化能力有待进一步提升的问题（Stacke et al., 2020；Aubreville et al., 2023）。

构建准确且可靠的计算病理学工具不仅依赖于先进的AI模型还需要充足的高质量数据。在计算病理学领域AI与ML技术的学习能力以及路径数据集的数量持续增长。相比之下在自然语言处理（Ghorbani et al. 2022）与计算机视觉（Zhai et al. 2022；Dehghani et al. 2023）等其他领域相比公开可获取的路径数据集数量相对较少这一现象主要源于路径切片图像通常体积庞大且复杂导致全球范围内的透明共享极为困难。此外患者隐私及伦理问题限制了相关领域的研究人员获取并分享这些敏感信息的努力。此外路径切片图像的一致性也制约了AI模型在这一领域的泛化能力（Stacke et al. 1998）。值得注意的是Kather19路径数据库包含来自不同结直肠组织类型的10万张图像块但这些图像块最初均来源于仅86张全切片图像而GLySAC数据库则包含来自仅8张全切片图像生成的3万8千七百五十五个细胞核实例进一步凸显出当前路径数据库建设中存在的局限性

虽然在提供多样化病理学数据集方面取得了一定进展,例如如PANDA dataset for prostate cancer Gleason scoring,它涵盖了来自六个不同机构的总共12,625份 Whole Slide Imaging（WSI）样本,使用了三种不同的数字切片扫描仪.然而,在针对特定计算病理学任务时,获取足够数量与特定任务相关的多样化数据仍然面临诸多挑战.这种大规模的数据收集无论如何都需要投入大量时间和人力资源.因此,开发专门针对特定任务的高效计算病理学模型与工具仍然是一个亟待解决的需求.

迁移学习是一种广泛应用的学习策略。它通过从一个任务或问题中继承或转移已有的知识来缓解数据稀缺性的问题。然而，在医学影像分析以及其他领域中已有一系列研究致力于这一技术的应用。传统的研究往往基于像ImageNet和JFT等公共自然图像数据库构建了预训练模型。尽管有研究表明，在医学影像分析中提取自公共自然图像数据库（如ImageNet）的经验特征仍能有效支持诊断相关任务。然而其效果往往受限于医学影像的独特性和多样性。随着公开医学影像数据集数量的增长及其预训练模型的应用逐渐普及；当前仍存在诸多不确定性：例如现有数据集的数量是否足以覆盖实际需求；以及这些数据集之间的多样性是否能够满足不同临床场景的需求

此外

本研究致力于应对计算病理学领域中数据与标注不足这一挑战，并设定构建能够在未知数据上准确且可靠地应用的计算病理学工具的目标。为实现这一目标,我们提出了一种高效的学习框架,并强调其有效性；这种框架通过基于高质量源数据集构建现有模型,并在较小的数据集上训练目标模型,最终实现了预期效果；该方法被命名为基于多头注意力的动量对比学习知识蒸馏（MoMA）,并遵循KD框架以从现有模型中转移相关知识；同时采用动量对比学习与注意力机制以获取一致、可靠且具有上下文感知的特征表示

本研究在多组织病理数据集上对MoMA进行了系统评估，并旨在模拟计算病理学工具研究与开发中的典型工作场景。相较于其他方法而言，在特定任务的目标模型学习方面展现了显著的优势。实验结果则为我们提供了关于如何有效利用有限目标数据集从预训练模型向学生模型迁移知识的重要参考依据。

我们的主要贡献如下：

研究团队成功构建了高效可靠的MoMA学习框架，在现有模型的基础上充分依托高质量数据集，并在有限样本条件下实现了精确可靠的计算病理学工具。

开发一种基于多头注意力机制的动量对比学习方法来进行知识蒸馏；通过稳定且可靠的策略实现来自现有模型的知识到目标模型的迁移。

通过对MoMA进行了系统性评估，并将其应用于多组织病理学数据集（如乳腺癌、肺癌等），验证了其在特定任务目标模型学习中的显著优势

研究了MoMA及其相关方法在不同场景下的效果，并对其基于有限数据集的计算病理学工具开发提供了有效指导

Aastract

摘要

Without a doubt, advanced artificial intelligence models and high-quality data are key factors in achieving success in developing computational pathology tools. While the overall volume of pathology data continues to grow, the lack of high-quality data poses challenges for specific tasks, primarily due to privacy concerns and ethical issues surrounding patient data. Within this study, we introduce a novel approach leveraging knowledge distillation—employing an existing model as a teacher to train a new student model—to address these challenges in computational pathology. This allows the target model to incorporate meaningful features from the teacher model while effectively adapting to distinctive characteristics of its own dataset. Specifically, we employ a student–teacher framework using momentum contrastive learning with multi-head attention mechanisms. This method ensures consistent and context-aware feature representations by distilling relevant knowledge from pre-trained models without direct access to source data. The proposed strategy enables effective knowledge transfer between domains and tasks through rigorous evaluation across diverse scenarios involving both related and unrelated classification tasks. Each scenario was designed using identical classification tasks for both teacher and student models. Experimental results demonstrate that our approach achieves superior accuracy compared with other related methods across various domains and tasks. Furthermore, these findings provide valuable insights into optimal learning strategies for different types of tasks and scenarios within computational pathology.

毫无疑问，在计算病理学领域开发工具的成功离不开先进的人工智能模型以及高质量的数据集作为基础支撑。然而，在具体应用某些任务时会面临一个普遍存在的挑战：由于隐私保护与患者隐私权等因素的影响，在特定任务中缺乏高质量的数据集是一个常见问题。尽管如此，在这一领域内所积累的病理学相关数据总量仍在持续增长中。本研究提出了一种创新的方法——采用知识蒸馏技术（即利用现有模型去学习并构建一个新的目标模型），以期有效解决上述难题。具体而言，在学生–教师框架下，在不直接依赖原始数据的情况下训练目标模型，并通过结合动量对比学习与多头注意力机制来提取并转移相关知识信息；从而使得目标模型能够获得具有语境感知特性的稳定特征表示。这种设计使得目标模型不仅能够无缝适应目标域的独特特征需求，并且能够充分继承教师模型中的信息性表征。

我们对多维度场景采用了系统性评估策略，并观察其性能表现。具体而言，在教师网络设计上进行了一系列优化工作：首先，在面对与目标网络相同的分类问题时均展开了训练；其次，在处理与目标网络相关的问题时也保持了良好的性能水平；最后，在面对完全无关的任务时仍能展现出一定的适应能力。实验数据显示，在知识迁移方面我们的方法表现出了显著的准确率和稳定性，并且这种优势尤其体现在跨领域任务中。这些结果不仅验证了我们的方法的有效性，并且还为其在计算病理学中的应用提供了重要的参考依据。

Method

方法

展示MoMA概述的图1和算法1可在附录A中找到。

该框架概述如图1及附录A中所示。定义为𝐷𝑆𝐶 = {(𝐱𝑖 , 𝐲𝑖 )}𝑖=1…𝑁_{SC}为源教师数据集,𝐷𝑇𝐺 = {(𝐱𝑖 , 𝐲𝑖 )}𝑖=1…𝑁_{TG}为目标学生数据集,其中\mathbf{x}_i与\mathbf{y}_i分别代表第i张病理图像及其真实标签, N_{SC}与N_{TG}分别表示源数据集与目标数据集中的样本数量（N_{SC}≫N_{TG}）。源教师数据集用于训练教师模型,而目标学生数据集则用于训练目标学生模型。令\mathcal{T}代表教师模型,\mathcal{S}代表学生模型。\mathcal{T}由教师编码器f_T与教师分类器g_T构成,\mathcal{S}则包含学生编码器f_S与学生分类器g_S。此外,该框架还涉及教师投影头p_T,教师注意力头h_T,学生投影头p_S,以及学生注意力头h_S等组件。对于输入图像\mathbf{x}_i, teacher编码器f_T与student编码器f_S均提取初始特征表示,随后这些特征通过相应的投影头与注意力头进行一系列处理（即通过p_T,h_T或p_S,h_S）,从而增强其特征表征能力.g_T,g_S接收这些初始特征表示并进行图像分类.值得注意的是,g_T^*仅在训练teacher网络时使用.由于受限于共享医学数据获取的问题,我们假设如下场景:已知teacher网络\mathcal{T}已在预训练阶段完成对source Teacher dataset D_{SC}上的训练,并且其预训练权重已可调用但无法直接访问source dataset.在提供pre-trained teacher encoder f_T的情况下,该框架的目标是在target dataset D_{TG}上实现稳健且高效的student model \mathcal{S}.为此,MoMA综合运用了知识蒸馏（KD）策略以及动量对比学习方法.结合这两种学习策略,MoMA允许从在高质量source dataset上预训练过的teacher encoder f_T向有限规模target student encoder $f_S进行知识迁移.

Conclusion

结论

Herein, we introduced a highly efficient and powerful learning framework named MoMA for establishing reliable classification models in pathology images. Leveraging the KD framework, momentum contrastive learning, and self-attention mechanisms (SA), MoMA successfully transferred knowledge from multiple source domains to target domains while developing robust classification models tailored for five distinct tasks. Moreover, the experimental outcomes of MoMA indicate an appropriate learning strategy for diverse distillation tasks and scenarios. We expect this approach to be highly beneficial in advancing computational pathology tools across various tasks. Future research will focus on further exploring the efficiency of the KD method and applying MoMA to additional datasets and tasks within computational pathology.

在本研究中,我们提出了一种高效而精确的学习框架,命名为MoMA,旨在生成具有高准确性和鲁棒性的病理图像分类模型。通过结合知识蒸馏框架、动量对比学习方法以及注意力机制,MoMA不仅实现了源域向目标域的知识迁移能力,并且能够训练出适用于五个不同任务的鲁棒分类模型。进一步地,MoMA在各向异性蒸馏任务及多场景下的实验结果表明其适应性。这一发现将有助于开发出更高效的计算病理学工具,在多个应用场景中得到广泛应用

Results

结果

Table 1 and Fig. 3 (and Figs. C.1 and C.2 in Appendix C) show theresults of MoMA and its competitors on the two TMA prostate datasets(Prostate USZ and Prostate UBC). On Prostate USZ, the teacher modelTC𝑃 𝐴𝑁𝐷𝐴, which was trained on PANDA only, achieved 63.4% ACC,0.526 F1, and 0.531 𝜅**𝑤, which is substantially lower to other studentmodels with 𝑇 𝐿, 𝐿𝐷, and 𝐹𝐷. Among the student models with 𝑇 𝐿, thestudent model with no pre-trained weights (FT𝑁𝑜𝑛𝑒) was inferior to theother two student models; the student model pre-trained on PANDA(FT𝑃 𝐴𝑁𝐷𝐴) outperformed the student model pre-trained on ImageNet(FT𝐼𝑚𝑎𝑔𝑒𝑁𝑒𝑡). These indicate the importance of pre-trained weights andfine-tuning on the target dataset, i.e., Prostate USZ. As for the KDapproaches, MoMA𝑃 𝐴𝑁𝐷𝐴, pre-trained on PANDA, outperformed allother KD methods, achieving ACC of 73.6%, which is 0.9% higher thanFT𝑃 𝐴𝑁𝐷𝐴, and F1 of 0.687 and 𝜅𝑤 of 0.670, which are comparable tothose of FT𝑃 𝐴𝑁𝐷𝐴.

On the independent test set, Prostate UBC, it is remarkable thatTC𝑃 𝐴𝑁𝐷𝐴 achieved 78.2% ACC and 0.680 𝜅**𝑤, which are superior tothose of all the student models with 𝑇 𝐿, likely suggesting that thecharacteristic of PANDA is more similar to Prostate UBC than ProstateUSZ. The performance of the student models with 𝑇 𝐿 and 𝐹𝐷 wassimilar to each other between Prostate USZ and Prostate UBC; forinstance, MoMA𝑃 𝐴𝑁𝐷𝐴 obtained higher ACC but lower F1 and 𝜅𝑤 onProstate UBC than on Prostate USZ. As MoMA and other student modelswith 𝐹𝐷 adopt vanilla KD by setting 𝛾 to 1 in , i.e., mimicking theoutput logits of the teacher model, there was, in general, a substantialincrease in the performance on Prostate UBC. MoMA𝑃 𝐴𝑁𝐷𝐴, in particular, achieved the highest ACC of 83.3% and 𝜅𝑤 of 0.763 overall modelsunder consideration, which are 11.1% and 0.145 higher than those onProstate USZ in ACC and 𝜅𝑤, respectively.By randomly sampling 25% and 50% of the training set, we repeatedthe above experiments using MoMA and other competing models toassess the effect of the size of the training set. The results of the sametask distillation using 25% and 50% of the training set are availablein Appendix B (Tables B.1 and B.2). The experimental results were moreor less the same as those using the entire training set. MoMA𝑃 𝐴𝑁𝐷𝐴 wascomparable to FT𝑃 𝐴𝑁𝐷𝐴 on Prostate USZ. KL+MoMA𝑃 𝐴𝑁𝐷𝐴 outperformed the competing models on Prostate UBC. These results validatethe effectiveness of MoMA on the extremely small target dataset.

这些图表（此外，在附录C中包含有图C.1和C.2）详细呈现了MoMA及其竞争对手的方法在两个TMA前列腺数据集（Prostate USZ和Prostate UBC）上的实验结果。

在Prostate USZ数据集上进行实验后发现，在仅使用PANDA数据集对教师模型进行微调时（即TC𝑃 𝐴𝑁𝐷𝐴），该模型达到了63.4%的准确率、0.526的F1分数以及加权Kappa系数为0.531的结果。这些数值均显著低于通过迁移学习（TL）、逻辑蒸馏（LD）以及特征蒸馏（FD）方法对学生成员进行微调后的模型表现。具体而言，在迁移学习方法下未使用任何预训练权重的学习者（FTNone）的表现低于其他两种学习者；而针对PANDA数据集进行过微调的学习者（FT𝑃 𝐴𝑁𝐷𝐴）在目标领域上的性能优于基于ImageNet预训练的学习者。这些实验结果凸显了预训练权重以及在其目标数据集（如Prostate USZ）上进行微调的重要性。

在知识蒸馏领域中，在经过PANDA预训练的情况下...

在独立测试集Prostate UBC中, 该教师模型TC𝑃 𝐴ﻨ𝐷𝐴达到了78.2%的准确率（ACC）并获得了0.680加权κ值（𝜅𝑤），展现出显著优势于所有通过𝑇𝐿方法训练的学生模型。这一结果暗示着PANDA特征与Prostate UBC之间的相似性可能比其与Proseat USZ之间的更为接近。在对Proseat USZ与Proseat UBC进行对比时, 使用𝑇𝐿与𝐹𝐷的学习策略所得到的学生模型表现相当接近；例如, MoMA𝑃 𝐴 NDP A 在该测试集中的准确率高于其在Proseat USZ上的表现, 其F1值及加权κ值则略低于后者。值得注意的是, 采用基于KD的学习策略（通过将γ设定为1以模仿教师模型的输出logits）后, Proseat U BC 的性能得到了显著提升；具体而言, MoMA-P-AN-D A 模型在这项评估中取得了83.3%的最高准确率（ACC），其加权κ值也达到了0.763, 这两项指标分别较其实现在另一个基准集上的数值高出11.1%及0.145。

为了探讨训练集大小对模型性能的影响，在现有数据基础上我们选取了训练集中约25%和50%的数据样本进行重复实验研究，并将基于这些不同比例的数据样本进行相同任务的蒸馏过程及其结果进行了详细记录（见附录B中的表B.1及表B.2）。实验结果显示所选方案所得模型性能与基于完整数据集所获得的结果接近无异。其中MoMA𝑃 𝐴𝑁𝐷𝐴在Prostate USZ基准测试中表现不相上下与FT𝑃 𝐴𝑁𝐷𝐴相当而KL+MoMA Pis算子则在Prostate UBC基准测试中展现出显著的优势这些实验结论充分验证了MoMA Pis算子在极小规模目标数据集上的有效性

Figure

图

Fig. 1. Overview of the MoMA: Attention-augmented momentum contrast knowledge distillation framework. A batch of input images is encoded by the student encoder (𝑓 𝑆 ), andthe momentum teacher (𝑓 𝑇 ), and each feature representation is re-weighted with regard to other images in the batch as the context. A classifier is added on top of the studentencoder. The student model is jointly optimized by contrastive loss and cross-entropy loss

图1展示了MoMA概览：一种基于注意力增强机制的动量对比知识蒸馏方法。经过学生编码器（𝑓 𝑆）和动量教师编码器（𝑓 𝑇）处理后得到一批输入图像特征向量。每个特征向量会根据批次中其他图像的信息进行加权融合以形成更加鲁棒的空间表征。为了提升模型性能，在学生编码器结构顶部附加了一个多分类识别模块。通过综合运用对比损失函数与交叉熵损失函数进行联合优化训练以实现目标模型参数的有效更新

In the context of knowledge distillation, supervised tasks are consistently carried out. Feature extraction through distillation is implemented when a robust teacher model is accessible. Finally, vanilla 𝐿𝐾𝐷 serves as the standard approach when both the teacher and student models are engaged in identical tasks. SSL denotes self-supervised learning.

图2展示了不同任务与数据集间的蒸馏流程概览。（1）该流程始终采用监督学习方式。（2）当具备经过良好训练的教师模型时，则可采用特征蒸馏方法。（3）当教师模型与学生模型在相同的任务场景下运行时，则应采用基础形式的𝐿𝐾𝐷（_SSL保留自监督学习特性）。

Figure 3 illustrates box plots employed in same-task distillation, indicating that all KD models make use of the pre-trained parameters from PANDA.

Figure 3 presents a box plot illustrating the performance of knowledge蒸馏 (KD) models on the same-task distillation task. All KD models utilize pre-trained weights from PANDA.

Fig. 4. Box plot showing task-oriented distillation. Each of these KD models adopts the pre-trained weights from PANDA.

Chart 4: Boxplots of Relevant Task Distillation, all knowledge distillation (KD) models leverage pre-trained weights from PANDA.

Figure 5 illustrates bar charts for unrelated task distillation, where all KD models extract pre-trained weights from ImageNet.

图5. 不相关任务蒸馏的柱状图：所有知识蒸馏（KD）模型均利用基于ImageNet的预训练权重。

Fig. 6. The figure illustrates the correlation coefficients between teacher and student feature representations.

图6. 教师网络和学生网络特征表示之间的相关系数矩阵。

Fig. 7. 使用t-SNE进行特征表示可视化分析，并结合轮廓系数对ImageNet以及PANDA教师模型与5个学生数据集的特征进行对比研究：(a)前列腺USZ组,(b)前列腺UBC组,(c)前列腺AGGC组,(d)结直肠K16组,(e)乳腺BRACS组.

图7. 采用t-SNE进行特征表示的可视分析及其轮廓系数评估，并分别涉及两个主要的教师模型框架（ImageNet与PANDA）以及五个代表性学生数据集集合：包括（a）Prostate USZ类型的样本（USZ），（b）Prostate UBC类型的样本（UBC），（c）Prostate AGGC类型的样本（AGGC），（d）Colon K16类型的数据（K16），以及（e）Breast BRACS类型的案例（BRACS）。

Figure C-1 illustrates the confusion matrices for Prostate-USZ (Test I). Each of these matrices summarizes the averaged results across five independent runs.

第C.1章中的Prostate USZ（对照组I）上的confusion matrix展示了各分类器性能的对比结果

Figure C.2 presents the confusion charts for Prostate UBC (Test II). Each chart corresponds to the average across five iterations.

图表C.2展示了Prostate UBC（第二测试集）上的混淆矩阵。每个混淆矩阵体现了5次运行结果的平均值。

Fig. C.3 illustrates the confusion matrices for Prostate AGGC CV (Test I). Each matrix corresponds to the average results across five-fold cross-validation runs.

图C.3. 基于Prostate AGGC CV测试集I的混淆矩阵。每一个这样的混淆矩阵代表了五折交叉验证过程的平均表现。

Fig. C.4 illustrates the confusion matrices generated from the Prostate AGGC test (Test II). Each confusion matrix reflects the average results across 5-fold cross-validation experiments.

图表C.4展示了Prostate AGGC测试集（测试集II）的相关混淆矩阵。每个混淆矩阵具体反映了五折交叉验证实验的结果均值。

Fig. C.5. The confusion matrices for Colon K16 SN (Test I). The each confusion matrix reflects the average across five iterations.

图表C.5. Colon K16 SN（测试数据集1）上的混淆矩阵。每个confusion matrix represents the average outcome across five runs.

Fig. C.6. Confusion charts for Colon K16 (Test II). Each confusion chart shows the average performance over five runs.

该图表展示的是测试集II上所对应的混淆矩阵。其中每个混淆矩阵均反映其5次运行过程中的平均表现水平。

Fig. C.7. Confusion matrices for irrelevant task distillation: classification of breast carcinoma subtypes. The average confusion matrix across 5 runs represents each instance.

图表C.7展示了一个通过无关任务蒸馏构建的混淆矩阵及其在乳腺癌亚型分类中的应用

Fig. C.8. Arrays of confusion matrices in the context of irrelevant task distillation for gastric microsatellite instability estimation.

图 C.8 的无关任务蒸馏相关混淆矩阵用于评估胃部微卫星不稳定性的预测性能；每一个这样的混淆矩阵均代表了五折交叉验证实验中各项指标的平均值

Table

表

Table 1 Outcomes of task-level knowledge distillation. The KL denotes the utilization of KL divergence loss.

表1 相同任务蒸馏的结果。KL表示使用KL散度损失。

Table 2Results of relevant task distillation.

表2 相关任务蒸馏的结果。

Table 3 Results of irrelevant-task-distillation: large-intestine tissue-type classification.

表3 无关任务蒸馏的结果：结肠组织类型分类。

Table 4 Results of irrelevant task distillation: breast carcinoma sub-type classification.

表4无关任务蒸馏的结果：乳腺癌亚型分类。

Table 5 Ablation experiments of MoMA comparing the performance with and without Multi-head Attention across three distillation tasks.

表5带有和不带多头注意力的MoMA在三个蒸馏任务上的消融实验结果。

Table 6Outcomes of three specialized distillation processes utilizing MoMA framework and self-supervised learning approaches.

表6使用MoMA和自监督学习在三个蒸馏任务上的结果。

Table 7 Findings from irrelevant task distillation: Prediction of gastric microsatellite instability

表7无关任务蒸馏的结果：胃部微卫星不稳定性预测。

Table B.1评估结果表明同一任务分层训练在50%的数据集上表现良好。其中KL表示使用了KL散度损失函数.

表B.1在50%训练数据上的相同任务蒸馏结果。KL表示使用KL散度损失。

Table B.2 presents the results of uniform task-based distillation from a 25% training dataset. KL represents the utilization of KL divergence loss.

表B.2在25%训练数据上的相同任务蒸馏结果。KL表示使用KL散度损失

Table 3 Outcomes from irrelevant-task distillation: Prostate cancer categorization.

表B.3无关任务蒸馏的结果：前列腺癌分类。