[논문 리뷰] Swin MAE: Masked Autoencoders for Small Datasets

Zi'an Xu, Yin Dai외 6명|arXiv (Cornell University)|2022. 12. 28.

AI in cancer detection인용 수 11

한 줄 요약

Swin MAE는 소규모 의료 영상 데이터셋에서 비지도 표현 학습을 위한 스위프 트랜스포머 백본을 갖춘 마스킹된 오토인코더를 제안한다. 사전 훈련된 가중치 없이도 ImageNet에서 사전 훈련된 스위프 트랜스포머 모델과 동일하거나 略으로 뛰어난 전이 학습 성능를 달성하며, 마스킹된 영상 복원을 통해 단 몇천장의 이미지로도 강력한 특징 학습을 가능하게 한다.

ABSTRACT

The development of deep learning models in medical image analysis is majorly limited by the lack of large-sized and well-annotated datasets. Unsupervised learning does not require labels and is more suitable for solving medical image analysis problems. However, most of the current unsupervised learning methods need to be applied to large datasets. To make unsupervised learning applicable to small datasets, we proposed Swin MAE, which is a masked autoencoder with Swin Transformer as its backbone. Even on a dataset of only a few thousand medical images and without using any pre-trained models, Swin MAE is still able to learn useful semantic features purely from images. It can equal or even slightly outperform the supervised model obtained by Swin Transformer trained on ImageNet in terms of the transfer learning results of downstream tasks. The code is publicly available at https://github.com/Zian-Xu/Swin-MAE.

연구 동기 및 목표

• 딥 러닝에서 제한적이고 잘 레이블링된 의료 영상 데이터셋 문제를 해결한다.
• 비전 트랜스포머가 대규모 사전 훈련에 의존하는 문제를 해결하기 위해 스위프 트랜스포머를 통해 인도크티브 바이어스를 도입한다.
• 데이터 증강이 의료 영상에서는 종종 불가능하므로, 데이터 증강 없이도 소규모 데이터셋에서 효과적인 비지도 표현 학습을 가능하게 한다.
• 소규모 의료 영상에서의 자기지도 학습만으로도 하류 작업에서 경쟁적인 전이 학습 성능를 달성한다.
• 마스킹된 오토인코더를 사용한 Swin MAE가 사전 훈련된 가중치 없이도 ImageNet에서 미세조정된 지도 학습 모델과 동일하거나 뛰어난 성능를 보임을 입증한다.

제안 방법

• 사전 훈련 중 75%의 영상 패치를 무작위로 마스킹하는 마스킹 오토인코더(MAE) 프레임워크를 채택한다.
• 스위프 트랜스포머를 인코더 백본으로 사용하여 이격된 창문을 통해 국소적인 인도크티브 바이어스를 도입함으로써 소규모 데이터셋에서의 훈련 안정성을 향상시킨다.
• 16×16 패치와는 다름없이 4×4 겹치지 않는 패치를 사용한 패치 기반 영상 토크나이제이션을 적용하여 의료 영상 특성에 더 적합하게 한다.
• 전체 4×4 패치를 단위로 마스킹하는 윈도우 기반 마스킹 전략을 구현하여 복원 작업에서 공간 일관성을 유지한다.
• 원본 영상과 재구성된 영상 간의 평균 제곱오차(MSE) 손실을 사용하여 마스킹된 영역을 예측하는 데 중점을 둔다.
• 마스킹된 잠재 표현에서 전체 영상을 재구성하기 위해 경량 디코더 헤드를 구현하며, MAE와 Swin MAE 간의 공정한 비교를 위해 공통 아키텍처를 사용한다.

실험 결과

연구 질문

• RQ1 스위프 트랜스포머 백본을 갖춘 마스킹 오토인코더는 사전 훈련된 가중치 없이도 소규모 의료 영상 데이터셋에서 유용한 의미적 특징을 학습할 수 있는가?
• RQ2 동일한 소규모 데이터셋과 동일한 초모수를 사용할 때, 스위프 MAE는 표준 MAE보다 수렴 속도와 훈련 안정성 측면에서 어떻게 다른가?
• RQ3 스위프 MAE를 통한 비지도 사전 훈련은 지도 학습 모델의 ImageNet 미세조정 성능에 비해 하류 작업에서 얼마나 향상시킬 수 있는가?
• RQ4 스위프 트랜스포머의 인도크티브 바이어스는 표준 비전 트랜스포머에 비해 소규모, 자원이 제한된 의료 영상 데이터셋에서 자기지도 학습을 어떻게 향상시키는가?
• RQ5 의료 영상은 종종 단일 채널이고 손상되지 않은 상태이므로, 데이터 증강에 의존하지 않고도 마스킹 오토인코딩을 효과적으로 적용할 수 있는가?

주요 결과

• Swin MAE는 몇천 장의 의료 영상만으로도 하류 작업에서 ImageNet에서 사전 훈련된 지도 학습 스위프 트랜스포머 모델과 동일하거나 略으로 뛰어난 전이 학습 성능를 달성한다.
• 동일한 소규모 데이터셋과 동일한 초모수로 훈련했을 때, 표준 MAE에 비해 더 빠른 수렴 속도, 더 낮은 훈련 손실, 더 부드러운 훈련 곡선을 보인다.
• 스위프 트랜스포머의 인도크티브 바이어스 덕분에 스위프 MAE는 사전 훈련 없이도 훈련 안정성과 특징 학습 성능 향상을 달성하며, MAE를 능가한다.
• 재구성 결과는 75% 마스킹 비율 조건에서도 스위프 MAE가 손실된 영역, 특히 경계 윤곽과 색상 일관성을 효과적으로 복원함을 보여준다.
• 사전 훈련된 가중치가 없더라도 성능에 영향을 주지 않으며, 스위프 MAE는 소규모 데이터에서 자기지도 복원만으로도 강력한 표현을 학습한다.
• 업스트림 훈련 손실 곡선은 스위프 MAE가 소규모 데이터셋에서 MAE보다 더 나은 최적화 안정성과 더 빠른 수렴 속도를 보이며, 비지도 학습 능력이 뛰어남을 시사한다.