[24’ CVPR] Training Like a Medical Resident: Context-Prior Learning Toward Universal Medical Image Segmentation
카테고리: Medical
🔍 Abstract
Medical Dataset은 굉장히 heterogeneous하다. 이는 다양한 modalities와 anatomical structures를 포함하고 있기 때문이다. 이러한 특성 때문에 universal medical image segmentation은 매우 어려운 task이다. 이 논문에서는 이러한 문제를 단일 모델로 해결하기 위해 context-prior learning을 제안한다. 여기서는 특히 각 task와 modality에 대한 prompt learning을 통해 fine-tuning을 대체하고, 통일된 segmentation model을 학습하여 multi-dataset의 leverage를 극대화하고자 한다. 논문 제목의 굉장히 도발적인 “Training Like a Medical Resident”는 여러 데이터셋을 단일 모델로 잘 학습한다는 것을 의미한다.
1. Architecture
전체 구조는 다음과 같다. Backbone은 nnU-Net을 사용하였고, Context Prior 즉 일종의 Prompt를 task-specific 및 modality-specific하게 학습한다. 이때 데이터셋은 여러 task와 modality에 대해 사용할 수 있으므로 amortization을 수행할 수 있다는 장점이 있다. 자세한 과정 몇 가지를 살펴보면 다음과 같다.
먼저 Prior Selection은 이미지에 따라 Automatic하게 이루어진다. Prior를 모델이 선택하는 것이 아니라, 이미지의 메타 정보로부터 선택하는 것이다. 그리고 Posterior Prototype을 얻을 때 MLP를 거치게 되는데, 이 과정에서 Multi-scale feature가 aggregation된다. 이를 저자들은 Hierarchical modeling이라고 부른다. 이러한 Multi-scale approach도 성능 향상에 도움이 되었다. 마지막으로 Mask prediction이 아닌 modality prediction이 도움이 되는지 직관적으로 와닿지 않았는데, 일종의 modality를 더 잘 구분하기 위한 전략이라고 생각하면 될 것 같다.
참고로 데이터셋은 다음과 같은 3D 데이터셋이다. 3D 치고는 많은 편이다.
2. Experiments
2.1. Previous Methods
결과를 보기에 앞서, 기존의 방법론들을 이해해야 이 모델의 Contribution을 더 자세히 이해할 수 있기에 기존의 방법론들 몇 가지를 소개한다.
첫 번째로 소개할 모델은 2023년 ICCV에 소개된 CLIP-Driven이다. 이 논문에서는 Pretrained CLIP을 사용하여 Organ에 대한 prior knowledge를 얻고, 이를 이용하여 organ segmentation을 수행한다. 본 논문에서는 기존의 organ segmentation에 대한 embedding, 즉 일종의 task condition에 modality condition을 추가하여 구현하였다. 그러나 CLIP의 medical task prior knowledge가 부족했기에 성능이 좋지 않았다.
다음으로 소개할 모델은 2023년 MICCAI에 소개된 MultiTalent이다. 이 논문에서는 동일한 nnU-Net으로 다양한 데이터를 훈련했고, multiple segmentation head를 사용하여 각 task에 대한 segmentation을 수행하였다. 그러나 prior knowledge를 inference 때 사용할 방법이 없었기에 성능이 좋지 않았다.
마지막으로 소개할 모델은 2023년 MICCAI에 소개된 UniSeg이다. 현 논문과 가장 비슷한 Approach로 볼 수 있다. 위와 같이 Image sample-specific prompt를 사용하는 방법인데, prompt를 lowest resolution에서 얻어 multi-scale information을 얻지 못한다는 점에서 성능이 좋지 않았다.
2.2. Main Results
전체 결과는 위와 같다. 기존의 Task-specific model에 비해 amortized training으로 성능을 높일 수 있었고, 기존의 universal medical segmentation 모델에 비해 prompt learning을 사용하는 방법이 더 robust하다는 것을 입증하였다.
2.3. Further Analysis
다음으로 주목하여 본 것은 prior knowledge prompt 분석이다. 일반적인 CLIP embedding에 비해 Hermes의 prior는 서로 뚜렷한 특징을 가지고 있고, 실제 유사도를 잘 반영하고 있다. 이는 prompt learning이 prior knowledge를 잘 반영하고 있다는 것을 보여준다.
💡 Summary
지금까지 내용을 정말 한 마디로 요약하면 다음과 같다. Prompt learning을 통해 task, modality-specific condition을 제공하면서 amortized training을 수행하는 방법론이 universal medical image segmentation에서 좋은 성능을 보였다. 이러한 방법이 prior knowledge를 잘 반영한다는 점에서는 솔직히 동의하기 어렵다. 논문 제목을 보고 어떤 anatomical knowledge를 잘 leverage하는 방법론을 제안한다고 생각했으나 결국은 dataset-driven prompt learning에 불과했기에 아쉬움이 남는다.
📃 Reference
- [24’ CVPR] Training Like a Medical Resident: Context-Prior Learning Toward Universal Medical Image Segmentation
- [23’ ICCV] CLIP-Driven Universal Model for Organ Segmentation and Tumor Detection
- [23’ MICCAI] MultiTalent: A Multi-dataset Approach to Medical Image Segmentation
- [23’ MICCAI] UniSeg: A Prompt-Driven Universal Segmentation Model as Well as A Strong Representation Learner
댓글 남기기