[논문 리뷰] Multi-Level Branched Regularization for Federated Learning

Jinkyu Kim, Geeho Kim외 1명|arXiv (Cornell University)|2022. 07. 14.

Privacy-Preserving Technologies in Data인용 수 11

한 줄 요약

이 논문은 비독립identical(Non-IID) 환경에서 모델 드리프트를 줄이기 위해 로컬 및 글로벌 하위네트워크를 다수의 하이브리드 경로에 통합하는 다중 수준 분기 정규화 기법인 FedMLB을 제안한다. 온라인 지식 정련을 통해 로컬 메인 경로 표현을 하이브리드 브랜치 표현과 정렬함으로써, 추가 통신 또는 메모리 비용 없이 정확도와 수렴성을 향상시킨다. 이는 다양한 백본과 데이터 이질성 수준에서 최신 기법들을 능가한다.

ABSTRACT

A critical challenge of federated learning is data heterogeneity and imbalance across clients, which leads to inconsistency between local networks and unstable convergence of global models. To alleviate the limitations, we propose a novel architectural regularization technique that constructs multiple auxiliary branches in each local model by grafting local and global subnetworks at several different levels and that learns the representations of the main pathway in the local model congruent to the auxiliary hybrid pathways via online knowledge distillation. The proposed technique is effective to robustify the global model even in the non-iid setting and is applicable to various federated learning frameworks conveniently without incurring extra communication costs. We perform comprehensive empirical studies and demonstrate remarkable performance gains in terms of accuracy and efficiency compared to existing methods. The source code is available at our project page.

연구 동기 및 목표

• 데이터 이질성과 비독립identical 클라이언트 데이터 분포로 인한 피드포워드 학습에서의 모델 드리프트 문제를 해결한다.
• 통신 또는 메모리 오버헤드를 증가시키지 않고도 학습 중 로컬 모델과 글로벌 모델 간의 분리도를 줄인다.
• 글로벌 표현 지식을 유지하면서도 클라이언트별 적응을 가능하게 하는 정규화 기법을 개발한다.
• 다양한 피드포워드 학습 프레임워크와 백본 아키텍처에서 호환성과 효율성을 확보한다.
• 낮은 참여도와 높은 데이터 이질성 환경에서 수렴 속도와 최종 정확도를 향상시킨다.

제안 방법

• 다양한 아키텍처 수준에서 글로벌 모델의 하위네트워크를 로컬 모델에 통합하여 다수의 하이브리드 브랜치를 구성한다.
• 로컬 모델의 메인 경로를 학습하는 동안 각 하이브리드 브랜치 내의 글로벌 하위네트워크 파라미터를 고정한다.
• 모든 브랜치에 대해 교차 엔트로피 손실을 적용하여 분류 능력을 유지하고, 메인 경로와 하이브리드 브랜치 출력 간의 KL 발산을 통한 지식 정련을 실시한다.
• 지식 정련에서 온도 조정된 소프트맥스를 사용하여 출력 분포를 부드럽게 하고 정련 안정성을 향상시킨다.
• 역전파 동안은 오직 메인 경로 파라미터만 업데이트하고, 하이브리드 브랜치 파라미터는 고정된 학생-선생 모델로 유지한다.
• 통신 또는 집계 메커니즘을 수정하지 않고도 표준 피드포워드 학습 파이프라인(FedAvg 등)에 통합한다.

실험 결과

연구 질문

• RQ1 통신 또는 메모리 비용 증가 없이 다중 수준 아키텍처 정규화가 비독립identical 피드포워드 학습에서 모델 드리프트를 줄일 수 있는가?
• RQ2 메인 경로와 하이브리드 브랜치 간의 온라인 지식 정련이 기존 정련 또는 정규화 방식보다 수렴성과 정확도 향상에 얼마나 기여하는가?
• RQ3 FedMLB가 다양한 신경망 아키텍처와 데이터 분포 설정에서 얼마나 일반화되는가?
• RQ4 하이브리드 브랜치의 수와 네트워크 깊이에 따른 배치가 성능 및 효율성에 미치는 영향은 어떠한가?
• RQ5 낮은 클라이언트 참여율과 높은 데이터 이질성 조건에서도 FedMLB의 성능 유지 여부는 어떠한가?

주요 결과

• Dirichlet(0.3) 비독립identical 데이터 하에서 CIFAR-100에서 ResNet-18을 사용할 경우 FedMLB는 54.58%의 Top-1 정확도를 기록하며, FedAvg(47.83%), FedLS-NTD(49.29%), FedGKD(47.96%)를 능가한다.
• VGG-9에서 FedMLB는 동일 조건 하에 54.54%의 정확도를 기록하며, FedAvg(47.04%)와 FedGKD(48.62%)를 뚜렷이 앞서간다.
• 모든 테스트된 백본에서 일관되게 정확도 향상을 보였으며, 경량 모델인 MobileNet(48.34% 대비 FedLS-NTD의 47.66%)과 ShuffleNet(42.29% 대비 FedLS-NTD의 39.27%)에서도 유의미한 성능 향상을 입증했다.
• 하이브리드 브랜치 수를 늘릴수록 성능 향상이 이루어지며, 표 8에서는 브랜치 수 증가에 따라 정확도 향상 추세가 명확히 나타난다.
• 로컬 계산량을 높일 경우(예: 100개의 로컬 에포크)에도 FedMLB는 FedAVG와 FedLS-NTD가 반복 횟수 증가에 따라 일관성 없이 행동하는 것과는 달리 강력한 성능 유지를 보였다.
• 제거 실험 결과, 하이브리드 브랜치 내 교차 엔트로피 손실과 KL 발산 손실 모두 성능 기여에 기여하며, 특히 KL 손실이 더 핵심적인 역할을 한다(표 9 참조).