기술 요약: ProMoE-FL

문제 정의

본 논문은 의료 환경에서의 결측 모달리티가 존재하는 멀티모달 연합 학습(Federated Learning, FL) 문제를 다룹니다. 실제 임상 환경에서는 개인정보 보호 규제와 장비 가용성의 차이로 인해 데이터가 여러 기관에 파편화되어 있습니다. 결과적으로, 참여하는 클라이언트들은 불완전한 모달리티 세트(예: 어떤 곳은 X-ray만 있고, 어떤 곳은 텍스트 보고서만 있으며, 어떤 곳은 둘 다 보유함)를 갖게 됩니다.

기존의 접근 방식들은 다음과 같은 상당한 한계에 직면해 있습니다:

• 중앙 집중식 방법론 (예: 프롬프트 기반, 생성형, 또는 드롭아웃 기반 방식): 엄격한 개인정보 보호 제약과 통신 효율성 요구 사항 때문에 연합 학습(FL)에 적용하기 부적합합니다.
• 공공 데이터 의존성: CAR-MFL과 같은 방식은 결측된 클라이언트 데이터를 보완하기 위해 추가적인 공공 멀티모달 데이터셋을 필요로 하는데, 이는 개인정보 보호 문제와 도메인 격차로 인해 의료 분야에서는 현실적이지 않습니다.
• 단순 임퓨테이션(Naive Imputation): Zero-filling과 같은 방식은 결합된 특징 공간(joint feature space)을 왜곡하고 모델이 자주 관찰되는 모달리티에 편향되게 만듭니다.
• 특징 합성(Feature Synthesis)의 한계: FeatImp와 같은 방식은 보조적인 정보를 모델링하지 않고 오직 관찰된 모달리티로부터만 특징을 합성하며, PmcmFL과 같은 프로토타입 기반 방식은 클래스 내 변동성(intra-class variability)을 포착하지 못하고 클래스 내 모든 샘플에 동일한 표현을 할당하여 의료 데이터에 필수적인 정밀함을 놓칩니다.

방법론: ProMoE-FL

저자들은 외부 공공 데이터에 의존하지 않고도 견고한 결측 모달리티 특징 합성을 수행하도록 설계된 프레임워크인 ProMoE-FL(Prototype-conditioned Mixture-of-Experts for Federated Learning)을 제안합니다.

1. 클라이언트 인지형 글로벌 프로토타입 뱅크 (Client-Aware Global Prototype Bank)

클라이언트 간의 이질성과 분포 변화를 해결하기 위해, 본 프레임워크는 기관 간의 임상적으로 유의미한 모달리티 사전 지식(prior)을 포착하는 글로벌 프로토타입 뱅크를 구축합니다.

• 프로토타입 구축: 각 클라이언트는 각 클래스에 대해 학습 가능한 모달리티별 프로토타입($P^k_m$)을 유지합니다. 이는 투영된 특징 표현을 클래스 중심점(centroid)에 정렬시키는 프로토타입 정렬 손실($L_{proto}$)을 통해 로컬에서 업데이트됩니다.
• 글로벌 집계: 단순 평균과 달리, 서버는 클라이언트 프로토타입을 축적하여 글로벌 뱅크($P^{global}_m$)를 형성합니다. 이 뱅크는 잠재 공간(latent space)에서 "클라이언트 인지형 사전 지식" 역할을 수행하여, 클라이언트 간의 분포적 가변성을 고려한 특징 합성을 가이드합니다.

2. 프로토타입 조건부 전문가 혼합 (Prototype-Conditioned Mixture of Experts, MoE)

모달리티가 누락된 클라이언트를 위해, **프로토타입 조건부 디코더(Prototype-Conditioned Decoder, PCD)**가 누락된 특징을 합성합니다.

• 구조: PCD는 트랜스포머(Transformer) 디코더로 구현됩니다. 이는 가용한 모달리티의 잠재 특징을 쿼리(query)로 사용하고, 타겟(누락된) 모달리티의 글로벌 프로토타입 뱅크에 어텐션(attention)을 수행하여 필요한 컨텍스트를 추출합니다.
• 전문가 혼합(Mixture of Experts): 이질적인 모달리티 쌍 사이의 복잡하고 비선형적인 관계를 처리하기 위해, 시스템은 $\Phi_{MoE}$라는 전문가 혼합 구조를 채택합니다. 이는 여러 개의 독립적인 PCD($\phi_e$)로 구성됩니다.
• 방향 인지형 라우팅(Direction-Aware Routing): 모달리티 인지형 라우터($R$)는 인스턴스 컨텍스트와 모달리티 인덱스(소스 $i$, 타겟 $j$)를 조건으로 게이팅 분포를 계산합니다. 이는 데이터를 최적의 전문가 조합으로 동적으로 라우팅하여, 특정 결측 모달리티 시나리오에 대해 네트워크가 최적의 합성 전략을 적응적으로 선택할 수 있게 합니다.

3. 연합 학습 프로토콜 (Federated Training Protocol)

