https://arxiv.org/abs/2410.01256
ParallelSFL: A Novel Split Federated Learning Framework Tackling Heterogeneity Issues
Mobile devices contribute more than half of the world's web traffic, providing massive and diverse data for powering various federated learning (FL) applications. In order to avoid the communication bottleneck on the parameter server (PS) and accelerate th
arxiv.org
오늘 리뷰할 논문은 MobiCom 2024에 게재된 ParallelSFL이다. 현실적으로 구성된 테스트베드 환경에서 "split federated learning" 전략이 얼마나 좋은 성능을 내는지 다루는 논문이다. 우선 논문의 용어들에 조금 친숙해져 보자.
- 연합학습에 참여하는 클라이언트를 worker라고 칭한다.
- split federated learning = SFL, 이 논문에서는 SFL 프레임워크를 제안하고 있다.
- parameter server = PS, 연합학습 환경에서의 글로벌 모델 즉 서버
- edge computing = EC
언제나 그렇듯 엣지 컴퓨팅 환경에서는 worker heterogeneity를 고려한 효율적인 학습이 중요한데, ParallelSFL은 모델을 bottom model과 top model로 나누어, 두 서브모델을 따로 학습시킨다. 또한 훈련에 참여하는 worker들을 리소스가 비슷한 애들끼리 클러스터링하는 전략을 사용해서, 학습 시간을 줄이고 효율적으로 목표 성능에 달성할 수 있도록 한다.
이제 ParallelSFL의 클러스터링 전략은 어떻게 되는지, 마주치는 주요 챌린지에는 어떤게 있는지 차근차근 알아볼 것이다. 그리고 흥미로웠던 experiment 환경에 대해서도 알아볼 것이다.
Introduction
가용할 수 있는 리소스가 제한적인 클라이언트들은 보통 1~30TOPS(Tera Operations Per Second), 1~8GB 메모리를 가지고 있다. 따라서 LLM같은 대형 모델을 학습시키기에는 무리가 있다. 그래서 경량화된 모델을 넘겨받고 모델의 마지막 레이어만 떼어서 학습시키기도 한다. split federated learning 기술은 large-scale model의 학습을 위해, 모델을 서브모델로 나누어 data parallelism, model parallelism을 실현하는 식으로 제안되어 왔다.
앞서 언급한 것처럼 모델을 top, bottom으로 나누게 된다면 전자는 resource-constrained worker에서 학습되고 후자는 PS로 넘겨진다. 이때 PS는 worker들과 주기적으로 소통해야 하는데, worker가 많은 학습 환경일수록 커뮤니케이션 병목(bottleneck) 현상이 생기게 된다. 이 병목 현상을 피하기 위해서 ParallelSFL에서는 worker들을 다수의 클러스터로 묶어서 각 클러스터가 협력적으로 SFL training을 할 수 있도록 하는 것이다.
이쯤에서 ParallelSFL의 contribution을 정리해보면,
1. 리소스 제한적인 클라이언트 환경에서 클라이언트를 효과적으로 클러스터링하여 SFL 학습을 수행하도록 한다.
2. 클러스터링 할 때, training efficiency와 model accuracy 간 밸런스를 맞출 수 있는 utility function을 수식적으로 정의한다.
3. 기존 벤치마크 데이터에 비해 traffic consumption을 최소 21% 절감하는 등 효율성을 챙김과 동시에 accuracy도 5% 이상 올리는 성능을 보인다.
Method
SFL 환경에서의 학습 능력을 개선하기 위해 어떤 챌린지를 마주쳐야 했을까?
- System Heterogeneity: 각 worker들이 서로 다른 계산 능력, 메모리 크기, 통신 속도 등의 시스템 자원을 보유하고 있어 학습 과정에서 균일한 처리 성능을 유지하지 못하는 현상. 여기서는 성능이 유사한 디바이스들을 클러스터로 그룹화하여 같은 클러스터 내에서 작업을 수행하여 문제를 해결한다. 그렇지 않으면 빠른 디바이스가 느린 디바이스의 학습 종료 시점까지 기다려야 하기 때문에 효율이 저하될 테니까
- Statistical Heterogeneity: 각 worker가 보유한 데이터 분포가 서로 다르거나, 특정 클래스에 치우쳐 있어 데이터가 독립적이고 동일한 분포(i.i.d.)를 따르지 않는 현상. 논문에서는 클러스터를 구성할 때 데이터 분포가 서로 비슷하도록 클러스터링하고, KL-divergence를 사용하여 클러스터 내 데이터 분포와 이상적인 i.i.d. 분포 간의 차이를 최소화하였다.
아래는 일반적인 SFL 아키텍처와 ParallelSFL의 클러스터링 아키텍처의 차이를 보여준다.
기존 SFL 프레임워크에서는 모든 worker가 PS와 직접적으로 통신하며 bottom submodel과 top submodel 간의 데이터(Features, Gradients)를 교환하는 것을 볼 수 있다. 하지만 Parallel 방식에서는 그룹 별로 PS와 통신하여 통신비용이 절감되며, 각 클러스터의 top worker가 클러스터 내 통신을 담당하며 PS로 데이터를 전달하는 것을 볼 수 있다.
Discussion of Privacy Protection
데이터 프라이버시는 어떻게 보호하고 있을까? 기존 SFL에서의 프라이버시 보호 방식은 1. worker가 라벨이 포함된 피처를 서버에 전달해 top 모델에서 forward propagation을 수행하고, 2. 만약 worker가 라벨을 공유하지 않는 경우에 top 모델이 logit(모델 출력값)을 반환하면 worker가 이를 사용해서 계산한 loss를 다시 top모델에 전송하는 방식을 취했다. 이렇게 전송된 loss값으로 gradient를 계산하고 backward propagation을 하는 것이다. (복잡해...)
ParallelSFL에서는 기존 SFL에서 사용된 라벨 공유 방식과 Logit 반환 방식을 모두 사용하다고 한다. 뿐만 아니라 DP(Differential Privacy, 데이터 샘플을 개별적으로 식별할 수 없도록 통계적으로 처리) 또는 Homomorphic Encryption(데이터가 암호화된 상태에서도 연산이 가능하도록 하는 기술) 등을 사용해서 추가적으로 보호할 수도 있다. 다만 여기서는 프라이버시 보호를 주로 다루진 않았다.
Evaluation
ParallelSFL은 방대한 테스트베드를 실제로 구축하여 실험을 진행했다. 총 80개의 NVIDIA Jetson kits(30 Jetson TX2, 40 Jetson NX, 10 Jetson AGX)를 사용하였는데 자세한 스펙은 아래 나와있다.
실험 태스크는 Speech Recognition, Object Recognition, Image Classification, Natural Language Processing 이며, System Heterogeneity를 위해 아래 configuration을 진행했다.
1. Computation: 모든 디바이스를 다른 모드로 설정했으며 time-varying on-device resource를 반영하고자 학습 20 라운드마다 모드를 바꿨다. 가용 CPU 수 등의 컴퓨팅 capacity를 configuration하는 것이다.
2. Communication: 모든 엣지 디바이스가 와이파이 라우터를 통해 PS와 로 연결되어 있다. 재밌는 건 random channel noise를 구현하고자 디바이스 그룹 별로 라우터로부터의 거리를 다르게 했다.(2m, 8m, 14m 등)
Statistical Heterogeneity 세팅에서는 non-IID 데이터셋을 만들기 위해 Dirichlet distribution을 사용했다.
실험 결과 IID와 non-IID 데이터셋에서 ParallelSFL을 포함한 다섯 가지 접근법 성능을 비교한 결과이다. Figure 4는 IID 환경에서의 test acc인데, ParallelSFL이 모든 데이터셋에서 가장 빠른 수렴 속도를 보였다. non-IID 환경인 Figure 5 또한 4개의 베이스라인 모델들과 비교했을 때 가장 높은 time-to-accuracy 성능을 보이는 것을 알 수 있다. 이는 KL-divergence를 최소화한 클러스터링과 모델 분할 방식이 효과적으로 작동했기 때문이다.
이번 논문은 수식이나 알고리즘 설명은 생략하고 대강 흐름만 공부했다. 요즘 하나를 자세히 파고들기가 쉽지 않다... 얼렁 에너지가 충전되어야 할 텐데 말이다.
설날에 여행이나 가면 좋을 것 같은데 노는 것도 부지런해야 할 수 있나보다. 내가 얼른 힘을 찾길 바라며...
다들 새해 복 많이 받으세요 ^0^
