[논문 리뷰] SPATL: Salient Parameter Aggregation and Transfer Learning for Heterogeneous Clients in Federated Learning

Sixing Yu, Phuong Nguyen외 4명|arXiv (Cornell University)|2021. 11. 29.

Privacy-Preserving Technologies in Data참고 문헌 25인용 수 5

한 줄 요약

SPATL는 유관한 모델 파라미터를 선택적으로 집계하고 클라이언트별 예측기를 통해 지식을 전달하는 방식으로 이종적이고 비독립적 동일 분포(non-IID) 환경에서 통신 오버헤드를 줄이고 학습 안정성을 향상시키는 분산 학습 프레임워크를 제안한다. 이는 최대 86.45%의 통신 비용 절감과 추론 시 39.7% 감소한 FLOPs를 달성하며, 최신 기법들에 비해 수렴성과 성능이 향상된다.

ABSTRACT

Federated learning~(FL) facilitates the training and deploying AI models on edge devices. Preserving user data privacy in FL introduces several challenges, including expensive communication costs, limited resources, and data heterogeneity. In this paper, we propose SPATL, an FL method that addresses these issues by: (a) introducing a salient parameter selection agent and communicating selected parameters only; (b) splitting a model into a shared encoder and a local predictor, and transferring its knowledge to heterogeneous clients via the locally customized predictor. Additionally, we leverage a gradient control mechanism to further speed up model convergence and increase robustness of training processes. Experiments demonstrate that SPATL reduces communication overhead, accelerates model inference, and enables stable training processes with better results compared to state-of-the-art methods. Our approach reduces communication cost by up to $86.45\%$, accelerates local inference by reducing up to $39.7\%$ FLOPs on VGG-11, and requires $7.4 imes$ less communication overhead when training ResNet-20.

연구 동기 및 목표

• 빈번한 모델 가중치 전송으로 인한 높은 통신 오버헤드 문제 해결.
• 비독립적 동일 분포(non-IID), 이종 클라이언트 환경에서의 학습 안정성 및 수렴성 향상.
• 자원 제약이 있는 엣지 디바이스에서 처리해야 할 파라미터 수를 최소화하여 로컬 추론 비용 절감.
• 데이터 이종성에 대응하기 위해 공유 인코더에서 클라이언트별 예측기로 효과적인 지식 전달 가능화.
• 대규모 탈중앙화된 AI 구현에 적합한 확장성 있고 효율적인 분산 학습 프레임워크 개발.

제안 방법

• 그래프 신경망(GNNs)과 다층 퍼셉트론(MLPs) 기반의 중요한 파라미터 선택 에이전트를 도입하여 모델 인코더에서 가장 중요한 파라미터만 선택적으로 전송.
• 모델을 공유 인코더와 클라이언트별 예측기 헤드로 분해하여, 미세조정된 로컬 예측기를 통해 비독립적 동일 분포 데이터에 지식 전달 가능.
• 온라인 강화학습을 사용해 각 클라이언트에서 중요한 파라미터 선택 에이전트를 미세조정하여 로컬 데이터 분포에 적응시키며 최소한의 오버헤드 유지.
• 로컬 및 서버 환경에서 제어 변수(control variates)를 사용한 기울기 제어 메커니즘을 구현하여 기울기 이탈을 보정하고 로컬 업데이트를 글로벌 최적화 방향과 일치시킴.
• 강화학습 에이전트의 사전 훈련 단계에서 모델 압축 및 프루닝 기법을 적용하여 엣지 배포에 적합한 낮은 추론 비용(26 KB 메모리, V100 기준 0.36 ms) 확보.
• 다양한 모델 아키텍처(예: ResNet-18, ResNet-56)에 대해 사전 훈련된 중요한 파라미터 선택 에이전트를 재사용하여 최소한의 미세조정으로 전이 학습 활용.

실험 결과

연구 질문

• RQ1 중요한 파라미터의 선택적 집합이 모델 정확도를 훼손하지 않으면서도 분산 학습에서 통신 비용을 크게 줄일 수 있는가?
• RQ2 클라이언트별 예측기를 통한 지식 전달이 비독립적 동일 분포, 이종 데이터 분포에서 모델 성능 향상에 얼마나 효과적인가?
• RQ3 가벼운 강화학습 에이전트가 엣지 디바이스에서 효율적으로 미세조정되어 로컬 데이터에 맞게 중요한 파라미터 선택을 적응시킬 수 있는가?
• RQ4 기울기 제어 기법이 탈중앙적, 비독립적 동일 분포 분산 학습 환경에서 학습 안정성과 수렴 속도 향상에 어느 정도 기여하는가?
• RQ5 SPATL는 특히 과도하게 파rameter화된 모델들(VGG, ResNet 등)에 대해 어떻게 성능을 발휘하는가?

주요 결과

• SPATL는 최신 기법 대비 통신 비용을 최대 86.45% 감소시키며, ResNet-20 기준 최대 7.4배의 감소를 기록.
• VGG-11 기준 로컬 추론을 위해 최대 39.7%의 FLOPs를 절감하여 엣지 디바이스에서의 효율성 향상.
• 클라이언트별 예측기를 통한 전이 학습은 비전이 학습 기반 모델 대비 최대 19.86% 향상된 정확도 달성.
• 기울기 제어 메커니즘이 학습을 안정화시키고, 특히 기울기 분산이 높은 비독립적 동일 분포 환경에서 수렴 속도 향상에 기여.
• 사전 훈련된 중요한 파라미터 선택 에이전트는 약 40회의 RL 업데이트 내에 신속하게 수렴하며, 미세조정 후 ResNet-18 및 ResNet-56 등 다양한 아키텍처에서 유사한 성능 기록.
• 강화학습 에이전트의 추론 비용은 극히 낮아(26 KB 메모리, V100 기준 0.36 ms) 엣지 디바이스에의 배포가 가능하다.