[MLSys 2020] FedProx - (3)

논문 제목: Federated Optimization in Heterogeneous Networks

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

 

 이전 포스트에서, 우리는 non-convex한 $F_k$에 관한 FedProx의 convergence를 증명하였습니다. (이전 글 보기) 이번 포스트에서는 FedProx의 convergence rate를 구해보고, convex case와 $\gamma$가 device 별로, 또 round 별로 달라지는 case에 관해서도 살펴볼 것입니다.

 

6. non-convex case의 convergence rate

 

 $\text{Theorem 6}$ (Convergence rate: FedProx)

  $\epsilon > 0$이 주어졌을 때, $B > B_{\epsilon}$, $\mu$, $\gamma$, 그리고 $K$에 대해서 $\text{Theorem 4}$의 assumption들이 FedProx algorithm의 매 iteration마다 성립한다고 가정하자. 그리고 $f(w^0) - f^* =: \Delta$라고 정의하자. 이때, FedProx algorithm의 $T := O(\frac {\Delta} {\rho \epsilon})$회 iteration이 진행되었다면, $\frac {1} {T} \sum_{t = 0}^{T - 1} \mathbb{E} [||\nabla f(w^t)||^2] \leq \epsilon$이다.

 

 $\text{Theorem 6}$에 따르면, $\epsilon$을 작게 잡을수록 수렴하는 데에 더 많은 iteration이 필요하고, 대신 더 정확한 결과가 도출된다는 것을 알 수 있습니다. 이는, overfitting이 바람직한 상황은 아니라는 것을 감안했을 때, 저자가 앞서 이야기한 "굳이 tight하게 $\epsilon$을 잡을 필요 없다"는 주장과 일맥상통하는 부분입니다. 이를 증명하기 위해서는 추가적인 assumption이 한 가지 필요한데, $\mathbb{E}_k [||\nabla F_k(w) - \nabla f(w)||^2] \leq \sigma ^2$를 만족하는 $\sigma$가 존재한다, 즉, variation이 bounded되어 있다는 것입니다. 이는 앞서 살펴본 $\text{Assumption 1}$과 마찬가지로 크게 문제가 되는 가정은 아닙니다.

 

$\text{Corollary 10}$ (Bounded variance equivalance)

 $\text{Assumption 1}$이 성립하고, $\mathbb{E}_k [||\nabla F_k(w) - \nabla f(w)||^2] \leq \sigma ^2$를 만족하는 $\sigma$가 존재할 때, $B_\epsilon \leq \sqrt{1 + \frac {\sigma ^2} {\epsilon}}$이다.

 

$\text{Proof}$

 $Var[X] = \mathbb{E} [X^2] - \mathbb{E} [X]^2$임을 이용하면  $\mathbb{E}_k [||\nabla F_k (w)||^2] - ||\nabla f(w)||^2 = \mathbb{E}_k [||\nabla F_k(w) - \nabla f(w)||^2] \leq \sigma ^2$임을 알 수 있다.  이를 다시 정리하면 $\mathbb{E}_k [||\nabla F_k (w)||^2] \leq \sigma ^2 + ||\nabla f(w)||^2$로 나타낼 수 있으며, $B$의 정의에 따라서 $B_\epsilon = \sqrt{\frac {\mathbb{E}_k[||\nabla F_k(w)||^2]} {||\nabla f(w)||^2}} \leq \sqrt{\frac {\sigma ^2 + ||\nabla f(w)||^2} {||\nabla f(w)||^2}} =  \sqrt{1 + \frac {\sigma ^2} {||\nabla f(w)||^2}} \leq \sqrt{1 + \frac {\sigma ^2} {\epsilon}}$이다. $\square$

 

$\text{Proof of Theorem 6}$

 $\text{Corollary 10}$의 결과를 $\text{Theorem 4}$의 $\rho$에 대입하면 알 수 있다. $\square$

 

7. convex case의 convergence

 

$\text{Corollary 7}$ (Convergence: Convex case)

 $\text{Theorem 4}$의 assumption이 모두 성립하고, $F_k(\cdot)$들이 모두 convex하다고 가정하자. 더 나아가, 모든 $k$와 $t$에 대해서 $\gamma_k^t = 0$을 만족한다고 하자. (이는 곧 모든 local update가 exact하게 이루어짐을 뜻한다.) 이때, 만약 $1 \ll B \leq 0.5 \sqrt{K}$라면, 최적의 $\mu$는 $\mu \approx 6LB^2$이며, 따라서 $\rho \approx \frac {1} {24LB}$이다.

 

$\text{Proof}$

 $\rho = \left( \frac {1} {\mu} - \frac {\gamma B} {\mu} - \frac {B(1 + \gamma) \sqrt{2}} {\bar{\mu} \sqrt{K}} - \frac {LB(1 + \gamma)} {\bar{\mu} \mu} - \frac {L(1 + \gamma)^2 B^2} {2 \bar{\mu}^2} - \frac {LB^2 (1 + \gamma)^2} {\bar{\mu}^2 K} (2 \sqrt{2K} + 2) \right)$임을 알고 있고, 가정에 따라 $L_- = 0$, $\bar{\mu} = \mu$, $\gamma = 0$이다. 따라서 식을 정리하면 $\rho = \left( \frac {1} {\mu} - \frac {B \sqrt{2}} {\mu \sqrt{K}} - \frac {LB} {\mu^2} - \frac {LB^2} {2 \mu^2} - \frac {LB^2} {\mu^2 K} (2 \sqrt{2K} + 2) \right)$이다. 또한, $B \leq 0.5 \sqrt{K}$이므로 $\rho \geq \frac {(2 - \sqrt{2})} {2 \mu} - \frac {LB + (1 + \sqrt{2})LB^2} {2 \mu^2}$이다. 따라서,

 

