[논문리뷰] Scaffold: Stochastic Controlled Averaging for Federated Learning

Scaffold: Stochastic Controlled Averaging for Federated Learning

Client Drift 현상을 보여주는 그림

Abstract & Introduction

FedAVG는 간단하고 적은 communication cost로 인하여 많이 사용되어왔다. 하지만, 최근 연구보고에 의하면, 그러한 성능들이 완벽하게 이해가 되지 않는다. 우리는 FedAVG를 위한 수렴 비율을 얻을 것이며, 데이터가 heterogeneous (non-iid) 상태에서 "client-drift"가 발생하고, 그 결과 수렴이 불안정하거나 매우 느려지는 단점을 증명한다.

이에 대한 해결책으로, 우리는 새로운 알고리즘 SCAFFOLD, local updates 과정에서 발생하는 “client-drift” 현상을 수정하기 위해 control variates를 사용한다. 우리는 SCAFFOLD가 상당히 적은 communication rounds를 필요로하고, data heterogeneity로 부터 영향을 받지 않는다는 것을 증명한다. 더욱이, 우리는 SCAFFOLD가 빠른 수렴에도 장점이 있다는 것을 보여준다.

 

Contribution

• heterogeneous data 환경에서의 convex, non-convex 함수들로 부터 이전에 잘 알려진 방법보다 FedAVG의 더 엄격한 수렴 비율을 다룬다.
• client sampling, full batch gradient 상황을 제외하고 수렴이 느리다는 것을 증명
• heterogeneous data 안에서 수렴이 빠른 SCAFFOLD 제안
• SCAFFOLD가 추가적으로 communication round 줄여주는 것도 확인
• 클라이언트 샘플링 비율과 무관하게 Federated Learning에 적합하다.

 

Method

correction term을 추가하여 local update 수행. 방법은 아래 수식과 같음

 

$y_i$ ← $y_i - \eta_l(g_i(y_i) + c - c_i)$

$y_i$는 i번째 클라이언트의 local model, $g_i$는 i번째 클라이언트의 기울기, $c, c_i$는 각각 global model의 correction term, i번째 클라이언트의 local model correction term.

 

Update correction term

$c_i^+$ ← $c_i -c + \frac{1}{K\eta_l}(x-y_i)$

$c$ ← $c + \frac{1}{N} \sum_{i\in S}(c_i^+ - c_i)$

 

correction 값만 빼면 FedAVG와 동일하다. cost 걱정이 없으면 매번 글로벌모델을 업데이트 하는게 성능이 가장 좋지만, cost 문제를 제외할 수는 없다.

 

 

 

Paper: https://arxiv.org/pdf/2001.06202.pdf

GIithub:https://github.com/ki-ljl/Scaffold-Federated-Learning