[논문 리뷰] Achieving Personalized Federated Learning with Sparse Local Models

Tiansheng Huang, Shiwei Liu외 4명|arXiv (Cornell University)|2022. 01. 27.

Privacy-Preserving Technologies in Data인용 수 13

한 줄 요약

이 논문은 통신 및 계산 비용을 낮추면서도 높은 정확도를 달성하기 위해 개인화된 마스크를 사용하는 희소 로컬 모델을 활용하는 개인화된 플랫폼 학습 프레임워크인 FedSpa를 제안한다. 훈련 전반에 걸쳐 고정된 희소성 수준을 유지함으로써(희소에서 희소로의 훈련), FedSpa는 효율적이고 개인화된 모델 업데이트를 가능하게 하며, 1/√T 속도로 국소 최소화점을 수렴한다.

ABSTRACT

Federated learning (FL) is vulnerable to heterogeneously distributed data, since a common global model in FL may not adapt to the heterogeneous data distribution of each user. To counter this issue, personalized FL (PFL) was proposed to produce dedicated local models for each individual user. However, PFL is far from its maturity, because existing PFL solutions either demonstrate unsatisfactory generalization towards different model architectures or cost enormous extra computation and memory. In this work, we propose federated learning with personalized sparse mask (FedSpa), a novel PFL scheme that employs personalized sparse masks to customize sparse local models on the edge. Instead of training an intact (or dense) PFL model, FedSpa only maintains a fixed number of active parameters throughout training (aka sparse-to-sparse training), which enables users' models to achieve personalization with cheap communication, computation, and memory cost. We theoretically show that the iterates obtained by FedSpa converge to the local minimizer of the formulated SPFL problem at rate of $\mathcal{O}(\frac{1}{\sqrt{T}})$. Comprehensive experiments demonstrate that FedSpa significantly saves communication and computation costs, while simultaneously achieves higher model accuracy and faster convergence speed against several state-of-the-art PFL methods.

연구 동기 및 목표

• 이질적인 클라이언트 간에 일반화되지 않는 비균일 데이터 분포(non-IID) 문제를 해결하여 단일 글로벌 모델이 전역적으로 성능을 내지 못하도록 하는 과제를 해결한다.
• 기존의 밀도 모델 또는 동적 희소성에 의존하는 개인화된 플랫폼 학습(PFL) 방법이 유발하는 높은 통신 및 계산 오버헤드를 극복한다.
• 훈련 중에 일정한 희소성 수준을 유지함으로써 엣지 디바이스에서의 메모리 및 계산 비용을 줄이는 PFL 솔루션을 개발한다.
• 개인화 및 모델 성능을 유지하면서도 희소 서브모델을 효율적으로 집계할 수 있도록 한다.
• 이론적으로 수렴 성능를 분석하고, 최신 기술 대비 정확도, 속도, 효율성 면에서 뛰어난 성능을 실증적으로 검증한다.

제안 방법

• 글로벌 모델에 개인화된 희소 마스크를 적용하여 클라이언트별 희소 서브모델을 생성하는 새로운 희소 개인화 플랫폼 학습(SPFL) 문제 정식화를 도입한다.
• 각 클라이언트가 훈련 전반에 걸쳐 일정한 희소성 수준을 유지하는 희소 모델을 유지하고 훈련하는 PFL 프레임워크인 FedSpa를 설계한다.
• 현지 기울기와 데이터를 기반으로 새로운 개인화된 마스크를 탐색하기 위해 동적 희소 훈련(DST) 메커니즘을 사용하여 모델 개인화를 보장한다.
• 마스크 업데이트 중에 기울기 기반 가중치 복구를 적용하여, 상위-αₜ 크기의 기울기 좌표를 선택해 잘린 가중치를 복원한다.
• 서버에서 희소 기울기 업데이트와 새로운 마스크를 집계하고, 표준 평균화를 통해 글로벌 모델을 업데이트한다.
• 모든 로컬 모델이 훈련 전반에 걸쳐 희소성을 유지하여 통신 및 계산 오버헤드를 최소화한다.

실험 결과

연구 질문

• RQ1 비균일 데이터 분포(NON-IID) 환경에서 통신 및 계산 비용을 낮추면서도 높은 개인화 성능을 달성할 수 있는가?
• RQ2 개인화된 마스크를 사용하는 희소에서 희소로의 훈련이 밀도 기반 PFL 기준선 대비 더 빠른 수렴과 더 높은 정확도를 제공하는가?
• RQ3 기울기 기반 가중치 복구가 마스크 재훈련 중에 모델 성능 향상에 얼마나 효과적인가?
• RQ4 FedSpa를 통해 훈련된 글로벌 모델은 데이터 이질성 하에서 표준 FedAvg와 경쟁 가능한 성능을 달성할 수 있는가?
• RQ5 마스크 초기화 전략이 이질적인 데이터 환경에서 FedSpa의 최종 성능에 미치는 영향은 무엇인가?

주요 결과

• FedSpa는 라운드당 전송되는 희소 기울기와 마스크만으로도 밀도 기반 PFL 기준선 대비 통신 오버헤드를 50% 감소시킨다.
• FedSpa는 비균일 데이터 분포 환경에서 FedAvg 대비 최대 21.9% 높은 테스트 정확도를 달성했고, Ditto 대비 4.4% 높은 정확도를 기록했다.
• FedSpa가 동일한 정확도 수준에 도달하기 위해 FedAvg 대비 76.2% 적은 통신 라운드를 필요로 하며, Ditto 대비 38.1% 적은 라운드를 요구한다.
• FedSpa의 마스크 탐색 과정(DST)은 라운드당 총 월 타임의 5%에서 20%에 불과하여 최소한의 런타임 오버헤드를 유발한다.
• FedSpa의 기울기 기반 가중치 복구(DST)는 무작위 복구 대비 더 빠른 수렴과 더 높은 최종 정확도를 달성한다.
• 비균일 데이터 분포 환경에서 FedSpa의 글로벌 모델은 개인화된 모델보다 성능이 뚜렷이 열 劣하므로, 데이터 왜곡 하에서 서브넷이 전체 모델보다 우수하다는 점을 확인한다.