$$\begin{align*} \mathbb{E}_{S_t} [f(w^{t+1})] &\leq f(w^t) - \rho ||\nabla f(w^t)||^2 \\&\leq f(w^t) - \frac {2 - \sqrt{2}} {2 \mu} ||\nabla f(w^t)||^2 + \frac {LB + (1 + \sqrt{2})LB^2} {2 \mu^2} ||\nabla f(w^t)||^2 \\&\lessapprox f(w^t) - \frac {1} {2 \mu} ||\nabla f(w^t)||^2 + \frac {3LB^2} {2 \mu^2} ||\nabla f(w^t)||^2        \end{align*}$$

 

이며, 더 나아가 $\mu \approx 6LB^2$로 잡을 경우, $\mathbb{E}_{S_t} [f(w^{t+1})] \lessapprox f(w^t) - \frac {1} {24LB^2} ||\nabla f(w^t)||^2$이다. $\square$

 

 convex case의 convergence는 최소 $T := O(\frac {LB^2 \Delta} {\epsilon})$회의 iteration 후에 진행되며, 이는 $\text{Theorem 6}$과 $\text{Corollary 7}$으로부터 자명합니다. 또한, bounded variance assumption에 의하여 $O (\frac {L \Delta} {\epsilon} + \frac {L \Delta \sigma ^2} {\epsilon ^2})$로도 표기할 수 있습니다. 이는 SGD의 시간복잡도와 정확히 일치하며, 저자들은 이러한 점을 근거로 해당 논문이 의미있다고 주장하고 있습니다. 물론, 이것이 FedProx가 distributed SGD보다 우위에 있다는 것을 의미하지는 않습니다. 그렇다기보다는, 연합학습 체계 중에서 convergence를 엄밀하게 증명해보인 첫 논문이라는 점, 그리고 그 속도가 SGD에 준한다는 점에서 해당 논문이 의미가 있다고 생각합니다.

 또한, $\text{Corollary 7}$을 통해서 알 수 있는 것이 한 가지 더 있는데, 비록 convex case로 한정지었지만, $\mu$를 어떻게 잡는지가 model의 convergence에 큰 영향을 준다는 것입니다. (물론, non-convex case에서도 똑같이 증명하기는 어렵겠지만 비슷한 결과가 나오리라 예상할 수 있고, 이 부분은 empirical test에서 확인하도록 하겠습니다.) 저는 처음에 이 부분을 명확하게 이해하기 어려웠는데, 분명 논문 초반부에서부터 계속해서 $\gamma_k^t$의 중요성을 언급하였음에도, 저자들은 $\gamma$가 아닌 $\mu$를 hyperparameter로 잡았기 때문입니다. (심지어, 이 $\mu$는 proximal term에서부터 정의도 되지 않은 채 쓰이고 있습니다.) 그렇다면, $\gamma$는 어디에서 사용되는 것이고, 또 hyperparameter가 아니라면 어디에서부터 주어지는 것일까요? 이에 관해서는 다음 포스트에서 ablation study와 함께 알아보도록 하고, 이에 관한 이야기를 하기 위해서 $\gamma$를 $\gamma_k^t$로 generalize하는 것으로 이번 포스트를 마무리짓겠습니다.

 

8. $\gamma$가 달라지는 경우의 convergence

 

$\text{Corollary 9}$ (Convergence: variable $\gamma$'s)

 local function $F_k$들이 모두 non-convex하고, $L$-Lipschitz smooth하다고 가정하자. 그리고 $L_- > 0$가 존재하여 $\nabla^2 F_k \succeq -L_- \textbf{I}$, $\bar{\mu} := \mu - L_{-} > 0$를 만족한다고 하자. $w^t$가 non-stationary한 solution이고 각 device의 local functions $F_k$가 $B$-dissimilar할 때 (즉, $B(w^t) \leq B$일 때), $\mu$, $K$, $\gamma_k^t$가 $$\rho^t := \left( \frac {1} {\mu} - \frac {\gamma^t B} {\mu} - \frac {B(1 + \gamma^t) \sqrt{2}} {\bar{\mu} \sqrt{K}} - \frac {LB(1 + \gamma^t)} {\bar{\mu} \mu} - \frac {L(1 + \gamma^t)^2 B^2} {2 \bar{\mu}^2} - \frac {LB^2 (1 + \gamma^t)^2} {\bar{\mu}^2 K} (2 \sqrt{2K} + 2) \right) > 0$$ 를 만족한다면, FedProx의 $t$번째 iteration에서 global objective function $f$의 expected decrease는 $\mathbb{E}_{S_t} [f(w^{t + 1})] \leq f(w^t) - \rho^t ||\nabla f(w^t)||^2$를 만족한다. (여기에서, $S_t$는 $t$번째 iteration에 참여하는 device들의 집합이고, $\gamma^t := max_{k \in S_t} \; \gamma_k^t$이다.)

 

$\text{Proof}$

 $\gamma^t$의 정의와 $\text{Theorem 4}$로부터 자명하다. $\square$

 

 이번 포스트에서는 FedProx의 convergence rate에 대한 증명과 convex case에서의 convergence 증명을 살펴보고, $\gamma$가 달라지는 case로 증명을 확장해보았습니다. 그리고 이 과정에서 hyperparameter $\mu$의 중요성도 확인해보았습니다. 다음 포스트에서는 도대체 $\gamma$가 무엇인지 (알아보자고 몇 번째 말하고 있는 것 같지만) 알아보고, 각종 ablation test를 살펴보겠습니다. (다음 글 보기)