[논문 리뷰] Adapting Pre-trained Vision Transformers from 2D to 3D through Weight Inflation Improves Medical Image Segmentation

Yuhui Zhang, Shih-Cheng Huang외 3명|arXiv (Cornell University)|2023. 02. 08.

Advanced Neural Network Applications인용 수 4

한 줄 요약

이 논문은 사전 훈련된 2D 비전 트랜스포머를 3D 의료 영상 분할에 적응시키기 위한 단순하면서도 효과적인 가중치 팽창 전략을 제안한다. 이는 전이 학습을 가능하게 하면서 깊이 정보를 유지한다. 중심화 초기화를 통해 2D 가중치를 3D로 팽창하고, 중심 영역 예측을 수행하는 작은 3D 창을 사용함으로써, 계산 비용이 단지 0.75% 증가하는 데 그치면서도 12개의 다양한 3D 의료 영상 데이터셋에서 최신 기술 수준의 성능을 달성한다.

ABSTRACT

Given the prevalence of 3D medical imaging technologies such as MRI and CT that are widely used in diagnosing and treating diverse diseases, 3D segmentation is one of the fundamental tasks of medical image analysis. Recently, Transformer-based models have started to achieve state-of-the-art performances across many vision tasks, through pre-training on large-scale natural image benchmark datasets. While works on medical image analysis have also begun to explore Transformer-based models, there is currently no optimal strategy to effectively leverage pre-trained Transformers, primarily due to the difference in dimensionality between 2D natural images and 3D medical images. Existing solutions either split 3D images into 2D slices and predict each slice independently, thereby losing crucial depth-wise information, or modify the Transformer architecture to support 3D inputs without leveraging pre-trained weights. In this work, we use a simple yet effective weight inflation strategy to adapt pre-trained Transformers from 2D to 3D, retaining the benefit of both transfer learning and depth information. We further investigate the effectiveness of transfer from different pre-training sources and objectives. Our approach achieves state-of-the-art performances across a broad range of 3D medical image datasets, and can become a standard strategy easily utilized by all work on Transformer-based models for 3D medical images, to maximize performance.

연구 동기 및 목표

• 사전 훈련된 2D 비전 트랜스포머를 3D 의료 영상 분할에 적응시키는 데 있어 전이 학습의 이점을 유지하면서 깊이 정보도 보존하는 데 도전하는 것.
• 가중치 팽창이 2D 자연 영상과 3D 의료 스캔 간의 도메인 갭을 효과적으로 메울 수 있는지 조사하는 것.
• 3D 의료 비전 작업에서 사전 훈련된 2D 모델을 활용하기 위한 일반화 가능하고 표준화된 전략을 수립하는 것.
• 다양한 사전 훈련 소스와 목표가 3D 분할 성능에 미치는 영향을 평가하는 것.
• 최소한의 하이퍼파라미터 튜닝으로도 다양한 3D 의료 영상 데이터셋에서 일관된 성능 향상을 보여주는 것.

제안 방법

• 중심화 전략을 사용하여 사전 훈련된 2D 비전 트랜스포머의 학습된 가중치를 팽창함으로써 2D 모델을 3D로 적응시키는 방식. 이때 중심 슬라이스의 가중치는 복사하고, 나머지는 0으로 초기화한다.
• 3D 입력 볼륨을 몇 개의 인접 슬라이스를 포함하는 작은 3D 창으로 분할하여 공간적 맥락을 유지한다.
• 각 창 내에서 중심 슬라이스의 분할 예측만을 훈련 목표로 삼음으로써 잔차 학습을 모방하고 깊이 정보를 점진적으로 통합할 수 있도록 한다.
• 팽창된 2D 트랜스포머를 인코더로 사용하는 표준 3D U-Net 스타일 아키텍처를 사용하여 엔드 투 엔드 분할을 수행한다.
• 다양한 팽창 전략을 비교하기 위해 추론 실험을 수행하며, 이는 예측 대상, 가중치 초기화 방식, 스트라이드, 입력 슬라이스 수를 포함한다.
• 다양한 해부학적 영역과 영상 모odalities(CT/MRI)를 가진 11개의 추가 3D 의료 영상 데이터셋에서 일반화 성능을 평가한다.

Figure 1: Approach overview. Large-scale pre-trained Transformers are used as the encoder in the segmentation model for transfer learning, in which weights are adapted using the inflation strategy to support 3D inputs. Each 3D image is split into windows, which contain a small number of neighbor sli

실험 결과

연구 질문

• RQ1 아키텍처 수정 없이도 가중치 팽창 전략이 사전 훈련된 2D 비전 트랜스포머를 3D 의료 영상 분할에 효과적으로 적응시킬 수 있는가?
• RQ2 사전 훈련 소스의 선택(자기지도 학습 대비 지도 학습)이 3D 의료 분할에서 전이 성능에 어떤 영향을 미치는가?
• RQ3 가중치 팽창을 통한 3D 적응에 있어 최적의 창 크기와 예측 전략(예: 중심 슬라이스만)은 무엇인가?
• RQ4 제안된 방법이 다양한 해부학적 구조와 영상 프로토콜을 가진 다양한 3D 의료 영상 데이터셋에 일반화되는가?
• RQ5 기본적인 2D 또는 3D 무작위 초기화 대비 가중치 팽창이 성능 향상에 얼마나 기여하는가?

주요 결과

• 제안된 가중치 팽창 전략은 계산 비용이 단지 0.75% 증가하는 데 그치면서도 강력한 베이스라인 대비 평균 DSC를 1.95% 향상시킨다.
• 사전 훈련된 가중치로 초기화된 모델은 다기관 분할 데이터셋에서 무작위 초기화된 3D 모델보다 11.18% 높은 성능을 보이며, 전이 학습의 중요성을 입증한다.
• 중심화 팽창 전략—2D 가중치를 중심 슬라이스에 복사하고 나머지는 0으로 초기화—가 가장 높은 성능을 달성함으로써 점진적인 깊이 통합 전략이 효과적임을 시사한다.
• 몇 개의 인접 슬라이스(예: 3~5개)만 포함시켜도 성능이 충분히 확보되며, 인접 슬라이스는 먼 슬라이스보다 더 관련성이 높은 맥락을 제공한다.
• 12개의 데이터셋에서 최신 기술 수준의 성능을 달성하였고, 주요 평가 외 11개의 추가 데이터셋 중 10개에서도 성능을 확보하여 강력한 일반화 능력을 입증한다.
• 자기지도 학습 대비 지도 학습 소스가 더 높은 전이 성능을 보이며, 레이블이 풍부한 사전 훈련이 3D 의료 분할에 더 유리함을 시사한다.

Figure 2: Transfer learning reduces noise in local predictions, and depth information further improves prediction consistency across slices. We show five consecutive input images, ground-truth labels, predictions from Ours , Ours w/o D , Ours w/o T&D .