[논문 리뷰] Boosting Federated Learning Convergence with Prototype Regularization

Yu Qiao, Huy Q. Le외 1명|arXiv (Cornell University)|2023. 07. 20.

Privacy-Preserving Technologies in Data인용 수 5

한 줄 요약

이 논문은 프로토타입 기반 정규화를 도입하여 비독립identically distributed (non-IID) 환경에서 수렴 속도를 가속화하고 정확도를 향상시키는 FedPR이라는 페더레이티드 러닝 프레임워크를 제안한다. 클라이언트는 각 클래스의 프로토타입을 국지적으로 계산하며, 서버는 이를 집계하여 전역 프로토타입을 형성한다. 이 전역 프로토타입은 L2 거리 최소화를 통해 국지 학습을 정규화하는 데 사용되며, 이로 인해 MNIST와 Fashion-MNIST에서 각각 3.3%와 8.9% 높은 테스트 정확도를 달성한다.

ABSTRACT

As a distributed machine learning technique, federated learning (FL) requires clients to collaboratively train a shared model with an edge server without leaking their local data. However, the heterogeneous data distribution among clients often leads to a decrease in model performance. To tackle this issue, this paper introduces a prototype-based regularization strategy to address the heterogeneity in the data distribution. Specifically, the regularization process involves the server aggregating local prototypes from distributed clients to generate a global prototype, which is then sent back to the individual clients to guide their local training. The experimental results on MNIST and Fashion-MNIST show that our proposal achieves improvements of 3.3% and 8.9% in average test accuracy, respectively, compared to the most popular baseline FedAvg. Furthermore, our approach has a fast convergence rate in heterogeneous settings.

연구 동기 및 목표

• 클라이언트 간 데이터 분포가 비독립identically distributed (non-IID)일 경우 모델 성능 저하 문제를 해결한다.
• 데이터 프라이버시를 훼손하지 않으면서 이질적 데이터 환경에서 수렴 속도와 테스트 정확도를 향상시킨다.
• 전역 클래스 프로토타입을 활용해 국지 모델 학습을 이끌어, 클라이언트 간 일반화 능력을 향상시킨다.
• 프로토타입을 페더레이티드 학습 루프에 통합하는 저복잡도, 확장 가능한 정규화 메커니즘을 개발한다.

제안 방법

• 클라이언트는 각 클래스의 평균 임베딩을 사용하여 자신의 훈련 샘플에서 국지 프로토타입을 계산한다.
• 서버는 모든 클라이언트의 이러한 국지 프로토타입을 단순 평균을 통해 집계하여 클래스별 전역 프로토타입을 형성한다.
• 전역 프로토타입은 클라이언트로 전송되어 국지 학습을 정규화하는 데 사용되며, 이는 국지 특징 표현과 전역 프로토타입 간의 L2 거리 손실을 최소화하는 방식이다.
• 국지 목표 함수는 표준 교차 엔트로피 손실과 프로토타입 정규화 항을 조합한다: $\mathcal{L}_{i}(\omega_{i}) = \mathcal{L}_{i}(\mathcal{F}(\omega;\boldsymbol{x}_{i}),y_{i}) + \ell_{2}(f_{e}(\omega_{e};\boldsymbol{x}_{i}) - \overline{y}_{j})$.
• 모델 파라미터와 전역 프로토타입은 동기화된 라운드를 갖는 FedAvg 스타일의 통신 체계에서 반복적으로 업데이트된다.
• 프레임워크는 4층 컨볼루션 신경망을 사용하여 구현되었으며, 강한 non-IID 조건을 시뮬레이션하기 위해 Dirichlet 기반 데이터 스케일($\alpha = 0.05$) 하에서 테스트되었다.

실험 결과

연구 질문

• RQ1 프로토타입 기반 정규화는 non-IID 페더레이티드 러닝 환경에서 수렴 속도와 테스트 정확도 향상에 기여하는가?
• RQ2 전역 프로토타입을 국지 학습에 통합할 경우, 이질적 데이터를 가진 클라이언트 간 모델 일반화에 어떤 영향을 미치는가?
• RQ3 강한 데이터 스케일 조건 하에서 제안된 방법이 FedAvg에 비해 정확도와 수렴 속도 측면에서 뛰어나게 성능을 발휘하는가?
• RQ4 대규모 클라이언트 수에 걸쳐 프로토타입 집계를 효율적으로 계산하고 확장할 수 있는가? 통신 오버헤드가 크지 않은가?

주요 결과

• MNIST에서 FedPR은 $\alpha = 0.05$ 데이터 스케일 조건 하에서 평균 테스트 정확도 94.62%를 기록했으며, 이는 FedAvg의 91.57% 대비 3.3% 향상된 것이다.
• Fashion-MNIST에서 FedPR은 동일한 조건 하에 평균 테스트 정확도 86.05%를 달성했으며, 이는 FedAvg의 79.04% 대비 8.9% 향상된 것이다.
• 제안된 방법은 커뮤니케이션 라운드 수에 따라 테스트 정확도가 빨리 안정화되는 것으로 나타나, FedAvg에 비해 더 빠른 수렴 속도를 보였다.
• 프로토타입 정규화 메커니즘이 페더레이티드 환경에서 클래스 불균형과 데이터 이질성으로 인한 성능 저하를 효과적으로 완화했다.
• 이 프레임워크는 낮은 알고리즘 복잡도를 유지하여 실세계 페더레이티드 시스템에 실용적으로 구현 가능한 것으로 나타났다.
• 전역 프로토타입을 정규화 요소로 사용함으로써 국지 표현을 공통의 전역 집계 클래스 구조와 일치시켜 모델 일반화 능력을 향상시켰다.