[MLSys 2020] FedProx - (2)

논문 제목: Federated Optimization in Heterogeneous Networks

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

 

 이전 포스트에서, 우리는 FedProx 알고리즘이 어떠한 방식으로 작동하는지 알아보았습니다. (이전 글 보기) 이번 포스트에서는 FedProx의 Convergence 증명 과정을 자세하게 확인해보겠습니다.

 

4. Bounded Dissimilarity

 

 우선, 증명 과정에서 사용되는 assumption 한 가지를 확인하도록 하겠습니다.

 

$\text{Definitnion 3}$ ($B$-local dissimilarity)

 $\mathbb{E} [||\nabla F_k(w)||^2] \leq ||\nabla f(w)||^2 B^2$를 만족하는 local function $F_k$들을 $w$에서 $B$-locally dissimilar하다고 정의한다. 더 나아가 $||\nabla f(w)|| \neq 0$인 경우, $B(w) = \sqrt{\frac {\mathbb{E}_k[||\nabla F_k(w)||^2]} {||\nabla f(w)||^2}}$로 정의한다. (여기에서, $f(w) = \sum_{k = 1}^N p_k F_K(w) = \mathbb{E}_k [F_k(w)]$이고, $p_k = \frac {n_k} {n}$이다.)

 

※ 주의: $B$-local dissimilarity의 $B$는 FedAvg에 사용되는 mini-batch size $B$와 전혀 관련이 없습니다.

 

 $B$-local dissimilarity가 의미하는 바가 무엇인지 살펴봅시다. 만약 모든 $k \in K$에 대해서 local function $F_k$가 동일한 값을 갖는다면, $B = 1$이 될 것입니다. 하지만 연합학습의 경우 이러한 일은 쉽게 일어나지 않기 때문에 일반적으로 $B > 1$입니다. 또한, 만일 IID condition이 만족된다고 하더라도, SGD와 같은 stochastic algorithm을 사용하는 경우 sampling 과정에서 서로 다른 data를 사용하기 때문에 $B = 1$은 이론적으로만 가능한 수치입니다. ($B = 1$ case는 $w$가 stationary한 case로 이해하는 것이 적절합니다.) 정리하자면, $B(w) \geq 1$이며, $B$가 클수록 local function들 간의 dissimilarity가 크다는 것을 의미합니다. 그리고 이러한  $B$-local dissimilarity의 정의와 관련하여, 저자들은 한 가지 assumption을 언급합니다.

 

$\text{Assumption 1}$ (Bounded dissimilarity)

어떠한 $\epsilon > 0$에 대하여, 다음을 만족하는 $B_{\epsilon}$이 존재한다: 임의의 $w \in \mathcal{S}_{\epsilon}^{c} := \{w \; | \; ||\nabla f(w)||^2 > \epsilon \}$에 대하여, $B(w) \leq B_{\epsilon}$.

 

 말 그대로 dissimilarity가 bounded되어 있다는 것인데, 그렇게 강한 가정은 아니기 때문에 받아들이는 데에 있어 큰 문제가 되지는 않습니다. 물론, 연합학습의 경우 data들이 non-IID하기 때문에 dissimilarity가 다소 크겠지만, 현실적으로 dissimilarity가 무한히 클 수는 없기 때문입니다. (만약 그런 경우가 존재한다면, 모든 device의 모든 data가 서로 관련이 없다는 것인데, 이 경우에는 어차피 정상적으로 학습되지 않을 것입니다. 현실에서는 서로 관련이 있는 data가 device 별로 어느 정도 존재합니다.)

 추가적으로, 저자들은 $\epsilon$을 tight하게 잡을 필요 없다고 주장하고 있습니다. $\epsilon$를 과도하게 작게 잡는다는 것은 곧 $B(w)$가 $1$에 충분히 가까워진다는 것이며, 이는 곧 $w$가 stationary한 상태에 준한다는 것입니다. 따라서,  이는 overfitting 문제를 야기할 수 있고, 그러므로 적당한 $\epsilon$을 잡기만 하면 된다는 것입니다.

 

