Probabilistic Feature Imputation and Uncertainty-Aware Multimodal Federated Aggregation

1. 연구 배경 및 문제 정의 (Problem)

• 배경: 현대 의료 진단은 단일 정보원 (예: 영상) 보다 영상, 임상 기록, 텍스트 보고서 등 다양한 모달리티 (Multimodal) 를 결합했을 때 정확도가 높아집니다. 또한, 환자 데이터의 프라이버시 보호를 위해 연방 학습 (Federated Learning, FL) 이 필수적입니다.
• 핵심 문제 (모달리티 이질성): 실제 의료 환경에서는 기관마다 보유 가능한 데이터가 다릅니다. 대형 병원에는 영상과 텍스트 보고서가 모두 있지만, 소규모 클리닉이나 지역 병원은 영상만 보유하고 텍스트 데이터가 결여된 경우가 많습니다.
• 기존 방법의 한계: 결측 모달리티를 해결하기 위해 기존 연구들은 '특징 Imputation(대체)' 네트워크를 사용하여 누락된 데이터를 생성합니다. 그러나 기존 방법들은 결정론적 (Deterministic) 접근을 취하여, 불확실성이 높은 상황에서도 하나의 고정된 점 추정치 (Point Estimate) 만을 출력합니다. 이는 의료와 같은 안전이 중요한 분야에서 잘못된 진단을 내릴 위험 (Hallucination) 을 내포하며, 생성된 데이터의 신뢰도를 판단할 수 없다는 치명적인 결함이 있습니다.

2. 제안 방법론 (Methodology)

저자들은 확률론적 특징 Imputation 네트워크 (P-FIN) 와 불확실성 인식 연방 집계 (Fed-UQ-Avg) 를 제안합니다.

A. 확률론적 특징 Imputation 네트워크 (P-FIN)

• 개념: 누락된 모달리티 (예: 텍스트) 를 고정된 값이 아닌, 평균 ($\mu$) 과 분산 ($\sigma^2$) 으로 파라미터화된 가우시안 분포로 모델링합니다. 이를 통해 생성된 특징의 불확실성을 정량화합니다.
• 아키텍처:
  • 인코더: 영상은 ResNet-50, 텍스트는 BERT 를 사용하여 256 차원 특징 벡터로 변환합니다.
  • Imputation 모듈: 입력된 영상 특징을 기반으로 Transformer 인코더를 통해 누락된 텍스트 특징의 분포를 예측합니다.
  • 이중 출력 헤드: 평균 ($\mu$) 과 분산 ($\sigma^2$) 을 동시에 예측하는 두 개의 MLP 를 사용합니다.
• 학습 손실 함수 ($\beta$-NLL): 표준 NLL 손실은 모델이 손실을 줄이기 위해 무한히 큰 분산을 예측하는 '분산 폭발 (Variance Explosion)' 현상을 유발할 수 있습니다. 이를 방지하고 보정된 불확실성을 학습시키기 위해 $\beta$-NLL (Seitzer et al., 2022) 손실 함수를 사용하며, stop-gradient 연산을 통해 분산이 무의미하게 커지는 것을 막습니다.

B. 불확실성 인식 로컬 퓨전 (Local Uncertainty-Aware Fusion)

• 불확실성 게이팅 (Gating): Imputation 된 특징을 분류기에 바로 넣지 않고, 예측된 분산 ($\sigma^2$) 을 기반으로 시그모이드 게이트 ($g = \text{sigmoid}(-\log \sigma^2)$) 를 적용합니다.
  • 불확실성이 높으면 게이트 값이 0 에 가까워져 신뢰할 수 없는 특징 성분을 억제합니다.
  • 불확실성이 낮으면 게이트 값이 1 에 가까워져 특징을 통과시킵니다.
• 교차 모달 어텐션: 게이트가 적용된 텍스트 특징과 영상 특징을 상호 어텐션 (Cross-Modal Attention) 을 통해 융합하여 최종 분류에 활용합니다.

C. 불확실성 인식 연방 집계 (Fed-UQ-Avg)