• 로컬 학습: 클라이언트는 태스크 손실($L_{task}$)과 프로토타입 정렬 손실($L_{proto}$)의 결합된 목적 함수를 최적화합니다.
• 대칭적 재구성(Symmetric Reconstruction): 멀티모달 클라이언트는 대칭적 재구성 목적 함수를 사용하여 합성 네트워크를 학습하며, 이를 통해 가용한 특징으로부터 누락된 특징을 재구성하고 그 반대 방향으로도 학습합니다.
• 통신: 클라이언트는 업데이트된 모델 파라미터, 프로젝션 헤드, 그리고 로컬 프로토타입을 전송합니다. 서버는 이를 FedAvg를 통해 집계합니다. 통신 오버헤드는 전체 파라미터의 약 0.3%로 매우 미미합니다.

주요 기여

1. 프레임워크 제안: 외부 공공 데이터 없이도 연합 학습 환경에서 결측 모달리티 합성을 가능하게 하는 프로토타입 조건부 특징 생성 프레임워크인 ProMoE-FL을 도입했습니다.
2. 새로운 아키텍처: 클라이언트 간 이질성을 처리하기 위한 클라이언트 인지형 프로토타입 어텐션과, 공유된 파라미터 내에서 조합 폭발(combinatorial parameter growth)을 방지하면서 방향 인지형 합성을 수행하는 전문가 혼합(MoE) 아키텍처를 통합했습니다.
3. 견고성 검증: 실제 임상 데이터 사일로(silo)를 반영하는 이질적(non-IID) 환경에서 개선된 견고성을 입증하는 광범위한 실험을 수행하였으며, 정량적 및 정성적 분석을 통해 이를 뒷받침했습니다.

실험 결과

본 방법론은 네 가지 공개 흉부 X-ray 데이터셋(MIMIC-CXR, NIH Open-I, PadChest, CheXpert)에서 평가되었습니다.

• 성능: ProMoE-FL은 동질적(homogeneous) 및 이질적(heterogeneous) 설정 모두에서 최신 기술(state-of-the-art) 베이스라인들(Zero-filling, Uniform-filling, FeatImp, PmcmFL, CAR-MFL)을 일관되게 능가했습니다.
  • 동질적 설정에서 9개 구성 중 6개에서 최고의 성능을 달al했습니다.
  • 이질적 설정(현실적인 연합 학습 시나리오를 반영)에서 모든 구성에 걸쳐 모든 베이스라인을 일관되게 앞질렀습니다. 특히, 단일 클래스별 프로토타입에 의존하여 성능 저하를 겪은 PmcmFL에 비해 유의미한 이득을 보였습니다.
• 공공 데이터 비교: 공공 데이터를 활용하는 CAR-MFL과 비교했을 때도 ProMoE-FL은 경쟁력 있는 성능을 보여주었으며, 이는 내부 프로토타입 메커니즘이 외부 데이터의 부재를 효과적으로 보완하고 있음을 증명합니다.
• 정성적 분석:
  • t-SNE/UMAP: ProMoE-FL은 FeatImp에 비해 합성된 특징과 실제(ground-truth) 특징 간의 더 긴밀한 통합과 더 정밀한 클래스 중심점 정렬을 보여주었습니다.
  • 희귀 병리 현상(Rare Pathologies): 본 방법론은 희귀 병리(예: 심장 확대, 골절, 폐 병변)에 대해 명확한 클래스 중심점 분리를 유지한 반면, FeatImp는 서로 다른 병리가 하나로 수렴하는 모드 붕괴(mode collapse) 현상을 보였습니다.
  • 전문가 활용도: 게이팅 분포가 0.5 근처에서 균형 잡힌 전문가 활용도를 보였으며, 이는 단일 전문가로의 붕괴가 아닌 안정적인 다중 전문가 협업이 이루어지고 있음을 나타냅니다.
• 언어적 이질성: 본 연구는 영어 데이터셋과 함께 스페인어 방사선 보고서(PadChest)를 포함함으로써, 이 특정 연합 학습 맥락에서 언어적 도메인 이질성을 고려한 첫 사례가 되었습니다.

의의 및 주장

본 논문은 ProMoE-FL이 실제 임상 연합 학습의 데이터 이질성 문제를 해결하는 실용적인 접근 방식임을 주장합니다. 단순한 특징 임퓨테이션을 넘어 적응형 프로토타입 조건부 합성으로 나아감으로써, 이 프레임워크는 다음과 같은 성과를 거둡니다:

• 진단 구조(희귀 병리 포함)를 보존하여, 합성된 데이터가 신뢰할 수 있는 진단에 필요한 변별 신호를 유지하도록 합니다.
• 가용 가능한 공공 멀티모달 데이터셋이 있다는 비현실적인 가정 없이도 효과적으로 작동합니다.
• 비독립적 동일 분포(non-IID) 데이터와 의료 센터마다 상이한 모달리티 가용성 문제를 성공적으로 해결하여, 개인정보를 보호하는 멀티모달 의료 AI를 위한 견고한 솔루션을 제공합니다.