[论文解读] Multiple Instance Learning Framework with Masked Hard Instance Mining for Whole Slide Image Classification

Wenhao Tang, Sheng Huang+4 more|arXiv (Cornell University)|Jul 28, 2023

Image Retrieval and Classification TechniquesComputer Science被引用 3

一句话总结

本文提出MHIM-MIL，一种新型的多实例学习（MIL）框架，通过使用带有动量更新和一致性正则化的孪生教师-学生结构，隐式地进行掩码难样本挖掘（MHIM），从而提升全切片图像（WSI）分类性能。与仅关注显著性显著样本的现有方法不同，MHIM-MIL聚焦于不那么显著但更具挑战性的图像块，从而在CAMELYON-16和TCGA肺癌数据集上以更低的训练成本实现了最先进（SOTA）的性能表现。

ABSTRACT

The whole slide image (WSI) classification is often formulated as a multiple instance learning (MIL) problem. Since the positive tissue is only a small fraction of the gigapixel WSI, existing MIL methods intuitively focus on identifying salient instances via attention mechanisms. However, this leads to a bias towards easy-to-classify instances while neglecting hard-to-classify instances. Some literature has revealed that hard examples are beneficial for modeling a discriminative boundary accurately. By applying such an idea at the instance level, we elaborate a novel MIL framework with masked hard instance mining (MHIM-MIL), which uses a Siamese structure (Teacher-Student) with a consistency constraint to explore the potential hard instances. With several instance masking strategies based on attention scores, MHIM-MIL employs a momentum teacher to implicitly mine hard instances for training the student model, which can be any attention-based MIL model. This counter-intuitive strategy essentially enables the student to learn a better discriminating boundary. Moreover, the student is used to update the teacher with an exponential moving average (EMA), which in turn identifies new hard instances for subsequent training iterations and stabilizes the optimization. Experimental results on the CAMELYON-16 and TCGA Lung Cancer datasets demonstrate that MHIM-MIL outperforms other latest methods in terms of performance and training cost. The code is available at: https://github.com/DearCaat/MHIM-MIL.

研究动机与目标

• 为解决现有MIL方法过度优先关注显著且易分类样本而导致泛化能力差的问题。
• 通过隐式识别并聚焦于弱监督WSI分类中难以分类的样本，提升模型的判别能力。
• 开发一种训练框架，提升稳定性与效率，且无需额外参数或复杂梯度更新。
• 通过引入一致性约束和掩码实例挖掘策略，缓解过拟合问题，提升模型鲁棒性。
• 在真实世界的WSI数据集上，验证MHIM-MIL在不同基于注意力的MIL模型中的通用性与有效性。

提出的方法

• MHIM-MIL采用孪生教师-学生架构，其中学生模型使用由动量更新的教师模型挖掘出的难样本进行训练。
• 动量教师通过学生模型参数的指数移动平均（EMA）进行更新，确保知识蒸馏过程稳定且逐步推进。
• 基于注意力分数实施实例掩码策略——将高注意力区域的图像块进行掩码，迫使学生模型聚焦于不显著且难以分类的样本。
• 引入三种混合掩码策略，通过多样化难样本分布提升训练效率并降低过拟合风险。
• 在学生与教师预测之间施加一致性约束，以利用超出包级别标签的额外监督信号。
• 该框架可兼容任意基于注意力的MIL模型，支持即插即用，具备广泛适用性。

实验结果

研究问题

• RQ1 若将注意力从显著样本转向难以分类的样本，是否能提升基于MIL的WSI分类器的泛化能力与判别能力？
• RQ2 在WSI分类中缺乏实例级标注的情况下，如何有效挖掘难样本？
• RQ3 基于动量的教师-学生框架结合一致性正则化，是否能增强MIL中难样本挖掘的稳定性与性能？
• RQ4 与传统基于注意力的MIL方法相比，掩码难样本挖掘是否能降低过拟合风险并提升训练效率？
• RQ5 MHIM-MIL在不同基于注意力的MIL模型与WSI数据集上的泛化能力达到何种程度？

主要发现

• 在CAMELYON-16数据集上，MHIM-MIL实现了最先进性能，相比基线AB-MIL模型AUC提升0.97%。
• 在TCGA肺癌数据集上，MHIM-MIL相比基线AB-MIL模型AUC提升1.00%，表明在不同模型上均具有一致性增益。
• 采用EMA更新策略的动量教师表现最佳，优于固定初始化和非小批量梯度更新的教师模型。
• 可视化结果表明，MHIM-MIL生成的肿瘤概率图更全面、更准确，覆盖更多真实肿瘤区域，同时减少非肿瘤区域的误报。
• 该框架降低了对显著图像块的依赖，提升了对细微、难分类肿瘤区域的检测能力，从而增强模型的鲁棒性与泛化能力。
• MHIM-MIL在更低的训练成本下实现更优性能，证明其在不增加额外参数的前提下具备更高的效率与稳定性。