[논문 리뷰] 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)|2022. 11. 14.

AI in cancer detection인용 수 5

한 줄 요약

이 논문은 전체 슬라이드 이미지(WSIs)를 고정 길이의 순서 있는 특징 벡터 시퀀스로 압축하여, 텍스트 사전학습된 트랜스포머를 사용한 효율적이고 정확한 분류를 가능하게 하는 다중헤드 어텐션 기반 시퀀스 압축 레이어인 SeqShort을 제안한다. 이와 함께 동결된 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.

연구 동기 및 목표

• 장기적인 입력 시퀀스인 이미지 패치로 인해 트랜스포머를 전체 슬라이드 이미지(WSIs)에 적용할 경우 발생하는 높은 계산 비용을 해결하기 위해.
• 분야 특화의 미세조정이나 아키텍처 재설계 없이도 강력한 사전학습된 언어 모델을 디지털 병리학에 활용할 수 있도록 하기 위해.
• 학습 가능한 위치 인코딩과 어텐션 기반 요약을 통해 WSI 임베딩에서 공간적 및 맥락적 정보를 유지하기 위해.
• 성능을 유지하거나 향상시키면서도 학습 가능한 파라미터의 수를 줄이기 위해.
• 최소한의 미세조정으로 텍스트 사전학습된 트랜스포머가 조직병리학 분류에 효과적으로 재사용될 수 있는지 보여주기 위해.

제안 방법

• WSI 패치 수천 개를 고정 길이의 순서 있는 특징 벡터 시퀀스로 압축하는 데 사용되는 다중헤드 어텐션 기반 시퀀스 단순화 레이어(이하 SeqShort)를 도입한다.
• SeqShort은 전체 패치 임베딩을 대상으로 주목사의 기반을 가진 쿼리 벡터를 사용하여, 자기 어텐션 메커니즘을 통해 중요한 특징을 포착한다.
• 표준 MIL 접근 방식의 한계를 보완하기 위해 공간 순서를 유지하기 위해 학습 가능한 위치 임베딩을 통합한다.
• 사전학습된 RoBERTa-base 언어 모델을 트랜스포머 인코더로 사용하며, 총 파라미터의 0.04%만 미세조정한다.
• EfficientNetV2-L 백본을 사용해 ImageNet에서 사전학습된 모델을 통해 WSI에서 인스턴스 수준의 특징을 추출한다.
• 최종 분류 헤드는 트랜스포머 인코더의 [CLS] 토큰 표현을 사용하는 다층퍼셉트론(MLP)이다.

실험 결과

연구 질문

• RQ1 어떤 어텐션 기반 시퀀스 단순화 기법이 분류에 유의미한 정보를 유지하면서도 WSI 패치 시퀀스를 효과적으로 압축할 수 있는가?
• RQ2 RoBERTa와 같은 텍스트 사전학습된 트랜스포머를 사용할 경우, 분야 특화의 사전학습 없이도 WSI 분류 성능이 향상되는가?
• RQ3 0.1% 미만의 최소한의 파라미터만 미세조정해도 WSI 분류 과제에서 최신 기술 수준의 성능을 달성할 수 있는가?
• RQ4 표준 MIL 방법과 비교했을 때, 시퀀스 단순화 레이어에 학습 가능한 위치 임베딩을 포함시키면 성능에 어떤 영향을 미치는가?
• RQ5 일반화된, WSI에 특화되지 않은 트랜스포머 아키텍처가 효율적인 시퀀스 압축과 결합될 경우, 디지털 병리학에서 뛰어난 성능을 달성할 수 있는가?

주요 결과

• 제안된 방법은 CAMELYON16 림프절 전이 분류 벤치마크에서 AUC 0.948, 정확도 0.882로 최신 기술 수준의 성능을 달성하였다.
• TCGA-BRCA 침습성 유방암 아형 분류 과제에서 10겹 교차검증을 사용하여 AUC 0.917, 정확도 0.841을 기록하였다.
• 신장세포암 아형 분류 과제에서는 AUC 0.892, 정확도 0.815를 달성하여 다양한 암 유형 간의 강건성을 입증하였다.
• 학습 가능한 파라미터 수를 8500만 개에서 330만 개(총 파라미터의 0.04%)로 줄여 계산 비용을 크게 감소시켰다.
• RoBERTa-base 파라미터의 0.1% 미만을 미세조정함으로써, 전체 미세조정 또는 고유 설계된 아키텍처를 가진 모델들과 비교해도 성능이 유사하거나 뛰어나다.
• SeqShort에 학습 가능한 위치 임베딩을 포함시킴으로써, 위치 정보가 없는 표준 MIL 방법보다 성능 향상이 이루어졌다.