[论文解读] Language models are good pathologists: using attention-based sequence reduction and text-pretrained transformers for efficient WSI classification

Juan I. Pisula, Katarzyna Bożek|arXiv (Cornell University)|Nov 14, 2022

AI in cancer detection被引用 5

一句话总结

本文提出 SeqShort，一种基于多头注意力的序列压缩层，可将整个切片图像（WSIs）压缩为固定长度、有序的特征向量序列，从而利用文本预训练的 Transformer 实现高效且准确的分类。通过结合冻结的 RoBERTa 基础模型并仅微调 0.04% 的参数，该方法在淋巴结转移和癌症亚型分类任务中实现了最先进性能，且计算资源和可训练参数极少。

ABSTRACT

In digital pathology, Whole Slide Image (WSI) analysis is usually formulated as a Multiple Instance Learning (MIL) problem. Although transformer-based architectures have been used for WSI classification, these methods require modifications to adapt them to specific challenges of this type of image data. Among these challenges is the amount of memory and compute required by deep transformer models to process long inputs, such as the thousands of image patches that can compose a WSI at $ imes 10$ or $ imes 20$ magnification. We introduce extit{SeqShort}, a multi-head attention-based sequence shortening layer to summarize each WSI in a fixed- and short-sized sequence of instances, that allows us to reduce the computational costs of self-attention on long sequences, and to include positional information that is unavailable in other MIL approaches. Furthermore, we show that WSI classification performance can be improved when the downstream transformer architecture has been pre-trained on a large corpus of text data, and only fine-tuning less than 0.1\% of its parameters. We demonstrate the effectiveness of our method in lymph node metastases classification and cancer subtype classification tasks, without the need of designing a WSI-specific transformer nor doing in-domain pre-training, keeping a reduced compute budget and low number of trainable parameters.

研究动机与目标

• 为解决由于 WSI 的图像块序列过长，导致将 Transformer 应用于 WSI 时计算成本过高的问题。
• 在无需领域特定微调或架构重构的情况下，实现将强大预训练语言模型应用于数字病理学。
• 通过可学习的位置编码和基于注意力的摘要机制，保留 WSI 表征中的空间和上下文信息。
• 在保持或提升性能的同时，显著减少 WSI 分类模型中的可训练参数数量。
• 证明文本预训练的 Transformer 可通过极少微调被有效重用于组织病理学分类。

提出的方法

• 引入一种基于多头注意力的序列压缩层（SeqShort），将数千个 WSI 图像块压缩为固定长度、有序的特征向量序列。
• SeqShort 使用可学习的查询向量对全部图像块嵌入进行注意力计算，通过自注意力机制捕捉显著特征。
• 该方法引入可学习的位置嵌入以保留空间顺序，解决了标准 MIL 方法在缺乏位置信息方面的局限性。
• 使用预训练的 RoBERTa-base 语言模型作为 Transformer 编码器，仅微调其中 0.04% 的参数。
• 使用在 ImageNet 上预训练的 EfficientNetV2-L 主干网络提取 WSI 的实例级特征。
• 最终分类头为多层感知机（MLP），其输入为 Transformer 编码器输出的 [CLS] token 表征。

实验结果

研究问题

• RQ1 基于注意力的序列压缩能否在保留分类判别性信息的同时，有效压缩 WSI 图像块序列？
• RQ2 在无需领域特定预训练的情况下，利用类似 RoBERTa 的文本预训练 Transformer 是否能提升 WSI 分类性能？
• RQ3 仅微调少于 0.1% 的参数是否能实现 WSI 分类任务的最先进结果？
• RQ4 与缺乏位置信息的标准 MIL 方法相比，序列压缩层中引入可学习位置嵌入对模型性能有何影响？
• RQ5 当与高效的序列压缩结合时，通用的、非 WSI 特定的 Transformer 架构是否能在数字病理学中实现优异性能？

主要发现

• 所提方法在 CAMELYON16 淋巴结转移分类基准上达到最先进性能，AUC 为 0.948，准确率为 0.882。
• 在 TCGA-BRCA 侵袭性乳腺癌亚型分类任务中，模型在 10 折交叉验证下达到 AUC 0.917，准确率 0.841。
• 在肾细胞癌亚型分类中，模型达到 AUC 0.892，准确率 0.815，表明其在多种癌症类型中均具鲁棒性。
• 模型将可训练参数数量从 8500 万减少至 330 万（占总参数的 0.04%），显著降低计算成本。
• 微调少于 0.1% 的 RoBERTa-base 参数即可实现与全量微调或定制架构模型相当甚至更优的性能。
• 在 SeqShort 中引入可学习位置嵌入可提升性能，优于缺乏位置信息的标准 MIL 方法。