Conformalized Neural Networks for Federated Uncertainty Quantification under Dual Heterogeneity

본 논문은 데이터 및 모델의 이중 이질성 하에서 효율적인 통신과 동시에 개인별 및 전역적 예측 불확실성 보장을 달성하기 위해 단일 통신 라운드에서 에이전트와 서버 간 교정을 수행하는 새로운 연방 학습 프레임워크인 FedWQ-CP 를 제안합니다.

핵심

🏥 배경: 병원이 서로 다른 상황 (연방 학습의 문제)

imagine 하세요. 전 세계 여러 병원이 모여서 AI 를 함께 훈련시키고 있습니다.

1. 데이터의 차이 (Data Heterogeneity): A 병원은 어린이 환자만 많고, B 병원은 노인 환자만 많습니다. 데이터가 완전히 다릅니다.
2. 모델의 차이 (Model Heterogeneity): A 병원은 최신 고사양 컴퓨터로 훈련된 '슈퍼 AI'를 쓰고, B 병원은 구형 컴퓨터로 훈련된 '평범한 AI'를 씁니다.

이때 가장 큰 문제는 **"AI 가 자신의 답변을 얼마나 확신하는지 (불확실성)"**를 정확히 알기 어렵다는 것입니다.

• 슈퍼 AI는 "99% 확신"이라고 말하지만, 사실은 너무 자신만만해서 틀릴 수도 있습니다 (과신).
• 평범한 AI는 "50% 정도"라고 말하며 너무 소심해져서, 실제로는 맞을 수도 있는데 도움을 주지 못합니다 (소극).

기존 방법들은 이 두 가지 차이 (데이터와 모델) 를 따로따로만 해결하려다 보니, 전체 평균은 좋아 보여도 **약한 병원은 계속 잘못된 진단을 내리는 '침묵하는 실패'**가 일어났습니다.

---

💡 해결책: "FedWQ-CP" (공정한 점수 조정 시스템)

이 논문은 FedWQ-CP라는 새로운 방법을 제안합니다. 이 방법은 '공정하게 점수를 조정하는 중재자' 역할을 합니다.

1. 각 병원의 '자신만의 기준선' 만들기 (로컬 캘리브레이션)

중앙 서버가 모든 데이터를 모으는 건 불가능합니다 (개인정보 보호 때문). 대신 각 병원이 자신만의 데이터로 AI 가 얼마나 틀렸는지 (점수) 를 계산합니다.

• 슈퍼 AI 는 높은 점수 기준을, 평범한 AI 는 낮은 점수 기준을 스스로 정합니다.
• 이때 중요한 건, **각 병원이 얼마나 많은 데이터를 가지고 훈련했는지 (샘플 수)**도 함께 기록한다는 점입니다.

2. 서버로 보내기: "데이터만 보내지 마!"

각 병원은 원본 환자 데이터나 AI 의 복잡한 내부 구조를 보내지 않습니다. 오직 두 가지 숫자만 서버로 보냅니다.

1. 기준점 (Quantile Threshold): "내가 이 정도 점수 이상이면 확신한다"는 기준.
2. 데이터 크기 (Sample Size): "내가 이 기준을 세울 때 몇 명을 봤는지".

3. 서버의 마법: "가중 평균" (Weighted Aggregation)

서버는 받은 숫자들을 단순히 평균내는 게 아니라, **데이터를 많이 본 병원의 기준에 더 큰 무게 (가중치)**를 둡니다.

• 비유: 만약 1000 명을 본 '슈퍼 병원'이 "이 정도면 안전하다"고 하고, 10 명만 본 '작은 병원'이 "완전 안전하다"고 한다면, 서버는 슈퍼 병원의 의견을 더 신뢰해서 최종 기준을 정합니다.
• 이렇게 해서 약한 병원이 너무 소극적으로 나오지 않도록 잡아주고, 강한 병원이 너무 과신하지 않도록 조절해 줍니다.

---

🌟 왜 이 방법이 특별한가요?

1. 한 번에 끝내요 (One-shot): 기존 방법들은 서버와 병원이 수백 번 대화를 나누며 조정해야 했지만, 이 방법은 서로 한 번만 숫자를 주고받으면 끝납니다. 매우 빠르고 효율적입니다.
2. 약한 병원을 구해요: 기존 방법들은 전체 평균만 맞추려다 약한 병원이 계속 실패하는 경우가 많았습니다. 하지만 이 방법은 데이터가 적은 병원의 기준도 고려해서, 어느 병원이든 신뢰할 수 있는 진단을 내리게 해줍니다.
3. 정확하면서도 간결해요: 불필요한 "안전 장치"를 너무 많이 두지 않아서, 진단 결과를 더 간결하고 명확하게 줍니다. (예: "암일 확률이 95% 이상"이라고 명확히 알려주는 것)

