[论文解读] Attention-Challenging Multiple Instance Learning for Whole Slide Image Classification

Yunlong Zhang, Honglin Li+4 more|arXiv (Cornell University)|Nov 13, 2023

Advanced Image and Video Retrieval Techniques被引用 4

一句话总结

该论文提出了一种新型的多实例学习（MIL）框架——注意力挑战型多实例学习（ACMIL），用于全切片图像（WSI）分类，通过强制注意力机制聚焦于更具挑战性的实例来缓解过拟合问题。ACMIL结合了多分支注意力（MBA）以捕捉多样化的判别性模式，以及随机Top-K实例掩码（STKIM）以抑制显著实例，从而在三个WSI数据集上实现了最先进的性能，同时提升了泛化能力和鲁棒性。

ABSTRACT

In the application of Multiple Instance Learning (MIL) methods for Whole Slide Image (WSI) classification, attention mechanisms often focus on a subset of discriminative instances, which are closely linked to overfitting. To mitigate overfitting, we present Attention-Challenging MIL (ACMIL). ACMIL combines two techniques based on separate analyses for attention value concentration. Firstly, UMAP of instance features reveals various patterns among discriminative instances, with existing attention mechanisms capturing only some of them. To remedy this, we introduce Multiple Branch Attention (MBA) to capture more discriminative instances using multiple attention branches. Secondly, the examination of the cumulative value of Top-K attention scores indicates that a tiny number of instances dominate the majority of attention. In response, we present Stochastic Top-K Instance Masking (STKIM), which masks out a portion of instances with Top-K attention values and allocates their attention values to the remaining instances. The extensive experimental results on three WSI datasets with two pre-trained backbones reveal that our ACMIL outperforms state-of-the-art methods. Additionally, through heatmap visualization and UMAP visualization, this paper extensively illustrates ACMIL's effectiveness in suppressing attention value concentration and overcoming the overfitting challenge. The source code is available at \url{https://github.com/dazhangyu123/ACMIL}.

研究动机与目标

• 解决由于数据有限、分辨率高和类别不平衡导致的全切片图像（WSI）分类中多实例学习（MIL）持续存在的过拟合问题。
• 探究现有MIL模型中过拟合的根本原因，特别是通过热力图可视化揭示的注意力过度集中于少数判别性实例的现象。
• 开发一种方法，促使注意力机制聚焦于更具多样性和挑战性的实例，而非依赖于少数显著的图像块。
• 通过显式促进注意力多样性并抑制对顶级注意力实例的过度依赖，提升WSI分析中模型的泛化能力和鲁棒性。

提出的方法

• 提出多分支注意力（MBA），通过并行的多个注意力分支捕捉具有不同模式的判别性实例，增强特征多样性。
• 提出随机Top-K实例掩码（STKIM），一种类似数据增强的技术，通过随机掩码Top-K注意力加权的实例，并将其注意力分数重新分配给剩余实例，以防止注意力过度集中。
• 在统一的ACMIL框架中结合MBA与STKIM，强制注意力机制关注更具挑战性且不那么明显的实例，从而提升泛化能力。
• 在MBA中引入多样性损失 $\mathcal{L}_d$，显式鼓励每个注意力分支学习不同的判别性知识，防止冗余。
• 利用UMAP可视化与热力图分析验证模型改进的注意力分布和特征聚类行为。
• 在三个WSI数据集（Camelyon16、BRACS、LBC）上使用两种主干网络（ResNet18和ViT-S/16）进行训练与评估，采用ImageNet和SSL预训练。

实验结果

研究问题

• RQ1 在现有MIL模型中，注意力过度集中于少数判别性实例在多大程度上导致了全切片图像分类中的过拟合？
• RQ2 通过强制注意力聚焦于不那么显著、更具挑战性的实例，是否能提升模型在WSI数据集上的泛化能力和性能？
• RQ3 随机Top-K实例掩码（STKIM）在破坏注意力主导性并促进注意力多样性方面有多有效？
• RQ4 结合多样性损失 $\mathcal{L}_d$ 的多分支注意力（MBA）是否相比单分支注意力能带来更鲁棒和泛化的特征学习？
• RQ5 热力图与UMAP可视化是否能证实ACMIL学习到的特征比先前的SOTA方法更具判别性且分离得更好？

主要发现

• ACMIL在三个WSI数据集上显著优于当前最先进的MIL方法，在使用ResNet18和ImageNet预训练时，LBC数据集的F1-score提升了5.9%。
• 在MBA中引入多样性损失 $\mathcal{L}_d$ 在所有数据集中带来了5.3%至8.0%的性能增益，证实其在促进注意力多样性方面具有关键作用。
• STKIM在测试时轻微降低性能，表明其在推理阶段并非必需，更适合作为训练阶段的正则化技术，类似于cutout。
• UMAP可视化显示，ACMIL学习到的特征更具判别性且分离得更好，LBC数据集的V-measure达到0.316，优于ABMIL的0.224。
• 热力图可视化证实，ACMIL减少了对少数实例的注意力集中，使注意力更广泛地分布在各个图像块上，这与过拟合减少相关。
• ACMIL在Camelyon16、BRACS和LBC数据集上所有指标（F1-score、AUC）均达到最先进水平，且在ABMIL及其他SOTA基线方法上保持一致的性能提升。