5. Non-convex case의 증명

 

 이제 위 assumption을 가지고 Non-convex case에 대해서 FedProx의 convergence를 증명해봅시다.

 

$\text{Theorem 4}$ (Non-convex FedProx convergence: $B$-local dissimilarity)

 $\text{Assumption 1}$이 성립한다고 가정하자. 또, 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$가 $$\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) > 0$$를 만족한다면, FedProx의 $t$번째 iteration에서 global objective function $f$의 expected decrease는 $\mathbb{E}_{S_t} [f(w^{t + 1})] \leq f(w^t) - \rho ||\nabla f(w^t)||^2$를 만족한다. (여기에서, $S_t$는 $t$번째 iteration에 참여하는 device들의 집합이다.)

 

$\text{Proof}$

 $\nabla F(w_k^{t + 1}) + \mu (w_k^{t + 1} - w^t) = ||\nabla h(w_k^{t + 1}; w^t)|| =: e_k^{t + 1}$로 정의하자. $(1)$ 그러면, $\gamma$-inexactness의 정의에 의하여 $e_k^{t + 1} \leq \gamma ||\nabla h(w^t; w^t)|| = \gamma ||\nabla F_k(w^t)||$이다. $(2)$ 다음으로, $\bar{w}^{t + 1} := \mathbb{E}_k [w_k^{t + 1}]$로 정의하자. 그러면, $(1)$로부터

 

$$\mu (w_k^{t + 1} - w^t) = e_k^{t + 1} - \nabla F(w_k^{t + 1})$$

$$w_k^{t + 1} - w^t = \frac {1} {{\mu}} e_k^{t + 1}  - \frac {1} {\mu} \nabla F(w_k^{t + 1})$$

$$\mathbb{E}_k[w_k^{t + 1}] - \mathbb{E}_k[w^t] = \frac {1} {{\mu}} \mathbb{E}_k[e_k^{t + 1}]  - \frac {1} {\mu} \mathbb{E}_k[\nabla F(w_k^{t + 1})]$$

 

임을 알 수 있다. $w^t$는 $k$와 관련 없으므로 $\mathbb{E}_k[w^t] = w^t$이고, 방금 정의한 대로 $\mathbb{E}_k[w_k^{t + 1}] = \bar{w}^{t + 1}$이다. 그러므로 식을 정리하면 $\bar{w}^{t + 1} - w^t = - \frac {1} {\mu} \mathbb{E}_k[\nabla F(w_k^{t + 1})] + \frac {1} {{\mu}} \mathbb{E}_k[e_k^{t + 1}]$이다. $(3)$

 

---

$\text{Claim}$: $h_k$는 $\bar{\mu}$-strongly convex하다. $(4)$

$\text{Proof}$

 임의의 model $w$에 대하여, $h_k(w) = F_k(w) + \frac {\mu} {2} ||w - w^t||^2 = F_k(w_1) + \frac {L_-} {2} ||w - w^t||^2 + \frac {\bar{\mu}} {2} ||w - w^t||^2$이다. 여기에서, $\frac {\bar{\mu}} {2} ||w - w^t||^2$는 그 자체로 $\bar{\mu}$-strongly convex하다. 또한, 가정에 의하여 $\nabla^2 F_k \succeq -L_- \textbf{I}$이기 때문에,  $\nabla^2 F_k  + L_- \textbf{I}$는 positive semi-definite하고, 따라서 $F_k(w_1) + \frac {L_-} {2} ||w - w^t||^2$는 convex하다. 그러므로 이 둘의 합인 $h_k(w)$는 $\bar{\mu}$-strongly convex하다. $\square$

---

 

 이제, $\hat{w}_k^{t + 1} := argmin_w \; h_k(w; w^t)$로 정의하자. 그러면 $(2)$와 $(4)$에 의하여, 우리는 $||\hat{w_k}^{t + 1} - w^{t + 1}|| \leq \frac {\gamma} {\bar{\mu}} ||\nabla F_k(w^t)||$, $||\hat{w_k}^{t + 1} - w^t|| \leq \frac {1} {\bar{\mu}} ||\nabla F_k(w^t)||$임을 알 수 있고, Triangle Inequality에 의하여 $||w_k^{t + 1} - w^t|| \leq \frac{1 + \gamma} {\mu} ||\nabla F_k(w_t)||$이다. $(5)$ 그리고 결론적으로 다음과 같은 식을 얻을 수 있다. (여기에서, 마지막 부등식에 $\text{Assumption 1}$이 사용되었다.) $(6)$

 