📝 결론

이 논문은 **"서로 다른 능력과 데이터를 가진 AI 들이 함께 일할 때, 누구도 소외되지 않고 모두 신뢰할 수 있도록 만드는 공정한 규칙"**을 만들었습니다.

마치 다양한 실력을 가진 학생들 (병원) 이 모여 시험을 볼 때, 선생님이 (서버) 각자의 노력 정도 (데이터 양) 를 고려해서 공정한 합격 기준을 정해주는 것과 같습니다. 덕분에 모든 학생이 자신의 실력에 맞는 정확한 평가를 받게 됩니다.

기술 요약

1. 문제 정의 (Problem Statement)

연방 학습 (Federated Learning, FL) 은 데이터 프라이버시를 보호하면서 분산된 환경에서 모델을 학습시키는 중요한 패러다임이지만, 불확실성 정량화 (Uncertainty Quantification, UQ) 측면에서 심각한 도전에 직면해 있습니다.

• 이중 이질성 (Dual Heterogeneity) 의 문제: 기존 FL-UQ 방법론들은 주로 데이터의 이질성 (데이터 분포의 차이) 또는 모델의 이질성 (아키텍처와 학습 강도의 차이) 중 하나만을 고려하거나, 이를 분리하여 다룹니다. 그러나 실제 환경에서는 두 가지가 동시에 발생합니다.
  • 데이터 이질성: 각 에이전트 (예: 병원) 의 데이터 분포가 다르고 (레이블 불균형, 공변량 편향), 데이터 양이 천차만별입니다.
  • 모델 이질성: 하드웨어 제약으로 인해 각 에이전트는 서로 다른 아키텍처 (예: 얕은 CNN vs 깊은 ResNet) 와 학습 강도로 모델을 훈련시킵니다.
• 현재의 한계: 이러한 이중 이질성 하에서 기존의 연방 컨포멀 예측 (Federated Conformal Prediction) 방법들은 전역 평균 (Global Average) 은 만족하더라도, 개별 에이전트 수준 (Agent-wise) 에서 과도한 커버리지 (Over-coverage) 또는 **심각한 커버리지 부족 (Under-coverage)**을 보입니다. 특히 데이터가 적거나 모델 성능이 낮은 약한 에이전트 (Weak Agents) 는 시스템이 신뢰할 수 없는 예측을 내더라도 이를 감지하지 못하는 '침묵하는 실패 (Silent Failures)'가 발생할 위험이 큽니다.

2. 제안 방법: FedWQ-CP (Methodology)

저자들은 **FedWQ-CP (Federated Weighted Quantile Conformal Prediction)**라는 새로운 프레임워크를 제안합니다. 이는 데이터와 모델의 이질성을 동시에 처리하면서도 통신 효율성을 극대화하는 원샷 (One-shot) 연방 보정 방법입니다.

핵심 알고리즘 및 프로세스

1. 로컬 보정 (Local Calibration):

   • 각 에이전트는 공유된 훈련 데이터를 기반으로 자신의 로컬 예측기를 학습하고 고정합니다.
   • 각 에이전트는 자신의 로컬 보정 데이터 (Calibration Data) 를 사용하여 비동일성 점수 (Non-conformity Scores) 를 계산합니다.
   • 분류 작업에는 APS(Adaptive Prediction Set), 회귀 작업에는 CQR(Conformalized Quantile Regression) 을 사용하여 점수를 산출합니다.
   • 각 에이전트는 로컬 데이터 분포에 기반한 **로컬 컨포멀 양분수 임계값 (Local Quantile Threshold, $q_k$)**을 계산합니다. 이는 각 모델의 고유한 불확실성 스케일을 정규화하는 역할을 합니다.

2. 가중치 양분수 집계 (Weighted Quantile Aggregation):

   • 각 에이전트는 서버에 **로컬 임계값 ($q_k$)**과 **보정 샘플 크기 ($n_k$)**만 전송합니다. (원시 데이터나 모델 파라미터는 공유되지 않음).
   • 서버는 전송된 정보를 바탕으로 샘플 크기에 비례한 가중 평균을 통해 전역 임계값 ($\hat{q}$) 을 계산합니다.
   • 공식: $\hat{q} = \sum_{k=1}^{M} \frac{n_k}{N} q_k$ (여기서 $N$은 전체 샘플 수).
   • 이 가중치 방식은 데이터가 적어 통계적 노이즈가 큰 약한 에이전트의 임계값이 전역 결정을 왜곡하는 것을 방지하고, 데이터가 풍부한 강한 에이전트의 신뢰도를 반영합니다.

