기술 요약: 효율적인 2차 최적화를 통한 불확실성 및 개인화를 포함한 연합 학습

문제 정의

연합 학습(Federated Learning, FL)은 분산된 클라이언트 간에 원시 데이터를 공유하지 않고 협력적으로 모델을 학습할 수 있게 해줍니다. 그러나 표준적인 FL 방식은 두 가지 주요 과제에 직면해 있습니다:

1. 데이터 이질성(Data Heterogeneity): 클라이언트 데이터가 비독립 동일 분포(non-IID)인 경우, 단일 글로벌 모델은 개별 클라이언트의 분포를 효과적으로 나타내지 못해 하위 최적의 성능을 초래할 수 있습니다.
2. 불확실성 정량화의 부재: 표준 FL 방법은 일반적으로 모델 파라미터의 점 추정치(point estimates)만을 학습합니다. 이는 데이터가 적은 환경에서 과적합에 취약하게 만들며, OOD(Out-of-Distribution) 탐지와 같이 예측 불확실성을 제공하는 것이 중요한 작업에서 한계를 보입니다.

베이지안 FL은 사후 분포(posterior distributions)를 학습하여 불확실성을 포착하고 계층적 프레임워크를 통해 개인화를 가능하게 함으로써 원칙적인 해결책을 제시하지만, 기존 방법론들은 상당한 병목 현상을 겪습니다. 현재의 베이지안 FL 접근 방식은 계산 비용이 많이 드는 루틴(예: MCMC 샘플링, 명시적인 헤시안(Hessian) 계산이 필요한 라플라스 근사, 또는 딥 앙상블)에 의존하며, 전체 사후 분포를 전송함으로써 높은 통신 비용을 발생시킵니다. 이러한 요소들은 자원이 제한된 클라이언트에게 확장성 문제를 야기합니다.

방법론: FedIvon

저자들은 효율적인 2차 변분 추론(second-order variational inference)을 수행하기 위해 **개선된 변분 온라인 뉴턴(Improved Variational Online Newton, IVON)**을 활용하는 새로운 베이지안 FL 알고리즘인 FedIvon을 제안합니다. 이 방법은 베이지안 추론의 이점과 Adam 같은 1차 최적화 도구와 유사한 계산 및 통신 효율성 사이의 균형을 맞춥니다.

핵심 구성 요소

1. 클라이언트 측 변분 추론 (IVON):

   • 각 클라이언트는 점 추정치를 학습하는 대신, 대각 공분산 행렬을 가진 가우시안 분포 $N(\theta | m_k, \text{diag}(\sigma_k^2))$를 사용하여 로컬 사후 분포 $q_k(\theta)$를 근사합니다.
   • 2차 효율성: 제곱 기울기에 따라 학습률을 조정하는 Adam과 달리, IVON은 명시적인 계산 없이 대각 헤시안($h$)을 암시적으로 추정하기 위해 재매개변수화(reparameterization)를 사용합니다. 분산은 $\sigma^2 = 1/(h + \delta)$로 유도됩니다.
   • 최적화: 알고리즘은 자연 기울기(natural gradients)를 사용하여 평균 벡터($m$)와 분산($\sigma^2$)을 모두 업데이트합니다. 이는 자연스럽게 불확실성 추정치를 제공하며, 손실 함수의 기하학적 구조를 포착하여 더 빠른 수렴을 가능하게 하는 동시에 Adam과 유사한 계산 비용을 유지합니다.

2. 서버 측 사후 분포 집계:

   • 서버는 클라이언트의 사후 분포를 집계하여 글로벌 사후 분포를 계산합니다. 가우시안 분포의 곱 또한 가우시안 분포이므로, 글로벌 사후 분포는 가우시안 $N(\theta | \mu, \Lambda^{-1})$로 근사됩니다.
   • 집계 전략: 글로벌 평균($\mu$)과 정밀도 행렬($\Lambda$)은 클라이언트 데이터 크기($N_k$)에 비례하는 가중치를 적용하여 로컬 평균과 정밀도의 가중 합으로 계산됩니다. 이 전략은 전체 모델 가중치나 복잡한 분포를 전송하는 대신 평균과 대각 분산(또는 헤시안)만을 전송하므로 계산 효율적입니다.