• 문제: 기존 FedAvg 는 데이터 양에 비례하여 가중치를 부여하지만, Imputation 품질이 낮은 클라이언트 (높은 불확실성) 의 업데이트가 전체 모델을 오염시킬 수 있습니다.
• 해결책: 클라이언트의 가중치를 두 가지 요소로 결정합니다.
  1. 데이터 가중치 ($w_{data}$): 기존 FedAvg 방식 (데이터 양 비례).
  2. 신뢰도 가중치 ($w_{conf}$): 클라이언트의 평균 Imputation 불확실성 ($\bar{\sigma}^2$) 을 기반으로 계산. 불확실성이 낮을수록 가중치가 높아집니다.
• 최종 가중치: $W_k = (1-\alpha)w_{data} + \alpha w_{conf}$ 형태로 결합하여, Imputation 신뢰도가 높은 클라이언트의 업데이트를 우선적으로 반영합니다.

3. 주요 기여 (Key Contributions)

1. P-FIN 아키텍처 제안: 누락된 모달리티를 생성할 때 불확실성 (분산) 을 함께 출력하는 최초의 확률론적 연방 학습 프레임워크를 제시했습니다.
2. 이중 레벨 불확실성 활용:
   • 로컬 레벨: 게이팅 메커니즘을 통해 신뢰할 수 없는 Imputation 특징을 분류 전 단계에서 차단.
   • 글로벌 레벨: Fed-UQ-Avg 를 통해 Imputation 품질이 낮은 클라이언트의 업데이트 가중치를 낮춰 전역 모델의 성능을 보호.
3. $\beta$-NLL 손실 적용: 분산 붕괴 (Variance Collapse) 나 폭발을 방지하고 보정된 (Calibrated) 불확실성 추정을 가능하게 함.

4. 실험 결과 (Results)

• 데이터셋: CheXpert, NIH Open-I, PadChest 의 흉부 X-ray 데이터를 활용하여 연방 학습 환경을 시뮬레이션했습니다.
• 실험 설정: 단모달 (이미지만 보유) 클라이언트와 다모달 (이미지 + 텍스트 보유) 클라이언트의 비율을 8:2, 6:4, 4:6 으로 변경하며 데이터 부족 상황을 테스트했습니다.
• 성능:
  • 가장 극단적인 조건 (단모달 80% 인 8:2 비율) 에서 제안된 방법 (P-FIN + Fed-UQ-Avg) 은 기존 결정론적 방법 (FIN + FedAvg) 대비 AUC 5.36%p 향상 효과를 보였습니다.
  • 전체적으로 모든 설정에서 Zero-filling, Uniform-filling, 기존 Imputation 방법 등을 압도적으로 상회했습니다.
• 보정 분석 (Calibration):
  • ECE (Expected Calibration Error): 0.0422 로 매우 낮은 값을 기록하여 예측된 불확실성이 실제 오차와 잘 일치함을 증명했습니다.
  • 상관관계: 예측된 불확실성이 높을수록 실제 Imputation 오차도 증가하는 강한 양의 상관관계를 확인했습니다.

5. 의의 및 결론 (Significance & Conclusion)

• 안전성 강화: 의료 AI 에서 '학습된 불신 (Learning to Distrust)' 패러다임을 도입하여, 모델이 자신이 모르는 상황 (불확실성이 높은 Imputation) 을 스스로 인지하고 이를 분류나 집계 과정에서 배제함으로써 오진 위험을 줄였습니다.
• 규제 준수: FDA 등 규제 기관이 요구하는 AI 의료기기의 불확실성 정량화 (Uncertainty Quantification) 요구사항을 충족시키는 실용적인 솔루션을 제공합니다.
• 미래 전망: 단방향 (영상→텍스트) Imputation 에서 양방향 생성 및 인지적 불확실성 (Epistemic Uncertainty) 통합 등으로 확장할 수 있는 기반을 마련했습니다.

이 논문은 결측 데이터가 존재하는 연방 학습 환경에서, 불확실성을 정량화하고 이를 시스템의 의사결정 (로컬 퓨전 및 글로벌 집계) 에 적극적으로 활용함으로써 의료 AI 의 신뢰성과 안전성을 획기적으로 높였다는 점에서 의의가 큽니다.