[论文解读] Exploring Visual Prompts for Whole Slide Image Classification with Multiple Instance Learning

Yi Lin, Zhongchen Zhao+4 more|arXiv (Cornell University)|Mar 23, 2023

Digital Imaging for Blood Diseases被引用 4

一句话总结

本文提出视觉提示（visual prompts）以通过多实例学习（MIL）微调ImageNet预训练模型，实现全切片图像（WSI）分类，缓解自然图像与组织病理学图像之间的域偏移问题。通过冻结主干网络并仅在代表性切片上微调轻量级提示组件，该方法在Camelyon16和TCGA-NSCLC数据集上，对多种MIL模型和主干网络均实现了稳定且一致的性能提升。

ABSTRACT

Multiple instance learning (MIL) has emerged as a popular method for classifying histopathology whole slide images (WSIs). However, existing approaches typically rely on pre-trained models from large natural image datasets, such as ImageNet, to generate instance features, which can be sub-optimal due to the significant differences between natural images and histopathology images that lead to a domain shift. In this paper, we present a novel, simple yet effective method for learning domain-specific knowledge transformation from pre-trained models to histopathology images. Our approach entails using a prompt component to assist the pre-trained model in discerning differences between the pre-trained dataset and the target histopathology dataset, resulting in improved performance of MIL models. We validate our method on two publicly available datasets, Camelyon16 and TCGA-NSCLC. Extensive experimental results demonstrate the significant performance improvement of our method for different MIL models and backbones. Upon publication of this paper, we will release the source code for our method.

研究动机与目标

• 为缓解ImageNet预训练模型与组织病理学WSI之间的域偏移问题，该问题限制了基于MIL的WSI分类中特征质量的提升。
• 克服在有限切片级别标签下微调整个特征提取器的局限性，该方法存在过拟合风险并会损害预训练表征能力。
• 通过引入视觉提示，探索视觉领域中的提示学习，特别是针对组织病理学图像，以将预训练特征适配至病理数据。
• 开发一种高效、轻量级的训练方案，仅更新提示组件而冻结主干网络，从而降低计算成本。
• 在公开的WSI数据集上，验证视觉提示在多种主干网络（ResNet-18/50）和MIL模型（DTFD, ABMIL）上的有效性。

提出的方法

• 提出一种新颖的视觉提示学习框架，将可学习的提示组件添加至冻结的ImageNet预训练主干网络中，以适配组织病理学数据的特征表示。
• 采用代表性切片选择（RPS）策略，减少用于提示训练的切片数量，提升效率并降低过拟合风险。
• 仅训练提示模块和轻量级MIL分类器，同时保持主干网络权重冻结，实现高效的领域自适应。
• 采用两阶段流程：第一阶段基于特征多样性从WSI中选择代表性切片；第二阶段使用这些切片仅微调提示组件。
• 将提示模块集成至特征提取头中，支持提示与MIL分类器的端到端联合训练，同时保持主干网络固定。
• 采用基于注意力的MIL模型（如DTFD、ABMIL）在提示适配后，将切片级特征聚合为WSI级预测结果。

Figure 1: An illustration of three adapting schemes for WSI classification with MIL: (a) freezing the pretrained feature extractor while only training the MIL classifier; (b) fine-tuning the feature extractor and training the MIL classifier; (c) only training the prompt blocks and the MIL classifier

实验结果

研究问题

• RQ1 视觉提示能否在基于MIL的分类任务中有效缓解ImageNet与组织病理学WSI之间的域偏移？
• RQ2 当缺乏切片级别标签时，提示学习是否优于微调整个特征提取器？
• RQ3 提示模块的数量如何影响WSI分类中的模型性能与鲁棒性？
• RQ4 为实现高效提示训练且避免过拟合，代表性切片的最佳数量是多少？
• RQ5 所提方法是否能在不同主干网络与MIL架构之间实现良好泛化？

主要发现

• 在Camelyon16数据集上，使用ResNet-18时，该方法相比基线模型AUC提升3.83%；在TCGA-NSCLC数据集上，AUC提升1.28%。
• 使用ResNet-50时，该方法在所有评估指标上均优于冻结基线和全参数微调方法（RPS-FT），展现出更优的泛化能力与鲁棒性。
• 与基线相比，添加超过一个提示模块可使AUC提升约1%，表明提示组件具有显著有效性与稳定性。
• 该方法在不同MIL模型（DTFD与ABMIL）和主干网络（ResNet-18与ResNet-50）上均保持一致的性能增益，表明其具有广泛适用性。
• 在所有K值的代表性切片数量下，该方法性能均优于RPS-FT，且GPU资源消耗不足其50%，证明了其高效性。
• 消融实验表明，微调整个主干网络会导致过拟合并降低性能，而基于提示的适配方法可有效避免该问题。

Figure 2: Overview of our method. ( \@slowromancap i@) MIL classifier training with all patch features. ( \@slowromancap ii@) Representative patch selection. ( \@slowromancap iii@) Prompt learning with representative patches.