Paper: https://arxiv.org/abs/2410.10604
Multi-modal Vision Pre-training for Medical Image Analysis
Self-supervised learning has greatly facilitated medical image analysis by suppressing the training data requirement for real-world applications. Current paradigms predominantly rely on self-supervision within uni-modal image data, thereby neglecting the i
arxiv.org
요즘 Multi-modal, 여러 Data를 한번에 활용할 때 Fusion하는 방식에 대해 관심이 많아 정리중이며, CVPR 2025 Paper들을 훑어보고있다.
이 페이퍼 역시 CVPR 2025에 게재된 논문이다. 코드는 아래에서 확인 가능하다.
https://github.com/openmedlab/BrainMVP
1. Introduction
본 논문에서는 Medical Image Analysis에서 Self-supervised Learning(SSL)을 활용하는 데에 초점을 둔다.
SSL을 활용해 라벨이 없는 데이터에서도 효과적인 representation learning을 수행하고, 이후 Downstream task에서 높은 성능을 확보할 수 있도록 Generalization하는 게 목적이다.
저자들은 해당 도메인에서 SSL을 다룬 연구들이 unimodal에 의존하고 있고, 다양한 modal에 대한 정보를 고려하지 못한다는 한계를 지적한다. (이 페이퍼에서 말하는 multimodal은 MRI 이미지의 모달리티 (T1, T2, FLAIR 등)이지만 일반적인 multimodal data와 관련해서도 충분히 참고할만한 점이 있다고 생각된다.)
따라서 저자들은 새로운 Multimodal vision pre-training framework를 제안한다.
연구에서는 Fig 1에 요약된 바와 같이, 아래 세가지 proxy task를 활용한다.
- Cross-modal Image Reconstruction
- 한 modal의 일부를 다른 modal의 픽셀로 마스킹
- 원래 이미지를 복원하게 하여 modal간 상관관계를 학습
- Modality-wise Data Distillation
- 각 modal에 대해 학습 가능한 modality template을 생성하여 구조적 정보를 요약
- 사전학습과 다운스트림 간의 Bridge 역할
- Modality-aware Contrastive Learning
- Reconstruction과 Distillation으로 생성된 두 representation 간의 일관성(feature-level consistency)을 유지하도록 contrastive loss를 사용
2. Method
Fig 2는 세가지 proxy task로 구성된 프레임워크의 전체적인 구조를 보여준다.
논문에 설명이 기재된 순서대로 확인해보겠다.
Cross-modal Reconstruction
이름 그대로 서로 다른 modality 간의 관계를 학습하여 Reconstruction을 수행하며, representation을 학습하는 방법이다.
(Cross-Modal Masking) 한 Modality의 이미지 일부를 다른 Modality로 마스킹하고, 원래 Modality의 이미지를 복원하도록 학습한다. Loss는 아래와 같다. Masked Image를 인코딩 한 후 Decoding 해 원래 이미지를 Reconstruct 후 계산한 복원 오차이다.
저자들은 서로 다른 Modality간 공통 구조는 존재하며, 각 Modality가 보완적인 정보를 담고 있다는 점에서 이런 학습 방식을 활용했다. (즉, 이런 데이터에서 유효한 과정일 가능성이 크다!)
Modality-wise Data Distllation
(설명이 길어져 항목을 나눴다.)
Motivation
Foundation model은 highly generalizable latent representations를 추출하는 것이 목표다.
하지만 pre-training에 사용되는 proxy task는 downstream task와 관련이 없는 경우가 많고, generalization에 악영향을 미친다.
따라서, 저자들은 downstream 작업에 적용될 때 pre-trained model의 representation이 용이하기를 바라며
pre-train 단계에서 specific representation을 얻는데 필요한 guide를 제공하는 bridging component를 도입하기 위해 이 방법을 제시한다.
Method
Fig 2에도 나타나듯이, 이 작업은 Cross-modal reconstruction과 함께 수행된다.
우선, learnable 즉 학습 가능한 모달별 template 을 0으로 초기화한다.
Cross-modal reconstruction 설명시, 이미지 일부를 다른 Modality의 이미지로 마스킹 한 것과 같이 이번에는 template 으로 마스킹한다. 그리고 원래 이미지를 복원하도록 학습한다.
아래와 같이 Masking Function만 다르고 Loss도 같은 형태이다.
Result
Model은 distilled modality template을 형성하기 위해
1) 특정 modality의 structural Info를 학습해야 할 뿐만 아니라
2) modalities 간의 transformation relationship도 학습해야 한다.
궁극적으로, Pre-trained model에 의해 학습된 representation은 modality-agnostic으로 간주되며, 다양한 modality의 fused representation을 포함한다.
Epoch에 따라 학습된 Modality별 Templates은 Fig 4에서 확인 가능하다.
Modality-aware Contrastive Learning
앞서 서술한 서로 다른 방식(Cross-modal or template)으로 마스킹된 같은 환자의 이미지 쌍이 Feature level에서 consistency를 갖도록 유도하는 contrastive learning 방식이다.
예로, 같은 환자의 T1 Modal의 값을 기준으로 만들었지만, 하나는 T2 Modal로 마스킹한 것, 다른 하나는 Template으로 마스킹한 것이 있다고 하자. 그렇다면 이 둘은 "같은 환자의 T1 Modal의 값을 나타낸다"는 공통적인 의미가 있기에, Encoder를 통해 추출한 Feature는 Feature Space상 가까워야한다. 즉 유사해야한다.라는 가정으로 학습하는 것이다.
단방향 Contrastive Loss는 아래와 같다.
f는 Cross-modal 방식의 Embedding, g는 template 방식의 Embedding을 나타낸다.
최종적으로 학습에 사용하는 Loss는 아래와 같이 양방향 Loss다.
3. Conclusion
이 페이퍼에서는 여러 MRI Modality 간의 구조적인 관계를 활용한 self-supervised 방식의 Pre-training method를 제시했다.
곧 일반적인 의미의 Multimodal Data를 다뤄야할 것 같은데, 아래와 같이 이번 논문에서의 아이디어들이 유사하게 적용될 수도 있을 것 같다.
1) 각 Modality를 활용해 특정 Modality를 Reconstruction하는 방식으로 Representation Learning
2) 각 Modality의 대표적인 요약 벡터(Template)를 생성
3) Contrastive Learning을 활용한 Modality간 Aligned embedding