3. 전역 평가:

   • 서버는 계산된 전역 임계값 $\hat{q}$를 모든 에이전트에 브로드캐스트합니다.
   • 각 에이전트는 이 임계값을 사용하여 공유된 전역 테스트 세트에서 예측 집합 (Prediction Sets) 또는 구간을 생성합니다.

3. 주요 기여 (Key Contributions)

1. 이중 이질성 해결: 데이터 분포의 차이와 모델 아키텍처/성능의 차이를 동시에 고려하는 최초의 연방 UQ 프레임워크 중 하나입니다.
2. 통신 효율성 및 프라이버시:
   • 원샷 (One-shot) 통신: 보정을 위해 에이전트와 서버 간에 단 한 번의 통신만 필요합니다.
   • 최소 정보 공유: 각 에이전트는 2 개의 스칼라 값 (임계값과 샘플 수) 만 전송하므로 통신 오버헤드가 극도로 낮고 프라이버시가 완벽하게 보호됩니다.
3. 구조적 가정 불필요: 기존 방법들이 필요로 하는 분포 이동 (Distribution Shift) 모델링 (예: 밀도 비율 추정, 레이블 시프트 파라미터 추정) 이나 반복적인 최적화 (Iterative Optimization) 과정이 전혀 필요 없습니다.
4. 범용성: 분류 (Classification) 와 회귀 (Regression) 작업 모두에 적용 가능하며, 다양한 입력 모달리티와 아키텍처를 지원합니다.

4. 실험 결과 (Results)

저자들은 7 개의 공개 데이터셋 (MNIST, FashionMNIST, CIFAR-10 및 4 가지 의료 영상 데이터셋) 에서 분류 및 회귀 작업을 수행하여 FedWQ-CP 를 검증했습니다.

• 커버리지 정확도 (Coverage):
  • FedWQ-CP 는 **강한 에이전트 (Strong Agents)**와 약한 에이전트 (Weak Agents) 모두에서 목표 커버리지 (예: 95%) 를 일관되게 유지했습니다.
  • 반면, 기존 방법들 (DP-FedCP 등) 은 약한 에이전트에서 심각한 커버리지 부족 (Under-coverage) 을 보였거나, 전역 평균은 맞지만 개별 에이전트 신뢰도가 떨어지는 문제를 겪었습니다.
• 효율성 (Efficiency):
  • 커버리지 요구 사항을 만족하면서 가장 작은 예측 집합 크기 (Classification) 또는 **가장 짧은 구간 길이 (Regression)**를 생성했습니다.
  • 이는 불확실성을 불필요하게 과대평가하지 않고, 가장 정밀한 예측을 제공함을 의미합니다.
• ** Ablation Study:**
  • 샘플 크기 가중치를 제거한 경우 (단순 평균), 약한 에이전트에서 체계적인 커버리지 부족이 발생하여 가중치 집계 방식의 중요성을 입증했습니다.

5. 의의 및 결론 (Significance)

이 논문은 고위험 (High-stakes) 환경 (예: 의료 진단) 에서 연방 학습을 적용할 때 필수적인 신뢰할 수 있는 불확실성 정량화 문제를 해결합니다.

• 실용적 가치: 데이터와 모델이 매우 이질적인 현실 세계의 연방 학습 환경에서도 각 참여 기관이 자신의 예측에 대한 신뢰도를 정확히 파악할 수 있게 합니다.
• 확장성: 단일 통신 라운드와 최소한의 데이터 공유로 구현 가능하므로, 대용량 분산 시스템에서도 확장 가능하고 프라이버시 규제를 준수하는 솔루션을 제공합니다.
• 기술적 기여: 컨포멀 예측 이론을 이질적인 연방 환경에 성공적으로 적용하기 위해 '가중 양분수 집계'라는 새로운 메커니즘을 도입하여, 기존 방법론들의 한계를 극복했습니다.

요약하자면, FedWQ-CP는 이중 이질성 하에서도 개별 에이전트와 전역 시스템 모두에서 신뢰할 수 있고 효율적인 불확실성 정량을 가능하게 하는 획기적인 프레임워크입니다.