[논문 리뷰] Scaling Vision Transformers to Gigapixel Images via Hierarchical Self-Supervised Learning

Richard J. Chen, Chengkuan Chen외 5명|arXiv (Cornell University)|2022. 06. 06.

Cell Image Analysis Techniques인용 수 8

한 줄 요약

이 논문은 전임상 병리학에서 거대한 전체 슬라이드 이미지(WSI)를 표현하기 위해 다중 척도 자기지도 학습을 활용하는 계층적 비전 트랜스포머 아키텍처인 HIPT를 소개한다. 10,678개의 WSI와 33종의 암 유형을 대상으로 104M개의 $256\times256$ 패치와 408k개의 $4096\times4096$ 영역을 사용해 사전 훈련한 HIPT는 암 아형 분류 및 생존 예측과 같은 슬라이드 수준 작업에서 최신 기술 수준(SOTA) 성능을 달성하였으며, 종양 미세환경 내에서 중요한 공간적 및 표현형 인도적 편향을 계층적 자기지도 학습이 포착할 수 있음을 보여준다.

ABSTRACT

Vision Transformers (ViTs) and their multi-scale and hierarchical variations have been successful at capturing image representations but their use has been generally studied for low-resolution images (e.g. - 256x256, 384384). For gigapixel whole-slide imaging (WSI) in computational pathology, WSIs can be as large as 150000x150000 pixels at 20X magnification and exhibit a hierarchical structure of visual tokens across varying resolutions: from 16x16 images capture spatial patterns among cells, to 4096x4096 images characterizing interactions within the tissue microenvironment. We introduce a new ViT architecture called the Hierarchical Image Pyramid Transformer (HIPT), which leverages the natural hierarchical structure inherent in WSIs using two levels of self-supervised learning to learn high-resolution image representations. HIPT is pretrained across 33 cancer types using 10,678 gigapixel WSIs, 408,218 4096x4096 images, and 104M 256x256 images. We benchmark HIPT representations on 9 slide-level tasks, and demonstrate that: 1) HIPT with hierarchical pretraining outperforms current state-of-the-art methods for cancer subtyping and survival prediction, 2) self-supervised ViTs are able to model important inductive biases about the hierarchical structure of phenotypes in the tumor microenvironment.

연구 동기 및 목표

• 기존의 약한 지도 학습 다중 인스턴스 학습(MIL) 프레임워크가 $256\times256$ 패치 수준 표현에 국한되고 장거리 의존성에 대한 주의 메커니즘이 부족한 점을 해결하기 위해.
• 여러 해상도에서의 고유한 계층적 구조를 활용하여 거대한 전체 슬라이드 이미지(WSI)에서 효과적인 슬라이드 수준 표현 학습을 가능하게 하기 위해.
• 자기지도 학습을 통해 토큰화 및 집계 구성 요소를 동시에 사전 훈련하는 비전 트랜스포머 아키텍처를 개발하여, 데이터가 적은 환경에서의 일반화 능력을 향상시키기 위해.
• 계층적 자기지도 사전 훈련이 종양 미세환경 내에서 생물학적으로 의미 있는 인도적 편향, 예를 들어 종양-면역 세포 국소화 및 공간적 표현형 조직화를 포착하는지 검증하기 위해.

제안 방법

• 하나의 패치 수준에서 $16\times16$ 세포 수준 토큰에서부터 $256\times256$ 및 $4096\times4096$ 영역 표현으로까지 하향식 집합을 수행하는 3단계 아키텍처인 계층적 이미지 피라미드 트랜스포머(HIPT)를 제안한다.
• 두 단계 자기지도 사전 훈련: 먼저 $256\times256$ 패치에서, 그 다음 $4096\times4096$ 영역에서, 인스턴스 식별을 사용한 대비 학습을 활용한다.
• $16\times16$ 패치 임bedding을 기반 특징 추출기로 사용하는 비전 트랜스포머(ViT) 백본을 사용하며, $256\times256$ 및 $4096\times4096$ 척도에 대해 별도의 ViT 인코더를 사용한다.
• 요소 분할 주의 메커니즘을 적용하여 계층적 주의 맵을 시각화함으로써 종양 세포 및 면역 세포 국소화의 해석 가능성을 높인다.
• 단계별 사전 훈련 수행: 먼저 104M개의 인스턴스에 대해 $256\times256$ 패치에서, 그 다음 408k개의 인스턴스에 대해 $4096\times4096$ 영역에서 수행한다.
• 마지막 4개의 ViT 레이어에서의 [CLS] 토큰을 사용해 특징을 연결하지만, 성능 향상은 관찰되지 않았다.

실험 결과

연구 질문

• RQ1 표준 MIL 프레임워크에 비해 계층적 자기지도 사전 훈련을 통한 비전 트랜스포머의 성능 향상이 거대한 전체 슬라이드 이미지에서 슬라이드 수준 표현 학습에 기여하는가?
• RQ2 16×16 세포에서부터 4096×4096 조직 영역까지의 다중 척도 시각 토큰을 모델링함으로써 종양 미세환경 내 생물학적으로 관련성이 있는 인도적 편향을 포착할 수 있는가?
• RQ3 전암(pretraining)과 기관 특이적 사전 훈련 간의 일반화 능력과 주의 맵의 해석 가능성 측면에서 어떤 차이가 있는가?
• RQ4 자기지도 학습 비전 트랜스포머가 암 아형 분류 및 생존 예측 작업에서 지도 학습 또는 약한 지도 학습 기반 MIL 기준선을 능가할 수 있는가?
• RQ5 계층적 주의 메커니즘은 종양 둥지, 간질 침범, 면역 세포 침윤과 같은 임상적으로 중요한 특징을 얼마나 정확하게 국소화하는가?

주요 결과

• 계층적 사전 훈련을 적용한 HIPT는 9개의 슬라이드 수준 작업에서 최신 기술 수준 성능을 달성하였으며, 암 아형 분류 및 생존 예측에서 현재 최고 성능 기준선을 초월한다.
• 33종의 암 유형에서 전암 사전 훈련은 기관 특이적 사전 훈련보다 더 우수한 일반화 능력을 보였으며, 특히 다양한 조직병리학적 패atters에서 종양 세포와 림프구 세포의 국소화에 유리하다.
• 해석 가능성 면에서 뛰어난 성능을 보이며, 계층적 주의 맵이 디모플라스틱 간질 내 침습성 종양 세포와 침습성 유방암과 유선세포암을 구별하는 데 핵심적인 종양 둥지 패턴을 성공적으로 국소화하였다.
• BRCA 아형 분류에서 전암 사전 훈련을 적용한 HIPT는 BRCA 특이적 사전 훈련을 초월하였으며, 이는 유사한 표현형을 가진 아형을 구별하기 위해 필요한 간질 맥락을 더 잘 포착했기 때문일 것이다.
• 자기지도 학습 비전 트랜스포머는 조직 미세환경 전반에 걸친 장거리 공간적 의존성을 모델링할 수 있었으며, 종양-면역 세포 국소화에 주의를 기울임으로써 정확한 생존 예측이 가능했다.
• ViT 레이어 간 특징 연결은 성능 향상에 기여하지 않았으며, 이는 최종 레이어의 [CLS] 토큰만으로도 후행 작업에 충분하다는 것을 시사한다.