[논문 리뷰] Self Pre-training with Masked Autoencoders for Medical Image Classification and Segmentation

Lei Zhou, Huidong Liu외 4명|arXiv (Cornell University)|2022. 03. 10.

Radiomics and Machine Learning in Medical Imaging인용 수 8

한 줄 요약

이 논문은 외부 데이터셋(예: ImageNet)에 의존하지 않고 대상 데이터셋의 훈련 데이터만을 사용하여 의료 영상 분석을 위한 자기지도 학습 마스크된 오토인코더(MAE) 사전 훈련 방법을 제안한다. 3D 영상 패치의 12.5–75%를 마스킹하고 비전 트랜스포머를 통해 이를 재구성함으로써 맥락적 특징 학습을 향상시켜 흉부 X-ray 분류, Abdominal CT 세그멘테이션, MRI 뇌 종양 세그멘테이션에서 최신 기술 수준의 성능을 달성하며, 무작위 초기화 및 ImageNet 사전 훈련 대비 뚜렷한 성능 향상을 보였다.

ABSTRACT

Masked Autoencoder (MAE) has recently been shown to be effective in pre-training Vision Transformers (ViT) for natural image analysis. By reconstructing full images from partially masked inputs, a ViT encoder aggregates contextual information to infer masked image regions. We believe that this context aggregation ability is particularly essential to the medical image domain where each anatomical structure is functionally and mechanically connected to other structures and regions. Because there is no ImageNet-scale medical image dataset for pre-training, we investigate a self pre-training paradigm with MAE for medical image analysis tasks. Our method pre-trains a ViT on the training set of the target data instead of another dataset. Thus, self pre-training can benefit more scenarios where pre-training data is hard to acquire. Our experimental results show that MAE self pre-training markedly improves diverse medical image tasks including chest X-ray disease classification, abdominal CT multi-organ segmentation, and MRI brain tumor segmentation. Code is available at https://github.com/cvlab-stonybrook/SelfMedMAE

연구 동기 및 목표

• 의료 영상 분야에서 대규모 도메인 전용 사전 훈련 데이터의 부족 문제를 해결하기 위해, 대상 데이터셋의 훈련 데이터만을 사용하는 자기지도 사전 훈련 프레임워크를 제안하는 것.
• 해부학적 구조가 기능적·기계적으로 상호의존하는 의료 영상에서 마스크된 오토인코딩이 맥락적 특징 학습을 향상시킬 수 있는지 조사하는 것.
• 낮은 데이터 환경에서 다양한 3D 의료 영상 작업(분류 및 세그멘테이션 포함)에 대해 MAE 사전 훈련의 효과를 평가하는 것.
• 의료 영상 기준 평가 벤치마크에서 MAE 자기사전훈련이 ImageNet 사전 훈련 및 무작위 초기화 대비 최종 성능에서 어떻게 성능을 내는지 비교하는 것.
• 의료 영상 작업에 최적화된 하이퍼파rameter(예: 마스크 비율, 사전 훈련 에포크 수)를 결정하는 것.

제안 방법

• MAE 재구성 과정에서 간섭을 방지하기 위해, 사인-余弦 위치 임베딩을 사용한 비전 트랜스포머(ViT) 백본을 사용한다.
• 마스크된 오토인코더 아키텍처를 적용하여 3D 영상 패치의 12.5%에서 75%를 무작위로 마스킹하고, 가시 패치와 학습 가능한 마스크 토큰만을 사용해 이를 재구성한다.
• 에코더는 오직 가시 패치만 처리하여 패치 단위의 표현을 생성하고, 경량 디코더는 이러한 표현과 마스크 토큰에서 전체 영상을 재구성한다.
• 모델은 마스킹된 패치에 대해 평균 제곱오차 손실을 사용하여, 주변 패치의 맥락 정보를 활용해 누락된 영역을 추론하도록 유도한다.
• 사전 훈련 후, 에코더 가중치는 최종 모델을 초기화하는 데 사용되며, 분류 작업에는 선형 분류기, 세그멘테이션 작업에는 UNETR 유사 디코더를 사용한다.
• 피니튜닝은 계층별 학습률 감소와余弦 학습률 스케줄링을 적용하며, 각 데이터셋에 맞는 작업별 하이퍼파라미터를 설정한다.

Fig. 1 : Segmentation Pipeline with MAE Self Pre-training. Left: A ViT encoder is first pre-trained with MAE. A random subset of patches is input to the encoder and a transformer decoder reconstruct the full image. Right: The pre-trained ViT weights are transferred to initialize the segmentation enc

실험 결과

연구 질문

• RQ1 대상 데이터셋에서만 자기지도 마스크된 오토인코딩 사전 훈련을 수행할 경우, 의료 영상 분류 및 세그멘테이션 작업의 성능 향상이 가능할까?
• RQ2 소규모 의료 영상 데이터셋에 적용했을 때, MAE 사전 훈련이 ImageNet 사전 훈련을 능가할 수 있을까?
• RQ3 특히 3D CT 및 MRI 데이터에 대해 의료 영상 세그멘테이션 작업에서 MAE 사전 훈련에 최적의 마스크 비율은 무엇일까?
• RQ4 자기사전훈련은 종양 미세환경과 같이 해부학적으로 연결된 영역 간 맥락적 의존성을 포착하는 데 모델의 능력에 어떤 영향을 미칠까?
• RQ5 사용된 위치 임베딩 방식(사인-余弦 vs. 학습 가능한)이 의료 영상 MAE 사전 훈련에서 재구성 품질과 최종 성능에 영향을 미칠까?

주요 결과

• BTCV 데이터셋에서 복수의 복부 장기 세그멘테이션 작업에서 MAE 자기사전훈련은 평균 DSC를 베이스라인 78.8%에서 83.5%로 향상시켜 ImageNet 사전 훈련보다 뚜렷이 뛰어난 성능을 보였다.
• BraTS MRI 뇌 종양 세그멘테이션 작업에서 MAE 사전 훈련은 DSC를 77.4%에서 78.91%로 향상시키고, HD95를 7.78mm에서 7.22mm로 감소시켰다.
• CXR14 흉부 X-ray 분류 작업에서 MAE 자기사전훈련은 ImageNet 사전 훈련보다 0.8% 높은 성능을 기록했으며, 자기지도 사전 훈련을 적용한 최신 기술 수준의 CNN 기반 모델을 초월했다.
• 최적의 세그멘테이션 성능는 12.5% 마스크 비율에서 달성되었고, 높은 비율(예: 75%)은 성능 향상에 덜 효과적이었으며, 이는 작업에 맞는 최적화가 필요함을 시사한다.
• 더 긴 사전 훈련 에포크 수는 성능 향상에 기여했지만, 100–150 에포크 이후 과적합 현상이 발생했으며, 특히 소규모 데이터셋에서 두드러졌다.
• 사전 훈련 단계에서 사인-余弦 위치 임베딩을 사용할 경우, 학습 가능한 임베딩보다 재구성 품질과 최종 성능이 향상되었다.

Fig. 2 : MAE reconstruction. First row: Original image. Second row: Masked image where masked regions are colored with gray/black. Third row: Reconstructed images from unmasked patches. Every two columns show the results of one dataset, i.e., CXR, BTCV and BraTS from left to right.