3. 계층적 베이지안 프레임워크를 통한 개인화:

   • FedIvon은 개인화된 FL을 자연스럽게 지원합니다. 서버의 글로벌 사후 분포는 클라이언트의 로컬 최적화를 위한 사전 분포($p_k(\theta)$) 역할을 합니다.
   • 클라이언트는 개인화된 손실 함수 $L_k(q) = \mathbb{E}_{q_k}[\bar{\ell}_k(\theta)] + \beta D_{KL}(q_k(\theta) \| p_k(\theta))$를 최소화합니다.
   • 하이퍼파라미터 $\beta$는 로컬 적응과 글로벌 지식 사이의 트레이드오프를 조절합니다. $\beta=0$이면 모델은 완전히 개인화되며, 값이 높을수록 더 많은 글로벌 정보를 통합합니다.

주요 기여

• 효율적인 2차 베이지안 FL: IVON을 사용하여 1차 방법(Adam 등)과 유사한 계산 비용으로 변분 추론을 수행하는 FedIvon을 도입했습니다.
• 향상된 성능 및 불확실성: FedIvon이 최신(SOTA) 베이지안 및 비베이지안 FL 방법들에 비해 우수한 예측 정확도와 더 나은 불확실성 추정치를 달fulness함을 입증했습니다.
• 네이티브 개인화: 서버의 사후 분포를 사전 분포로 활용하여 로컬 데이터 특성과 글로벌 공유 지식 사이의 균형을 효과적으로 맞춤으로써, 클라이언트 수준의 모델 개인화를 자연스럽게 가능하게 하는 프레임워크를 제공합니다.

실험 결과

저자들은 EMNIST, SVHN, CIFAR-10 데이터셋을 사용하여 200개의 클라이언트(각 클라이언트당 데이터 <100개)가 존재하는 non-IID 설정에서 FedIvon을 평가했습니다.

표준 FL 결과

• 정확도 및 교정(Calibration): FedI안은 모든 데이터셋에서 테스트 정확도, 기대 교정 오차(ECE), 음의 로그 가능도(NLL), Brier Score 측면에서 베이스라인(FedAvg, FedLaplace)보다 우수한 성능을 보였습니다.
• OOD 탐지: OOD 탐지 작업에서 FedIvon은 베이스라인 대비 경쟁력 있거나 더 우수한 AUROC 점수를 기록하여, 예측 불확실성을 효과적으로 정량화하는 능력을 검증했습니다.
• 수렴: 가중치 기울기 추정을 위한 확률적 샘플링으로 인해 초기 수렴은 다소 느렸으나, 헤시안 근사가 개선됨에 따라 훈련이 진행될수록 다른 방법들을 앞질렀습니다.

개인화된 FL 결과

클래스 분포 왜곡(CIFAR-10) 및 클래스 개념 드리프트(CIFAR-100) 상황에서 FedIvon을 pFedME, pFedBayes, pFedVEM과 비교했습니다.

• CIFAR-10: FedIvon은 pFedVEM과 대등한 클라이언트 정확도를 달성하면서도, 글로벌 서버 정확도에서 유의미한 향상을 보여 이질적인 데이터의 집계 능력이 뛰어남을 입증했습니다.
• CIFAR-100: FedIvon은 개인화된 클라이언트 정확도와 글로벌 서버 정확도 모두에서 모든 베이스라인보다 크게 향상된 성능을 보였으며, 이는 개념 드리프트 처리에 대한 견고함을 보여줍니다.

의의 및 주장

본 논문은 FedIvon이 베이지안 FL의 이론적 이점(불확실성 정량화, 개인화)과 연합 시스템의 실질적 제약(계산, 통신) 사이의 간극을 메운다고 주장합니다. 효율적인 2차 최적화 접근 방식을 사용함으로써, 기존 연구에서 발견되는 무거운 계산 오버헤드(MCMC 또는 명시적 헤시안 계산)를 피할 수 있습니다.

저자들은 FedIvon이 다음과 같은 최초의 방법이라고 단언합니다:

1. 1차 최적화 도구(Adam 등)와 대등한 계산 효율성 제공.
2. 기존 베이지안 FL 방법의 확장성 문제 없이 완전한 베이지안 이점(불확실성 및 개인화) 제공.
3. 데이터 이질성이 높고 클라이언트 데이터가 제한적인 시나리오에서 우수한 성능 발휘.

이 연구는 FedIvon을 실제 환경의 견고하고, 불확실성을 인지하며, 개인화된 머신러닝 모델을 배포하기 위한 확장 가능하고 실용적인 솔루션으로 포지셔닝합니다.