[论文解读] Task-specific Fine-tuning via Variational Information Bottleneck for Weakly-supervised Pathology Whole Slide Image Classification

Honglin Li, Chenglu Zhu+6 more|arXiv (Cornell University)|Mar 15, 2023

AI in cancer detection被引用 6

一句话总结

本文提出了一种基于变分信息瓶颈（Variational Information Bottleneck）的微调框架，用于弱监督全切片图像（WSI）分类，通过将每张 WSI 的超过 10,000 个图像块蒸馏为少于 1,000 个关键实例，显著降低了计算成本。该方法仅使用 WSI 级别的标签即可实现主干网络的端到端训练，在性能和泛化能力上显著优于 ImageNet 预训练特征和自监督特征，且仅使用不到 1% 的图像块标注，性能已接近完全监督基线。

ABSTRACT

While Multiple Instance Learning (MIL) has shown promising results in digital Pathology Whole Slide Image (WSI) classification, such a paradigm still faces performance and generalization problems due to challenges in high computational costs on Gigapixel WSIs and limited sample size for model training. To deal with the computation problem, most MIL methods utilize a frozen pretrained model from ImageNet to obtain representations first. This process may lose essential information owing to the large domain gap and hinder the generalization of model due to the lack of image-level training-time augmentations. Though Self-supervised Learning (SSL) proposes viable representation learning schemes, the improvement of the downstream task still needs to be further explored in the conversion from the task-agnostic features of SSL to the task-specifics under the partial label supervised learning. To alleviate the dilemma of computation cost and performance, we propose an efficient WSI fine-tuning framework motivated by the Information Bottleneck theory. The theory enables the framework to find the minimal sufficient statistics of WSI, thus supporting us to fine-tune the backbone into a task-specific representation only depending on WSI-level weak labels. The WSI-MIL problem is further analyzed to theoretically deduce our fine-tuning method. Our framework is evaluated on five pathology WSI datasets on various WSI heads. The experimental results of our fine-tuned representations show significant improvements in both accuracy and generalization compared with previous works. Source code will be available at https://github.com/invoker-LL/WSI-finetuning.

研究动机与目标

• 解决现有弱监督 WSI 分类方法在使用冻结的 ImageNet 或自监督表征时面临的高计算成本与性能瓶颈问题。
• 在存在海量图像块的吉字节级 WSI 中，实现 MIL 框架下主干网络的端到端微调。
• 通过利用 WSI 级标签，将任务无关的自监督特征转化为任务特定表示，以提升特征的任务特异性和泛化能力。
• 在弱监督设置下支持多样化的训练时数据增强，以提升在领域分布偏移下的鲁棒性。

提出的方法

• 引入变分信息瓶颈（VIB）模块，将每张 WSI 的超过 10,000 个图像块蒸馏为少于 1,000 个最具代表性的实例，使训练成本降低逾十倍。
• 将 VIB 模块应用于学习 WSI 包的最小充分统计量，从而在存在大量图像块的情况下仍能实现通过主干网络的梯度反向传播。
• 将基于 VIB 的微调方法与自监督预训练相结合，仅使用 WSI 级标签，将任务无关的 SSL 特征转化为任务特定表示。
• 采用拉格朗日优化框架平衡信息压缩与预测准确率，确保模型保留关键的诊断特征。
• 在微调过程中引入多样化的训练时数据增强，以提升在存在领域偏移的数据集上的泛化能力。
• 采用可微的 top-K 选择机制，识别每张 WSI 包内最具支持度的实例，构成蒸馏后的包以实现高效训练。

实验结果

研究问题

• RQ1 变分信息瓶颈框架是否能有效降低在吉字节级 WSI 上训练的计算负担，同时保留诊断信息？
• RQ2 WSI 级弱监督在多大程度上能提升自监督或 ImageNet 预训练特征在全切片图像分类中的任务特异性？
• RQ3 使用蒸馏后的图像块表示对主干网络进行端到端微调，是否能带来优于冻结特征提取的性能提升？
• RQ4 所提方法在领域偏移下是否具备良好的泛化能力，尤其是结合多样化数据增强时？

主要发现

• 与冻结 ImageNet 预训练特征相比，所提微调方法在 Camelyon-16-C 上 F1 分数提升 18.43%，在 Camelyon-17 上提升 18.86%。
• 在 Camelyon-16-C 上，使用微调特征的 Max-pooling 方法在 F1 上分别优于 CLAM-SB 和 TransMIL 10.91% 和 6.28%。
• 该方法在领域偏移下表现出更强的泛化能力，DTFD-MIL 在 Camelyon-17 上微调后 F1 提升 11.30%，表明其具备出色的鲁棒性。
• F1 指标显著优于 AUC，表明该方法在提升分类性能（敏感性与特异性）方面优于排序能力，更具临床相关性。
• T-SNE 可视化结果表明，微调后的特征分布更具可分性与任务特异性，v-score 更高。
• CLAM-SB 的注意力图显示，使用微调特征的模型更聚焦于肿瘤区域，体现了更高的可解释性与诊断相关性。