[论文解读] Self-distillation Augmented Masked Autoencoders for Histopathological Image Classification

Yang Luo, Zhineng Chen+2 more|arXiv (Cornell University)|Mar 31, 2022

AI in cancer detection被引用 11

一句话总结

本文提出SD-MAE，一种用于组织病理学图像分类的自蒸馏增强掩码自编码器，通过将可见图像块的解码器特征知识迁移至编码器浅层，增强特征表示能力。该方法在标准MAE和对比学习基础上表现更优，在多个基准测试中实现最先进性能，准确率与稳定性更高，尤其在低多样性与跨域设置下表现突出。

ABSTRACT

Self-supervised learning (SSL) has drawn increasing attention in histopathological image analysis in recent years. Compared to contrastive learning which is troubled with the false negative problem, i.e., semantically similar images are selected as negative samples, masked autoencoders (MAE) building SSL from a generative paradigm is probably a more appropriate pre-training. In this paper, we introduce MAE and verify the effect of visible patches for histopathological image understanding. Moreover, a novel SD-MAE model is proposed to enable a self-distillation augmented MAE. Besides the reconstruction loss on masked image patches, SD-MAE further imposes the self-distillation loss on visible patches to enhance the representational capacity of the encoder located shallow layer. We apply SD-MAE to histopathological image classification, cell segmentation and object detection. Experiments demonstrate that SD-MAE shows highly competitive performance when compared with other SSL methods in these tasks.

研究动机与目标

• 为解决由于类别多样性有限和语义同质性导致的对比学习在组织病理学图像中高假阴性风险问题。
• 通过利用预训练期间通常被忽略的可见图像块，改进掩码自编码器（MAE）。
• 通过来自可见图像块解码器特征的自蒸馏，增强编码器浅层的表征能力。
• 在包括分类、分割与目标检测在内的多个下游任务中验证SD-MAE的有效性。
• 在低多样性与跨域设置下，展示SD-MAE相较于MAE及其他自监督学习方法的优越性能与鲁棒性。

提出的方法

• SD-MAE引入一种自蒸馏损失，将可见图像块的解码器特征知识迁移至编码器浅层的对应编码特征。
• 该方法将解码器输出特征作为“教师”，将可见图像块在编码器中的输出特征作为“学生”进行知识蒸馏。
• 采用解耦训练策略，仅对0.2%的可见图像块应用自蒸馏损失，提升训练稳定性和性能。
• 自蒸馏损失作用于特征表示而非像素级重建，支持更深层次的特征学习。
• 模型采用视觉Transformer（ViT）作为骨干网络，并在无标签全切片图像（WSIs）上通过掩码-预测的预训练任务进行预训练。
• 通过在下游任务（包括分类、语义分割与目标检测）上进行微调，评估迁移学习性能。

实验结果

研究问题

• RQ1 在组织病理学图像的掩码自编码器中，对可见图像块进行自蒸馏是否能提升编码器的表征能力？
• RQ2 在低类别与低多样性数据集下，SD-MAE相较于对比学习与标准MAE在准确率与稳定性方面表现如何？
• RQ3 SD-MAE能否在不同领域间有效泛化，例如从结直肠癌WSI到乳腺癌WSI的迁移？
• RQ4 与像素级重建相比，使用特征级自蒸馏是否能带来更优的注意力图与特征判别能力？
• RQ5 可见图像块知识蒸馏对分割与目标检测等下游任务的性能贡献如何？

主要发现

• 与MAE相比，SD-MAE在PCam上提升分类准确率1.6%，在NCT上提升0.6%，标准差降低至±0.8（原为±1.0），表明稳定性更高。
• 在跨域迁移中，SD-MAE相比MAE性能提升4%至5.7%，尤其在从高多样性数据集向低多样性数据集迁移时表现更优。
• 在MoNuSeg与Glas数据集上进行语义分割，SD-MAE分别取得92.3%与90.6%的aAcc，优于MAE的92.2%与90.4% aAcc。
• 在NuCLS数据集上进行目标检测，SD-MAE在PCam预训练下达到19.7 mAP，在NCT预训练下达到20.5 mAP，分别优于MAE 0.7与1.6个百分点。
• 可视化结果表明，由于注意力分布更广且关键特征激活更高，SD-MAE在低多样性设置下能学习到更聚焦且多样的注意力图。
• 消融实验表明，将可见图像块上的像素级MSE损失替换为自蒸馏损失，相比基于特征的MSE损失可提升0.3%准确率，相比标准MAE提升0.5%。