[논문 리뷰] ViT-AE++: Improving Vision Transformer Autoencoder for Self-supervised Medical Image Representations

Chinmay Prabhakar, Hongwei Li외 4명|arXiv (Cornell University)|2023. 01. 18.

AI in cancer detection인용 수 10

한 줄 요약

ViT-AE++는 보조 재구성 손실과 대조 손실로 비전 트랜스포머 오토인코더를 확장하여 2D 및 3D 의료 영상 모두에 대해 더 나은 자기지도 표현을 학습하고, 3D 부피로 접근 방식을 확장합니다.

ABSTRACT

Self-supervised learning has attracted increasing attention as it learns data-driven representation from data without annotations. Vision transformer-based autoencoder (ViT-AE) by He et al. (2021) is a recent self-supervised learning technique that employs a patch-masking strategy to learn a meaningful latent space. In this paper, we focus on improving ViT-AE (nicknamed ViT-AE++) for a more effective representation of 2D and 3D medical images. We propose two new loss functions to enhance the representation during training. The first loss term aims to improve self-reconstruction by considering the structured dependencies and indirectly improving the representation. The second loss term leverages contrastive loss to optimize the representation from two randomly masked views directly. We extended ViT-AE++ to a 3D fashion for volumetric medical images as an independent contribution. We extensively evaluate ViT-AE++ on both natural images and medical images, demonstrating consistent improvement over vanilla ViT-AE and its superiority over other contrastive learning approaches. Codes are here: https://github.com/chinmay5/vit_ae_plus_plus.git.

연구 동기 및 목표

• 레이블이 드문 의료 영상에서 자기지도 표현 학습을 개선하기 위한 동기 부여.
• 구조 및 에지 정보를 포착하는 보조 재구성 손실을 도입.
• 자동인코더 기반 학습과 대조학습(paradigms)을 결합하여 표현을 직접 최적화.
• MRI/CT 유사 모달리티를 다루기 위해 ViT-AE++를 3D 볼륨 데이터로 확장.
• ViT-AE와의 일관된 성능 향상과 대조 방법과의 경쟁력 있는 결과를 보여줌.

제안 방법

• 랜덤 마스킹된 패치와 누락된 패치를 위한 MASK 토큰을 사용하는 ViT 기반 오토인코더를 이용한다.
• 구조 정보를 포착하기 위해 보조 손실: L_per (VGG 기반 지각 손실) 및 L_edge (3D Sobel 에지 손실)을 도입한다.
• 두 개의 무작위 마스킹 뷰 사이에서 대조 손실 L_CL을 추가하여 표현을 직접 최적화한다.
• 손실들을 총 목적 함수에 합산: L = L_rec + λ1 L_per + λ2 L_edge + L_CL로 총 목적 함수를 구성하고, λ2 감소를 스케줄링한다.
• 3D 패치, 3D 위치 인코딩, 3D Transformer 디코더를 적용하여 프레임워크를 3D로 확장한다.
• 선형 탐색(linear probing)으로 여러 데이터셋에서 평가하고 AUC/정확도를 다운스트림 태스크 지표로 보고한다.

실험 결과

연구 질문

• RQ1 보조 구조 인식 재구성이 픽셀 단위 손실을 넘어서 의미론적 내용 캡처를 개선할 수 있는가?
• RQ2 오토인코더 재구성과 함께 대조 목표를 도입하면 의료 영상에 대해 더 나은 표현을 얻을 수 있는가?
• RQ3 ViT-AE++의 3D 확장이 볼륨 의료 데이터에서 지속적으로 개선을 보여주는가?
• RQ4 마스킹 비율과 에지/지각 손실 가중치가 학습 안정성과 표현 품질에 어떤 영향을 미치는가?

주요 결과

• ViT-AE++는 2D 데이터 세트에서 일반적으로 ViT-AE를 능가한다(CIFAR-10: 95.40 vs 94.10; CIFAR-100: 78.82 vs 75.61; Tiny ImageNet-100: 72.09 vs 70.42; Chest X-ray: 95.60 vs 95.20).
• 3D 데이터셋에서 ViT-AE++는 BraTS에서 선도적 대조 방법에 비해 경쟁력 있거나 우수한 결과를 달성하고 EGD에서 이를 능가한다(BraTS: 0.767 AUC; EGD: 0.846 AUC).
• 일화 연구는 보조 에지 및 지각 손실과 대조 손실이 ViT-AE 대비 개선을 jointly 주도한다.
• 마스킹 비율 p와 감소하는 에지 손실 가중치 λ2가 성능 및 학습 안정성에 중요하다; 최적 p는 대략 0.75이고 λ2의 선형 감소가 학습을 안정시킨다.
• 3D ViT-AE++는 의학 데이터셋에서 선형 탐색기로 사용될 때 강한 다운스트림 성능을 보여주고 최첨단 방법과 경쟁력을 유지한다.