[ICLR 2020] q-FFL, q-FedAvg - (4)

논문 제목: Fair Resource Allocation in Federated Learning

출처: https://arxiv.org/abs/1905.10497

 

 지난 포스트에서 q-FFL에 관한 이야기는 마무리지었고, 이번 포스트에서는 q-FFL의 solver인 q-FedAvg에 관한 내용을 다룰 것입니다. (이전 글 보기) 마치 FedSGD에서 FedAvg로 넘어가듯이, 우선 q-FedSGD부터 살펴보도록 하겠습니다.

 

7. q-FedSGD

 

 q-FFL의 핵심은 performance와 fairness 사이의 trade-off를 잘 조절할 수 있는 최적의 $q$를 찾는 것입니다. 다만, 특정 $q \in \mathbb{R}_{\geq 0}$가 어떠한 수준의 fairness와 대응하는지는 알 수 없기 때문에, 여러 가지의 $q$에 해당하는 $f_q$를 학습시킨 후, 그중 가장 목적에 부합하는 것을 사용하는 방식을 택하는 것이 하나의 방법이 될 것입니다. 하지만 이 방법은 hyperparameter를 각 $f_q$마다 tuning해주어야 한다는 문제점을 가지고 있습니다. 특히 step size(내지는 learning rate)의 경우 objective function의 gradient가 가지는 Lipschitz constant의 역수에 비례하기 때문에, 다양한 $q$에 대해서 다양한 step size를 탐색하는 것이 현실적으로는 불가능할 것입니다. 이러한 점을 고려하여, 저자들은 다음과 같은 방식을 제안합니다.

 

$\text{Lemma 3}$

 

 만약 nonnegative function $f(\cdot)$가 상수 $L > 0$에 대해서 $L$-Lipschitz continuous한 gradient를 가지고 있다면, 임의의 $q \geq 0$와 $w$에 대해서

$$L_q (w) := L f(w)^q + q f(w) ^{q - 1} ||\nabla f(w)||^2$$

는 $\frac {1} {q + 1} f^{q + 1} (\cdot)$의 $w$에서의 gradient에 대한 local Lipschitz constant의 upper bound이다.

 

$\text{Proof}$

 임의의 $w$에 대해서, $\frac {1} {q + 1} f^{q + 1} (w)$의 Hessian은 다음과 같다:

$$\nabla^2 \left( \frac {1} {q + 1} f^{q + 1} (w) \right) = q f^{q - 1} (w) \left( \nabla f(w) \right)^T  \nabla f(w)  + f^q(w) \nabla^2 f(w)$$ 이때, $\left( \nabla f(w) \right)^T  \nabla f(w) \preceq ||\nabla f(w)||^2 \times I$이고, 정의 상 $\nabla^2 f(w) \preceq L \times I$이므로, \begin{align*} \left| \left| \nabla^2 \left( \frac {1} {q + 1} f^{q + 1} (w) \right) \right| \right|_2 &= \left| \left| q f^{q - 1} (w) \nabla f(w) \left( \nabla f(w) \right)^T + f^q(w) \nabla^2 f(w) \right| \right|_2 \\&\leq q f(w) ^{q - 1} ||\nabla f(w)||^2 + L f(w)^q \\&=: L_q (w) \end{align*}

이다. $\square$

 

 $\text{Lemma 3}$ 덕분에, 이제는 더 이상 각 $q$마다 step size를 설정할 필요가 없어졌습니다. 대신, 특정한 한 가지의 $q$ 값(가령, $q = 0$)에 대해서만 step size를 tuning할 것이며, 그 결과 얻어낸 $L$의 근사치를 사용하여 해당 값 이외의 $q$ 값에 대하여 $\text{Lemma 3}$에서 정의된 $L_q (w)$를 구할 것입니다. 그리고 이것이 step size의 역할을 대체하게 됩니다. 특히, $q = 0$인 경우가 다른 $q$ 값들에 대하여 계산하는 과정에서 사용하기에 가장 용이하다는 점을 감안하면, 앞으로 우리가 할 일은 FedSGD(즉, zero fairness인 상황)에서 step size를 tuning하고, 이후 다양한 $q > 0$를 대입해보면서 최적의 $q$를 찾는 것입니다. 

 

 물론, 방금 설명하였듯이 처음 1회는 Lipschitz constant $L$을 추정하기 위하여 step size를 tuning할 필요가 있고, 사실 이 부분 때문에 "그렇다면, q-FedAvg는 $L$을 얻어내는 과정을 추가로 진행해야 하므로 다른 알고리즘들보다 더 많은 시간이 걸리는 것이 아닌가?"에 대한 논쟁이 있었습니다. 다만, 이는 hyperparameter tuning을 대신하는 과정이기 때문에, $L$을 추정하는 시간을 학습 시간에 포함하는 것이 타당한지는 조금 더 생각해보아야 할 부분입니다. 한편, $L_q (w)$는 supremum이 아니라 단순히 upper bound이기 때문에, local Lipschitz constant와의 간극을 얼마나 더 좁힐 수 있는지도 관건이 될 것 같은데, 이 부분에 대한 언급이 부족한 것은 다소 아쉽습니다. 여하튼, 아래의 pseudo code가 q-FedSGD 알고리즘을 잘 표현해주고 있는데, 이미 $q = 0$인 case를 계산하는 과정에서 $L$을 알아낸 상황이라고 가정합시다.

 

8. q-FedAvg

 

 FedSGD의 communication cost를 줄이고자 FedAvg를 고안하였듯이, q-FedAvg도 생각해볼 수 있습니다. 다만, 우리는 여러 paper를 통해서 local에서의 gradient를 단순히 합치는 것이 원래의 model을 적절하게 대표하지 못한다는 것을 확인하였습니다. 특히, q-FedSGD의 경우 $F_k^{q + 1}$의 지수 $q + 1$ 때문에, $q = 0$가 아닌 이상 더더욱 기존의 FedAvg 방식을 이용하기 어렵습니다. 저자들은 이 부분을 피해가기 위해서, $\nabla F_k (w^t)$ 대신 $\Delta w_k^t := L (w^t - \bar{w}_k^{t + 1})$을 사용하기로 하였습니다. 즉, local에서는 일반적인 SGD 방식의 update를 수행하고, global에서는 q-FedSGD 방식의 update를 수행하는 방식으로 q-FedAvg를 구성한 것입니다. 이러한 q-FedAvg의 pseudo code는 아래와 같습니다. (이때, step size $\eta$는 local SGD 계산 과정에서 사용됩니다.)

 

 q-FedAvg 역시, $q = 0$인 경우는 FedAvg에 해당하고, $q = \infty$인 경우는 Agnostic FL에 해당합니다만, local update 과정의 유무 때문에 $q = \infty$인 q-FedAvg가 기존의 Agnostic FL과 비슷한 성능을 보여주면서도 더 빨리 수렴한다고 저자들은 이야기합니다. 다음 포스트에서는 이 부분을 위시한 여러 experiments를 살펴보도록 하겠습니다.