논문 리뷰) (Mobisys 2022) FedBalancer: Data and Pace Control for Efficient Federated Learning on Heterogeneous Clients - (2)

이전 게시글

2023.01.11 - [Papers/Federated Learning] - 논문 리뷰) (Mobisys 2022) FedBalancer: Data and Pace Control for Efficient Federated Learning on Heterogeneous Clients - (1)

 

FedBalancer 리뷰 두번째 게시글이다. 생각보다 글이 자꾸 길어진다.,,

FedBalancer

FedBalancer의 overview이다. 총 7개의 step이 있다. (1) Server에서는 Client들에게 model weight와 loss threshold, deadline 총 세개의 정보를 보낸다. Client 쪽에서는 (2) Sample selection module에서 학습할 data를 선별하여 (3) Local model training을 수행한다. 이 과정에서 loss threshold와 deadline 설정에 활용되는 metadata와 update된 새로운 model weight가 나온다(2, 3). (4) 얘네들을 서버에 다시 보내준다. (5) Server는 받은 정보를 aggregation을 수행한다. 이후 metadata들은 (6) 새로운 Loss threshold를 결정하는 loss threshold selection module과 (7) 새로운 deadline을 설정하는 deadline selection module을 통과하여 client쪽에 뿌려줄 loss threshold와 deadline을 결정한다. 이러면 FL의 한 round가 결정된다.

이제 loss threhold가 무엇이고 어떻게 이걸 통해 sample을 선별하고 metadata는 무엇이고 어떻게 성능이 개선되는지를 보겠다. 이를 시작하기 앞서 저자들이 이야기한 challenge들을 살펴보면 다음과 같다.

 

첫째, sample을 선별하는 것이 성능을 떨어뜨리면 안된다. 학습에 활용되는 데이터 숫자를 줄이면 학습시간이 줄어드는 것은 당연한데, 여기서 성능을 유지하는 것은 당연한 일이 아니다. 즉 적게 보고도 많이 배울 수 있도록 해야 한다. 그런데 FL 상황에서는 다른 client들의 data distribution은 알 수 없으므로 global model 입장에서 유의미한 sample인지 여부를 아는 것은 매우 어려운 일이된다.

 

둘째, FL 알고리즘 위에서 privacy를 잘 지킬 수 있는 방식이어야 한다는 것. 학습에 용이한 sample을 선별하기 위해서는 sample이 가지고 있는 성능에 대한 정보를 알아야 하는데, 이 정보가 가지고 있는 정보량이 너무 커서 학습에 사용된 데이터를 다시 재건할 수 있는 경우, 즉 해당 정보를 가지고 model inversion attack을 수행할 수 있으면 안된다. Federated learning이 가지는 privacy guarantee를 지켜줘야 한다는 것. 결국 첫 challenge를 잘 지킬 수 있을만큼 정보를 가지면서 동시에 privacy guarantee도 할 수 있는 그러한 정보를 찾아 활용해야 한다는 것이다.

 

마지막은 최적화된 deadline을 adaptive하게 찾는 것이다. Motivational study에서 보았듯 적절한 deadline을 설정하는 것은 꽤나 중요한 일이다. 그러한데 deadline을 설정함에 있어서 학습에 사용되는 데이터의 크기를 고려하여야 하는데, FedBalancer는 학습에 사용되는 sample이 자꾸 바뀌기 때문에 deadline도 이걸 잘 조절해줘야 한다는 것이다. 논문에 직접적인 언급은 없지만, 결국 해당 논문에서 해결해야 하는 문제는 동적인 상황에서 적절히 deadline을 설정함과 동시에 목표는 time-to-accuracy를 높이는 것이기에 연산량도, 연산 시간도 잘 통제하는 것이다.

Sample Selection