\begin{align*} ||w_k^{t + 1} - w^t|| &\leq \mathbb{E}_k[||w_k^{t + 1} - w^t||] \leq \frac{1 + \gamma} {\mu} \mathbb{E}_k [||\nabla F_k(w_t)||] \\&\leq \frac{1 + \gamma} {\mu} \sqrt{\mathbb{E}_k [||\nabla F_k(w_t)||^2}] \leq \frac {B(1 + \gamma)} {\bar{\mu}} ||\nabla f(w^t)|| \; \end{align*}

 

 다음으로, $\bar{w}^{t + 1} - w^t = - \frac {1} {\mu} (\nabla f(w^t) + M_{t + 1})$을 만족하는 $M_{t + 1}$을 정의하자. $(3)$을 이용하여 다시 정리하면, $M_{t + 1} = \mathbb{E}_k [\nabla F_k(w_k^{t + 1}) - \nabla F_k(w^t) - e_k^{t + 1}]$이고, 따라서 다음과 같이 $||M_{t + 1}||$가 bounded되는 것을 확인할 수 있다. $(7)$

 

\begin{align*} ||M_{t + 1}|| &= \mathbb{E} [\nabla F_k(w_k^{t + 1}) - \nabla F_k(w^t) - e_k^{t + 1}] \\&\leq \mathbb{E} [L ||w_k^{t + 1} - w^t|| + ||e_k^{t + 1}||] \; (\because L \text{-Lipschitz smoothness of } f) \\&\leq (\frac {L(1 + \gamma)} {\hat{\mu}} + \gamma) \times \mathbb{E}_k [||\nabla F_k (w^t)||] \; (\because (2) \text{, } (5)) \\&\leq (\frac {L(1 + \gamma)} {\hat{\mu}} + \gamma) \times B ||\nabla f(w^t)|| \; (\because \text{Assumption 1}) \end{align*}

 

 이제 마지막으로, $f$의 local Lipschitz continuity를 이용하여 $f(w^{t + 1}) \leq f(\hat{w}^{t + 1}) + L_0 ||w^{t + 1} - \hat{w}^{t + 1}||$을 얻을 수 있다. 여기에서 $L_0$는 $f$는 local Lipschitz continuity constant인데, Lipschitz continuity constant인 $L$과의 차이에 의해서 다음과 같이 bounded된다. $(8)$

 

\begin{align*} L_0 &\leq ||\nabla f(w^t)|| + L \ max(||\bar{w}^{t + 1} - w^t||, ||w^{t+1} - w^t||) \\&\leq ||\nabla f(w^t)|| + L (||\bar{w}^{t + 1} - w^t|| + ||w^{t+1} - w^t||) \end{align*}

 

 

 더 나아가, 양변에 $\mathbb{E}_{S_t} [\cdot]$를 취해주면 $Q_t := \mathbb{E}_{S_t} [L_0 ||w^{t+1} - \bar{w}^{t+1}||]$일 때 $\mathbb{E}_{S_t} [f(w^{t + 1})] \leq f(\bar{w}^{t + 1}) + Q_t$이다. $(9)$ 여기에서 $Q_t$는 다음과 같이 bounded된다. $(10)$

 

