[논문 리뷰] Efficient Wireless Federated Learning with Partial Model Aggregation

Zhixiong Chen, Wenqiang Yi외 2명|arXiv (Cornell University)|2022. 04. 20.

Privacy-Preserving Technologies in Data인용 수 5

한 줄 요약

이 논문은 무선 피어드 학습에서 통신 부하를 줄이고 데이터 이질성에 대한 강건성을 향상시키기 위해 하위(특징 추출) 레이어만 파rameter 서버에서 집계하고 상위(예측) 레이어는 로컬에서 유지하여 개인화를 수행하는 부분 모델 집계 피어드 학습(PMA-FL)을 제안한다. 라플라스 최적화를 통해 장치 스케줄링, 대역폭, 계산 시간 분할을 공동으로 최적화함으로써 PMA-FL은 통신 부하를 감소시키고, 비독립 동일 분포(non-IID) 데이터 하에서 모델 수렴과 정확도를 향상시킨다. MNIST와 CIFAR-10에서 각각 3.13% 및 11.8% 높은 정확도를 달성하였으며, 기준 대비 최대 29%의 에너지 절감 또는 20%의 시간 절감을 이룬다.

ABSTRACT

The data heterogeneity across devices and the limited communication resources, e.g., bandwidth and energy, are two of the main bottlenecks for wireless federated learning (FL). To tackle these challenges, we first devise a novel FL framework with partial model aggregation (PMA). This approach aggregates the lower layers of neural networks, responsible for feature extraction, at the parameter server while keeping the upper layers, responsible for complex pattern recognition, at devices for personalization. The proposed PMA-FL is able to address the data heterogeneity and reduce the transmitted information in wireless channels. Then, we derive a convergence bound of the framework under a non-convex loss function setting to reveal the role of unbalanced data size in the learning performance. On this basis, we maximize the scheduled data size to minimize the global loss function through jointly optimize the device scheduling, bandwidth allocation, computation and communication time division policies with the assistance of Lyapunov optimization. Our analysis reveals that the optimal time division is achieved when the communication and computation parts of PMA-FL have the same power. We also develop a bisection method to solve the optimal bandwidth allocation policy and use the set expansion algorithm to address the device scheduling policy. Compared with the benchmark schemes, the proposed PMA-FL improves 3.13\% and 11.8\% accuracy on two typical datasets with heterogeneous data distribution settings, i.e., MINIST and CIFAR-10, respectively. In addition, the proposed joint dynamic device scheduling and resource management approach achieve slightly higher accuracy than the considered benchmarks, but they provide a satisfactory energy and time reduction: 29\% energy or 20\% time reduction on the MNIST; and 25\% energy or 12.5\% time reduction on the CIFAR-10.

연구 동기 및 목표

• 피어드 학습에서 데이터 이질성과 제한된 무선 통신 자원 문제를 해결한다.
• 비독립 동일 분포 데이터와 제약된 대역폭 및 에너지 조건 하에서 모델 수렴과 정확도를 향상시킨다.
• 무선 피어드 학습 시스템에서 통신 오버헤드와 에너지 소비를 감소시킨다.
• 장치 스케줄링, 대역폭 할당, 계산 시간 분할을 위한 공동 최적화 프레임워크를 개발한다.
• 균형 잡힌 계산과 통신을 통해 더 낮은 에너지 및 시간 비용으로 더 나은 학습 성능을 달성한다.

제안 방법

• 하위 레이어만 파라미터 서버에서 집계하고 상위 레이어는 로컬에서 개인화를 위해 유지하는 새로운 PMA-FL 프레임워크를 제안한다.
• 비볼록 손실 함수 하에서 PMA-FL의 수렴 경계를 유도하여 불균형한 데이터 크기의 영향을 분석한다.
• 라플라스 최적화를 사용하여 통신 및 계산 에너지를 균형 잡고, 스케줄링된 데이터 크기를 최대화하기 위한 공동 최적화 문제를 수립한다.
• 자원 제약 하에서 최적의 대역폭 할당 정책을 해결하기 위해 이분법적 방법을 설계한다.
• 최적의 장치 스케줄링 정책을 결정하기 위해 집합 확장 알고리즘을 구현한다.
• 통신 및 계산 단계 간의 동적 시간 분할을 통합하여, 두 단계의 전력이 균형을 이루었을 때 최적의 성능을 달성함을 보여준다.

실험 결과

연구 질문

• RQ1 데이터 이질성 하에서 부분 모델 집계는 무선 피어드 학습의 수렴과 성능에 어떤 영향을 미치는가?
• RQ2 PMA-FL에서 전역 손실을 최소화하기 위해 통신 및 계산 에너지 간 최적의 균형은 무엇인가?
• RQ3 장치 스케줄링, 대역폭 할당, 시간 분할을 공동으로 최적화하면 학습 정확도를 향상시키고 에너지 및 시간 비용을 절감할 수 있는가?
• RQ4 строг한 지연 및 에너지 제약 조건 하에서 PMA-FL은 기존 기준 대비 어떻게 성능을 발휘하는가?
• RQ5 불균형한 데이터 크기는 PMA-FL의 수렴 및 정확도에 어떤 영향을 미치는가?

주요 결과

• 비독립 동일 분포 데이터 하에서 PMA-FL은 기준 기법 대비 MNIST에서 3.13% 높은 테스트 정확도를 달성하고, CIFAR-10에서는 11.8% 향상되었다.
• 제안된 공동 최적화는 MNIST에서 최대 29%의 에너지 절감 또는 20%의 시간 절감을 이끌었으며, 정확도를 유지하거나 향상시켰다. CIFAR-10에서는 에너지 25% 절감 또는 시간 12.5% 절감을 달성하였다.
• 20%의 시간 예산 절감(예: 1.5초 대비 2.5초) 조건에서 제안된 알고리즘은 더 엄격한 지연 제약 조건 하에서 OCEAN보다 높은 정확도를 달성하였다.
• 더 높은 에너지 제약 조건(V=0.001) 하에서도 제안된 알고리즘은 OCEAN과 유사한 성능을 달성하면서도 더 에너지 효율적이었다.
• 통신 및 계산 전력이 균형을 이루었을 때 최적의 시간 분할 정책이 달성되었으며, 이는 전체 시스템 비용을 최소화한다.
• 라플라스 가중치 매개변수 V를 조정하면 정확도와 에너지 소비 간 트레이드오���이 가능하며, V를 정교하게 조정할 경우 최적의 성능을 달성할 수 있다.