FedMerge: 모델 병합을 통한 연합 학습 개인화

초록

FedMerge는 서버가 다수의 전역 모델을 클라이언트별 가중치로 병합해 하나의 맞춤형 모델을 생성하고, 클라이언트는 병합된 모델만을 로컬 학습한다. 이를 통해 다중 전역 모델의 장점을 유지하면서 통신·연산 비용을 크게 줄이고, 비IID 환경에서 기존 개인화 FL 방법보다 높은 성능을 달성한다.

상세 분석

FedMerge는 기존 연합 학습(Federated Learning, FL)에서 발생하는 비IID 데이터 문제를 모델 병합(model merging)이라는 지식 집합 기법으로 해결한다. 전통적인 FL은 하나의 전역 모델을 모든 클라이언트에 공유하지만, 데이터 분포가 크게 다를 경우 일반화 성능이 급격히 떨어진다. 최근 다중 전역 모델을 활용하는 MoE(Mixture‑of‑Experts) 기반 방법은 각 클라이언트가 여러 전문가 모델을 로컬에서 선택·조합해 개인화 성능을 높이지만, 모델 수가 늘어날수록 클라이언트의 통신·연산 부하가 선형적으로 증가한다는 근본적인 한계가 있다.

FedMerge는 이러한 한계를 “서버‑사이드 병합”이라는 설계로 극복한다. 서버는 d개의 전역 모델 Θ₁…Θ_d를 유지하고, 각 클라이언트 i에 대해 가중치 벡터 w(i,·)∈ℝ^d를 학습한다. 개인화 모델은 θ_i = Σ_j w(i,j)·Θ_j 로 정의되며, 여기서 w는 소프트맥스 정규화로 확률적 의미를 갖게 한다. 클라이언트는 병합된 θ_i만 받아 로컬 손실 ℓ(Y_i, f(X_i;θ_i))을 최소화하는 표준 SGD(또는 Adam) 업데이트를 수행한다. 로컬 업데이트 후 Δθ_i를 서버에 전송하면, 서버는 체인룰을 이용해 전역 모델과 병합 가중치의 그래디언트를 역전파한다. 구체적으로 ∂L/∂Θ_j = Σ_i (n_i/n)·w(i,j)·∂ℓ/∂θ_i, ∂L/∂w(i,j) = (n_i/n)·⟨Θ_j, ∂ℓ/∂θ_i⟩ 로 계산된다. 이는 전역 모델이 모든 클라이언트의 업데이트에 영향을 받고, 병합 가중치는 해당 전역 모델과 로컬 업데이트 간의 내적에 비례해 조정된다는 의미다.

핵심 기술적 기여는 다음과 같다. 첫째, 전역 모델과 병합 가중치를 공동 최적화하는 새로운 목적함수(식 3‑4)를 제시하고, 이를 “단일 레이어 완전 연결망”에 비유해 역전파식(5‑6)을 도출했다. 둘째, 클라이언트‑사이드 연산량을 병합 모델 하나로 제한함으로써 모델 수와 무관하게 일정한 통신·연산 비용을 보장한다. 셋째, 대규모 파라미터를 가진 모델(ResNet, LoRA 기반 파운데이션 모델)에서도 효율적인 가중치 업데이트를 위해 마지막 몇 개 레이어(분류 헤드)만을 사용해 내적을 계산하는 실용적 트릭을 도입했다. 넷째, 실험에서는 이미지 분류(ResNet), 자연어 처리(LoRA), 그리고 다양한 비IID 시나리오(디렉터리 분포, 라벨 스키우, 클러스터링 기반)에서 기존 클러스터링‑기반 PFL, pFed‑MoE, FedAvg 등과 비교해 평균 2‑5%p 이상의 정확도 향상을 기록했다.

또한 FedMerge는 로컬 학습 단계에서 기존 FedAvg와 동일한 SGD 흐름을 유지하므로, 기존 FL 파이프라인에 최소한의 코드 변경만으로 적용 가능하다. 수렴 분석(부록)에서는 가중치와 전역 모델이 각각 Lipschitz 연속성을 만족한다는 가정 하에, 전체 손실이 기대값 기준으로 감소함을 보이며, 실험적 수렴 속도도 FedAvg와 동등하거나 더 빠른 것을 확인했다.

전반적으로 FedMerge는 “다중 전역 모델을 활용하면서도 클라이언트 비용을 일정하게 유지”한다는 설계 목표를 성공적으로 달성했으며, 모델 병합이라는 비교적 단순한 연산을 통해 개인화 성능을 크게 끌어올린 점이 가장 큰 혁신으로 평가된다.