\begin{align*} Q_t &= \mathbb{E}_{S_t} [L_0 ||w^{t+1} - \bar{w}^{t+1}||] \\&\leq \mathbb{E}_{S_t} [(||\nabla f(w^t)|| + L (||\bar{w}^{t+1} - w^t|| + ||w^{t+1} - w^t||) \times ||w^{t+1} - \bar{w}^{t+1}||] \; (\because (8)) \\&\leq (||\nabla f(w^t)|| + L ||\bar{w}^{t+1} - w^t||) \mathbb{E}_{S_t} [||w^{t+1} - \bar{w}^{t+1}||] + L \mathbb{E}_{S_t} [||w^{t+1} - w^t|| \cdot ||w^{t + 1} - \bar{w}^{t+1}||]   \\&\leq (||\nabla f(w^t)|| + 2L ||\bar{w}^{t+1} - w^t||) \mathbb{E}_{S_t} [||w^{t+1} - \bar{w}^{t+1}||] + L \mathbb{E}_{S_t} [||w^{t + 1} - \bar{w}^{t+1}||^2] \end{align*}

 

 여기에서 $\mathbb{E}_{S_t} [||w^{t + 1} - \bar{w}^{t+1}||] \leq \sqrt{\mathbb{E}_{S_t} [||w^{t + 1} - \bar{w}^{t+1}||^2]}$이고,

 

\begin{align*} \mathbb{E}_{S_t} [||w^{t+1} - \bar{w}^{t+1}||^2] &\leq \frac {1} {K} \mathbb{E}_{k} [||w^{t+1} - \bar{w}^{t+1}||^2] \; (\because \text{random sampling from } K \text{devices}) \\&= \frac {1} {K} \mathbb{E}_{k} [||w^{t+1} - \mathbb{E}_k[w_k^{t+1}]||^2] \\&\leq \frac {2} {K} \mathbb{E}_k  [||w_k^{t+1} - w^t|| ^ 2] \\&\leq \frac {2} {K} \frac {(1 + \gamma)^2} {\bar{\mu}} \mathbb{E}_k [||\nabla F_k(w^t)||^2] \; (\because (6)) \\&\leq \frac {2B^2} {K} \frac {(1 + \gamma)^2} {\bar{\mu}^2} ||\nabla f(w^t)||^2 \; (\because \text{Assumption 1}) \end{align*}

 

이므로, $(10)$를 다시 정리하면 다음과 같은 결과를 얻을 수 있다. $(11)$

 

\begin{align*} Q_t &\leq (||\nabla f(w^t)|| + 2L ||\bar{w}^{t+1} - w^t||) \mathbb{E}_{S_t} [||w^{t+1} - \bar{w}^{t+1}||] + L \mathbb{E}_{S_t} [||w^{t+1} - \bar{w}^{t+1}||^2] \\&\leq (||\nabla f(w^t)|| + 2L ||\bar{w}^{t+1} - w^t||) \sqrt{\mathbb{E}_{S_t} [||w^{t+1} - \bar{w}^{t+1}||^2]} + L \frac {2B^2} {K} \frac {(1 + \gamma)^2} {\bar{\mu}^2} ||\nabla f(w^t)||^2 \\&\leq (||\nabla f(w^t)|| + 2L ||\bar{w}^{t+1} - w^t||) \sqrt{\frac {2B^2} {K} \frac {(1 + \gamma)^2} {\bar{\mu}^2} ||\nabla f(w^t)||^2} + \frac {2LB^2} {K} \frac {(1 + \gamma)^2} {\bar{\mu}^2} ||\nabla f(w^t)||^2 \\&\leq ||\nabla f(w^t)||(1 + 2L \frac {B(1+\gamma)} {\bar{\mu}}) \sqrt{\frac {2B^2} {K} \frac {(1 + \gamma)^2} {\bar{\mu}^2} ||\nabla f(w^t)||^2} + \frac {2LB^2} {K} \frac {(1 + \gamma)^2} {\bar{\mu}^2} ||\nabla f(w^t)||^2 \; (\because (6)) \\&= \left((1 + \frac {2LB(1+\gamma)} {\bar{\mu}}) \sqrt{\frac {2B^2} {K} \frac {(1 + \gamma)^2} {\bar{\mu}^2}} + \frac {2LB^2} {K} \frac {(1 + \gamma)^2} {\bar{\mu}^2} \right) ||\nabla f(w^t)||^2 \\&= \left((1 + \frac {2LB(1+\gamma)} {\bar{\mu}}) \frac {\sqrt{2}B(1+\gamma)} {\bar{\mu} \sqrt{K}} + \frac {2LB^2(1+\gamma)^2} {\bar{\mu}^2 K} \right) ||\nabla f(w^t)||^2 \\&= \left(\frac {\sqrt{2}B(1+\gamma)} {\bar{\mu} \sqrt{K}} + \frac {2 \sqrt{2} LB^2 (1+\gamma)^2} {\bar{\mu}^2 \sqrt{K}} + \frac {2LB^2(1+\gamma)^2} {\bar{\mu}^2 K} \right) ||\nabla f(w^t)||^2 \\&= \left(\frac {B(1+\gamma) \sqrt{2}} {\bar{\mu} \sqrt{K}} + \frac {LB^2 (1+\gamma)^2} {\bar{\mu}^2 K} (2\sqrt{2K} + 2) \right) ||\nabla f(w^t)||^2     \end{align*}

 

 따라서, 우리는 다음과 같은 결론을 내릴 수 있다.

 

\begin{align*} \mathbb{E}_{S_t} [f(w^{t + 1})] &\leq f(\bar{w}^{t + 1}) + Q_t \; (\because (9)) \\&\leq f(w^t) + \langle \nabla f(w)^t, \;  \bar{w}^{t+1} - w^t \rangle + \frac {L} {2} ||\bar{w}^{t+1} - w^t||^2 + Q_t \; (\because L\text{-Lipschitz smoothness of }f) \\&\leq f(w^t) + \langle \nabla f(w)^t, \;  \bar{w}^{t+1} - w^t \rangle + \frac {L(1+\gamma)^2 B^2} {2 \bar{\mu}^w} ||\nabla f(w^t)||^2 + Q_t \; (\because (6)) \\&= f(w^t) - \frac {1} {\mu} ||\nabla f(w^t)||^2 - \frac {1} {\mu} \langle \nabla f(w^t), \; M_{t+1} \rangle  + \frac {L(1+\gamma)^2 B^2} {2 \bar{\mu}^2} ||\nabla f(w^t)||^2 + Q_t \; (\because (3)) \\&\leq f(w^t) - \left( \frac {1 - \gamma B} {\mu} - \frac {LB(1+\gamma)} {\bar{\mu}\mu} - \frac {L(1+\gamma)^2 B^2} {2 \bar{\mu}^2}   \right) \times ||\nabla f(w^t)||^2 + Q_t \; (\because (7)) \\&\leq f(w^t) - \left( \frac {1} {\mu} - \frac {\gamma B} {\mu} - \frac {LB(1+\gamma)} {\bar{\mu}\mu} - \frac {L(1+\gamma)^2 B^2} {2 \bar{\mu}^2}   \right) \times ||\nabla f(w^t)||^2 \\& \quad + \left(\frac {B(1+\gamma) \sqrt{2}} {\bar{\mu} \sqrt{K}} + \frac {LB^2 (1+\gamma)^2} {\bar{\mu}^2 K} (2\sqrt{2K} + 2) \right) ||\nabla f(w^t)||^2 \; (\because (11)) \\&= f(w^t) - \rho ||f(w^t)||^2 \end{align*}

 

 이것으로 증명을 마친다. $\square$

 

 이전 포스트에서 proximal term이 또 다른 의미를 가지고 있다고 이야기하였는데, 그것은 수식 계산 상의 이점이라는 것을 해당 증명을 통해서 알 수 있습니다. 마지막으로, 증명 과정에서 한 가지 확인하지 않은 것이 있는데, 바로 $\rho > 0$임을 확신할 수 없다는 것입니다. 이를 위한 sufficient condition은 $\gamma B <1$, $\frac {B} {\sqrt{K}} < 1$이며, 이러한 제약식은 dissimilarity($B$)와 algorithm의 parameter($\gamma$, $K$) 간의 trade-off를 설명해줍니다.

 

 원래의 계획은 모든 증명을 포스트 하나로 끝내는 것이었는데, 생각보다 증명 과정이 길어졌네요. 다음 포스트에서는 FedProx의 convergence rate를 확인해본 후, 이번 포스트에서 증명한 방법을 활용하여 convex case의 convergence에 관해서, 그리고 해당 논문의 핵심이기도 한 $\gamma$가 달라지는 case의 convergence에 관해서 알아보겠습니다. (다음 글 보기)