FedBalancer의 시작점인 학습에 활용될 sample 선정하기의 단계이다. 앞서 overview figure를 통해 설명하면서 server 측에서 client에게 weight와 loss threshold, deadline 3개의 정보를 보낸다고 했는데, 그중 loss threshold를 활용하여 중요한 sample을 선정한다. 이러한 task가 importance sampling이라고 불리는 task와 유사하다고 논문에서는 언급하고 있으며, 일반적으로 현재 model에 중요한 sample을 고르는 상황에서 loss, gradient norm, gradient norm upper bound와 같은 criteria를 활용한다고 하는데, FedBalancer에서는 non gradient-based optimization 상황에도 활용될 수 있는 loss를 활용하겠다고 한다.

 

개인적인 해석을 덧붙이자면, loss의 경우 forwarding만을 통해 결과를 얻을 수 있으나, gradient norm등은 loss를 계산하면서 동시에 gradient를 계산해야 하다보니 연산량과 연산 시간의 측면에서 손해가 크다. Time-to-Accuracy를 높이고 싶은 상황에서 sample을 계산하기 위해서 시간과 연산 리소스를 투자하는 것이 적절하지 못하지 않은가? Loss를 사용하는 이유는 단순히 non gradient-based optimization에서도 적용 가능하게 하기 위한 것만은 아니지 않을까.

위 알고리즘은 Sample selection 알고리즘을 정리한 것인데 하나씩 살펴보자. 우선 기본적으로 아래 첨자 $R$과 위 첨자 $i$는각각 R번째 round와 i번째 client라는 의미이다. $D$, $B$는 각각 dataset, batch training latency이다. $loss$, $lt$, $ddl$는 모두 loss list, loss threshold, deadline이다. $E$는 local에서 학습할 epoch이다.

maxTrainableSize()는 해당 client가 주어진 시간내에 얼마나 많은 수의 batch를 training 할 수 있는지를 알려주는 함수이다. 이 크기가 전체 dataset보다 크다면, 전체 dataset을 모두 학습에 활용한다(line 3-4). 만약 그렇지 않다면, sample 일부를 버리는 선택을 한다. 우선 loss list에는 client가 가지고 있는 training set에 대한 loss 값들이다. $D$안의 모든 sample들 중에 주어진 threshold ($lt$)보다 큰 sample은 $OT$ (over-threshold), 작은 sample은 $UT$ (under-threshold)에 저장한다 (line5-9).

 

논문에서는 이들 중 threshold를 넘긴 $OT$ list 속의 sample들, 즉 loss 값이 loss threshold보다 큰 sample들이 학습에 있어 유의미한 sample들로 선정하고 있다. 뭐 loss값이 크다는 것은 아직 학습이 제대로 되지 않았다는 의미로 보겠다는 것이니 아주 틀린 가정은 아니라고 생각한다. 이제 $OT$ list에서 일부, $UT$ list에서 일부를 선정한다. 이 일부에 대한 비율은 parameter $p$로 결정한다. 앞선 가정을 따라 $OT$에서 더 많은 비율로 sample을 뽑는다. "그런데 왜 UT는 상대적으로 덜 중요하다고 가정했으면서 일정 비율을 뽑는가?" 하는 질문을 할 수 있다. 이는 catastrophic forgetting이라는 문제를 완화하기 위해서다.

 

catastrophic forgetting은 continual learning (CL) 논문에서 주로 등장하는 개념인데, 매우 간단하게 이야기 하자면 "새로운 것을 배우느라 과거에 학습한 것을 까먹는 것"이다. 즉, 어느 시점에 학습이 잘 되어있는 sample이라고 해서 앞으로도 학습이 잘 되어 있으리라는 보장은 없다는 것. round $R$에서는 학습이 잘 되어 $UT$에 속했던 어떤 sample이 학습에서 배제된 채로 계속해서 학습을 하지 않는 다면 model은 다른 sample들에 fitting되며 이 sample에 대한 정보는 잊어버려 기존의 성능을 보이지 못하게 된다는 것이다. 따라서 $UT$에서도 일부 sample들을 뽑아줄 필요가 있다.

