[논문 리뷰] Multi-modal Learning with Missing Modality via Shared-Specific Feature Modelling

Hu Wang, Yuanhong Chen외 4명|arXiv (Cornell University)|2023. 07. 26.

Multimodal Machine Learning Applications인용 수 5

한 줄 요약

이 논문은 공유-특수 특징 모델링(ShaSpec)을 제안한다. 이는 결손 모달리티가 있는 다중모달 학습을 위한 단순하면서도 효과적인 방법으로, 분포 정렬 및 도메인 분류 손실을 통해 공유 특징과 특수 특징을 동시에 학습함으로써, 훈련 또는 추론 시 모달리티가 결손된 경우에도 분류 및 세그멘테이션 작업 전반에서 견고한 성능을 보여준다. ShaSpec는 최신 기술 수준(SOTA)을 초월하며, BraTS2018에서 강화 종양의 세그멘테이션 Dice 스코어를 3% 이상 향상시키고, 종양 코어는 5%, 전체 종양은 3% 향상시켰다.

ABSTRACT

The missing modality issue is critical but non-trivial to be solved by multi-modal models. Current methods aiming to handle the missing modality problem in multi-modal tasks, either deal with missing modalities only during evaluation or train separate models to handle specific missing modality settings. In addition, these models are designed for specific tasks, so for example, classification models are not easily adapted to segmentation tasks and vice versa. In this paper, we propose the Shared-Specific Feature Modelling (ShaSpec) method that is considerably simpler and more effective than competing approaches that address the issues above. ShaSpec is designed to take advantage of all available input modalities during training and evaluation by learning shared and specific features to better represent the input data. This is achieved from a strategy that relies on auxiliary tasks based on distribution alignment and domain classification, in addition to a residual feature fusion procedure. Also, the design simplicity of ShaSpec enables its easy adaptation to multiple tasks, such as classification and segmentation. Experiments are conducted on both medical image segmentation and computer vision classification, with results indicating that ShaSpec outperforms competing methods by a large margin. For instance, on BraTS2018, ShaSpec improves the SOTA by more than 3% for enhancing tumour, 5% for tumour core and 3% for whole tumour. The code repository address is https://github.com/billhhh/ShaSpec/.

연구 동기 및 목표

• 기존 모델의 실제 적용 가능성을 제한하는 다중모달 학습에서의 모달리티 결손 문제를 해결한다.
• 기존 방법의 한계를 극복한다. 즉, 훈련 시 모든 모달리티가 반드시 필요하거나, 특정 작업 전용(예: 분류 또는 세그멘테이션 전용)인 방법들이다.
• 분류 및 세그멘테이션 작업 전반에 걸쳐 일반화 가능한 통합된 단순한 아키텍처를 개발하여, 훈련 및 추론 시 모달리티가 결손된 경우에도 대응할 수 있도록 한다.
• 특정 시나리오에 맞춘 별도의 훈련 없이도, 하나의 모델이 다양한 모달리티 조합의 결손을 처리할 수 있도록 한다. 이는 각 결손 시나리오별 별도의 모델을 구축할 필요를 줄인다.

제안 방법

• ShaSpec는 이중 브랜치 인코더를 사용한다. 하나는 모든 모달리티에 공통되는 특징을 위한 공유 인코더이며, 다른 하나는 모달리티별 특징을 위한 개별 특수 인코더이다.
• 공유 및 특수 특징은 학습 가능한 선형 투영 레이어를 사용한 잔차 학습 메커니즘을 통해 융합된다.
• 보조 작업으로는 분포 정렬(MMD 손실을 통한)과 도메인 분류가 포함되어, 분리된 의미 있는 특징 학습을 유도한다.
• 전반적인 손실은 주 작업 손실(예: 교차 엔트로피 또는 MSE)과 가중치 α 및 β를 통해 조정되는 보조 손실을 조합하여 구성된다.
• 이 방법은 엔드 투 엔드로 훈련되며, 분류 및 세그멘테이션 작업에 대해 최소한의 아키텍처 변경으로 적용 가능하다.
• 추론 시, 재학습 없이도 가용한 공유 및 특수 특징에 의존하여 동적으로 모달리티가 결손된 경우를 처리할 수 있다.

실험 결과

연구 질문

• RQ1 각 결손 시나리오별 별도의 모델이 필요 없이, 훈련 및 추론 모두에서 모달리티가 결손된 경우를 효과적으로 다룰 수 있는 통합된 다중모달 학습 프레임워크는 가능한가?
• RQ2 결손된 모달리티 상황에서, 분리된 공유 및 특수 특징을 학습함으로써 후행 작업 성능이 향상되는가?
• RQ3 분포 정렬 및 도메인 분류와 같은 보조 작업을 기반으로 한 단순한 아키텍처가 세그멘테이션 및 분류와 같은 다양한 작업에서 최신 기술 수준(SOTA)을 달성할 수 있는가?
• RQ4 보조 손실의 가중치를 조정하는 하이퍼파rameter α 및 β는 모델 성능과 훈련 안정성에 어떤 영향을 미치는가?

주요 결과

• BraTS2018에서 ShaSpec는 전체 모달리티 조건 하에서 강화 종양 분야에서 SOTA를 3.0% 향상시키고, 종양 코어는 5.0%, 전체 종양은 3.0% 향상시켰다.
• 추론 시 유일하게 T1 MRI만 이용 가능한 상황에서도 ShaSpec는 전체 종양의 Dice 스코어를 74.82%로 기록하여 베이스라인보다 1.5% 높였다.
• OpenI 데이터셋의 이진 흉부 X-ray 분류 작업에서 ShaSpec는 AUC 0.89를 달성했으며, 전체 모달리티 기반 베이스라인(AUC 0.90)을 초월했고, 30% 모달리티가 결손된 모델보다도 뚜렷이 뛰어난 성능을 보였다.
• 제거 실험 결과, α = 0.1 및 β = 0.02 조합이 최적의 성능을 내며, 이는 작은 보조 손실 가중치가 주 작업에 간섭을 최소화함을 시사한다.
• t-SNE 시각화 결과, 공유 특징은 다양한 모달리티 간에 잘 군집되어 있으며, 특수 특징은 분리되어 있음을 확인하여, 분리 목적의 타당성을 검증하였다.
• ShaSpec는 오직 0.22M 파라미터만을 사용하고, GPU 메모리는 일관되게 1421MiB를 소비하며, OOM 오류로 자주 다운되는 SMIL보다도 효율성과 메모리 안정성 면에서 뛰어난 성능을 보였다.