이리하여 OT와 UT에서 사전에 정의해둔 비율으로 sample들을 추출하고 이들을 학습에 사용할 training set으로 최종 선택한다. 그나저나 해당 알고리즘에서 인자로 넘어온 $w_{R}$은 사용되지 않는데 아마도 model weight겠지...?

Loss Threshold Selection

앞서 sample을 선별하는 과정에서 가장 중요한 기준이 되는 값은 바로 Loss Threshold이다. 아까는 그냥 그러려니 하고 넘어갔지만, 생각해보면 이 loss threshold라는 녀석을 어떻게 정해주느냐가 꽤나 중요한 문제임을 알 수 있다. 이 값이 너무 작다면 대부분의 sample들이 더 학습이 필요한 sample들이라고 판단되어 해결하고자 하였던 학습에 활용되는 sample의 수를 줄이는 문제를 해결할 수 없어지고, 이 값이 너무 크다면 sample들이 충분하지 못해 학습이 제대로 해결되지 못하는 문제가 발생할 수 있다.

 

심지어, model은 계속해서 학습이 되면서 충분히 학습된 sample의 수는 이상적으로는 점점 늘어나고, loss 값도 전반적으로 떨어질테니 이 loss threshold는 하나의 값으로 고정해 둘 것이 아니라 FL이 진행되는 과정에서 adaptive하게 값이 변해야 하는 것이다. 매 round마다 이 값을 manual하게 정해주고 자빠졌을수는 없는 노릇이니 이 loss threshold를 적절하게 설정해줄 필요가 있다.

FedBalancer는 다음과 같은 방식을 채택한다. R번째 라운드에서 각 client들로부터 $LLow^{i}$와 $LHigh^{i}$를 받아 list형태로 저장한다. 이때 $LLow^{i}$는 client $i$의 loss list중 가장 작은 minimum 값, $LHigh^{i}$는 가장큰 값을 100%라 할 때, 80% 위치에 있는 값이다. $LLow$는 min값이면서 $LHigh$는 max 값을 사용하지 않는 이유는, noisy한 data의 경우 loss의 값이 클 수 있고 따라서 실제의 분포를 예측하는 것에 방해가 될 수 있기에 이러한 sample들은 배제하기 위해 저자들이 설정한 기준이다.

 

이렇게 만든 $LLow$와 $LHigh$ list에서 각각 min값과 평균을 계산하여 각각 $ll$, $lh$라고 정의했다. 새로운 loss thhreshold $lt_{R+1}$은 $ll + (lh-ll)*ltr$이다. 여기서 $ltr$은 loss threshold와는 또 다른 값이라는 것이다. 풀어 쓰자면, 가장 작은 값보다는 기본적으로 크며 ($ll \leq lh$이고, $ltr$은 비율으로 0에서 1사이의 값을 가진다) $ll$과 $lh$ 사이를 $ltr$ 비율로 나눈 한 지점을 새로운 loss threshold로 잡는 것이다.

 

어려운 것은 없다. 다만 이 $ltr$ 값도 또 설정해줘야 하는 값이다. 그런데 생각해보면, $ll$과 $lh$가 변하고 학습이 진행됨에 따라서 점점 $ll$과 $lh$들은 값이 작아질 것이며, threshold는 점진적으로 키워줘야 학습에 사용되는 sample의 수를 합리적으로 줄일 수 있다. 따라서 $ltr$ 값도 adaptive하게 조정해줘야 하는 값이다.

 

FedBalance에서는 $ltr$값을 0.0으로 초기화하고 $lss$ (loss threshold step size)를 통해 점진적으로 증가시키는 방식을 채택하고 있다. 또한 $ltr$값은 deadline 설정에 사용되는 $ddlr$ (deadline ratio)을 조절하는 과정에서도 사용된다.

 

이번 글은 여기까지 하고 다음 게시글에서 $ltr$ 설정을 비롯하여 deadline 설정 등 남은 시스템의 구성 요소와 적용된 알고리즘을 살펴보겠다.