[ICLR 2020] Convergence of FedAvg - (3)

논문 제목: On the Convergence of FedAvg on Non-IID Data

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

 

 지난 포스트에 이어서 증명을 진행하겠습니다. (이전 글 보기)

 

3. Key Lemmas for $\text{Theorem 1}$ - 이어서

 

$\text{Lemma 2}$ [Bounding the variance]

 $\text{Assumption 3}$이 성립한다고 가정할 때, $\mathbb{E} [||g_t - \bar{g_t}||^2] \leq \sum_{k=1}^N p_k^2 \sigma_k^2$이다.

 

 $\text{Proof}$

 정의 상 $g_t = \sum_{k=1}^N p_k \nabla F_k(w_t^k, \xi_t^k)$, $\bar{g}_t = \sum_{k=1}^N p_k \nabla F_k(w_t^k)$이므로,

\begin{align*} \mathbb{E} [||g_t - \bar{g_t}||^2] &= \mathbb{E} \left[ \left| \left| \sum_{k=1}^N p_k \nabla F_k(w_t^k, \xi_t^k) - \sum_{k=1}^N p_k \nabla F_k(w_t^k) \right| \right|^2 \right] \\&= \mathbb{E} \left[ \left| \left| \sum_{k=1}^N p_k(\nabla F_k(w_t^k, \xi_t^k) - \nabla F_k(w_t^k)) \right| \right|^2 \right] \\&= \sum_{k=1}^N p_k^2 \mathbb{E} \left[ \left| \left| (\nabla F_k(w_t^k, \xi_t^k) - \nabla F_k(w_t^k)) \right| \right|^2 \right] \\&\leq \sum_{k=1}^N p_k^2 \sigma_k^2 \; (\because \text{Assumption 3}) \end{align*}

를 만족시키는 $\sigma_k$가 존재한다. $\square$

 

$\text{Lemma 3}$ [Bounding the divergence of $\{w_t^k\}$]

 $\text{Assumption 4}$가 성립한다고 가정하자. 만약 $\eta_t$가 non-decreasing하고 모든 $t \geq 0$에 대해서 $\eta_t \leq 2 \eta_{t + E}$라면, $\mathbb{E} \left[ \sum_{k=1}^N p_k ||\bar{w_t} - w_t^k||^2 \right] \leq 4\eta_t^2 (E - 1)^2 G^2$이다.

 

 $\text{Proof}$

 매 $E$ step마다 communication을 진행하기 때문에, 임의의 $t \geq 0$에 대하여 다음을 만족하는 $t_0 \leq t$가 존재한다.

•  $t - t_0 \leq E - 1$
• 모든 $k = 1, 2, \cdots, N$에 대해서 $w_{t_0}^k = \bar{w}_{t_0}$
• $\eta_{t_0} \leq 2 \eta_t$

 따라서,

\begin{align*} \mathbb{E} \left[ \sum_{k=1}^N p_k ||\bar{w_t} - w_t^k||^2 \right] &= \mathbb{E} \left[ \sum_{k=1}^N p_k || w_t^k - \bar{w_t}||^2 \right] =  \mathbb{E} \left[ \sum_{k=1}^N p_k || w_t^k - \bar{w}_{t_0} + \bar{w}_{t_0} - \bar{w_t}||^2 \right] \\&= \mathbb{E} \left[ \sum_{k=1}^N p_k || (w_t^k - \bar{w}_{t_0}) - (\bar{w_t} - \bar{w}_{t_0})||^2 \right]  \\&\leq \mathbb{E} \left[ \sum_{k=1}^N p_k || w_t^k - \bar{w}_{t_0} ||^2 \right] \; (\because \mathbb{E} \left[ ||X - \mathbb{E} [X]||^2 \right] \leq \mathbb{E} \left[ ||X||^2 \right]) \\&= \sum_{k=1}^N p_k \mathbb{E} \left[ \left| \left| \sum_{t=t_0}^{t-1} \eta_t \nabla F_k (w_t^k, \xi_t^k) \right| \right|^2 \right] \\&= \sum_{k=1}^N p_k \mathbb{E} \left[ (t - t_0) \times \frac {\left| \left| \sum_{t=t_0}^{t-1} \eta_t \nabla F_k (w_t^k, \xi_t^k) \right| \right|^2} {t - t_0} \right] \\&\leq \sum_{k=1}^N p_k \mathbb{E} \left[ (t - t_0) \times \sum_{t=t_0}^{t-1} \left| \left| \eta_t \nabla F_k (w_t^k, \xi_t^k) \right| \right|^2 \right] \; (\because \text{Jensen's Inequality}) \\&\leq \sum_{k=1}^N p_k \mathbb{E} \left[ (E - 1) \times \sum_{t=t_0}^{t-1} \eta_t^2 \left| \left| \nabla F_k (w_t^k, \xi_t^k) \right| \right|^2 \right] \; (\because t - t_0 \leq E - 1) \\&\leq \sum_{k=1}^N p_k \mathbb{E} \left[ (E - 1) \times \sum_{t=t_0}^{t-1} \eta_{t_0}^2 \left| \left| \nabla F_k (w_t^k, \xi_t^k) \right| \right|^2 \right] \; (\because \eta_t \text{ is nondecreasing}) \\&\leq \sum_{k=1}^N p_k \eta_{t_0}^2 (E - 1) \sum_{t=t_0}^{t-1} G^2 \; (\because \text{Assumption 4}) \\&= \sum_{k=1}^N p_k \eta_{t_0}^2 (E - 1) (t - t_0) G^2 \\&\leq \sum_{k=1}^N p_k \eta_{t_0}^2 (E - 1)^2 G^2  \; (\because t - t_0 \leq E - 1) \\&\leq 4 \eta_t^2 (E - 1)^2 G^2  \; (\because \eta_{t_0} \leq 2 \eta_t) \end{align*}

를 만족시키는 $G$가 존재한다. $\square$

 

4. Full Device Participation Case

 

 앞서 확인한 $\text{Lemma 1 ~ 3}$를 이용하여 full device participation case의 convergence를 증명하겠습니다.

 

$\text{Theorem 1}$

 $\text{Assumption 1 ~ 4}$가 성립한다고 가정하자. 그리고 $\kappa := \frac {L} {\mu}$, $\gamma := max  \{ 8 \kappa, E \}$, $\eta_t := \frac {2} {\mu (\gamma + t)}$로 지정하자. 그러면 full device participation case의 FedAvg는 다음을 만족한다. (단, $L, \mu, \sigma_k, G$는 $\text{Assumption 1 ~ 4}$에서 정의된 것과 동일하며, $B:= \sum_{k=1}^N p_k^2 \sigma_k^2 + 6L \Gamma + 8 (E - 1)^2 G^2$이다.)

$$\mathbb{E} [F (w_T)] - F^* \leq \frac {\kappa} {\gamma + T - 1} \left( \frac {2B} {\mu} + \frac {\mu} {2} \mathbb{E} ||w_1 - w^*||^2 \right)$$

$\text{Proof}$

$\Delta_t := \mathbb{E} ||\bar{w}_t - w^*||^2$로 정의하자. 이때, \begin{align*} \Delta_{t+1} &= \mathbb{E} ||\bar{w}_{t+1} - w^*||^2 = \mathbb{E} ||\bar{v}_{t+1} - w^*||^2 \; (\because \bar{w}_{t+1} = \bar{v}_{t+1}) \\&\leq (1 - \eta_t \mu) \mathbb{E} [||\bar{w} - w^*||^2] + \eta_t^2 \mathbb{E} [||g_t - \bar{g}_t||^2] + 6L \eta_t^2 \Gamma + 2 \mathbb{E} \sum_{k = 1}^N p_k ||\bar{w}_t - w_t^k||^2 \; (\because \text{Lemma 1}) \\&= (1 - \eta_t \mu) \Delta_t + \eta_t^2 \mathbb{E} [||g_t - \bar{g}_t||^2] + 6L \eta_t^2 \Gamma + 2 \mathbb{E} \sum_{k = 1}^N p_k ||\bar{w}_t - w_t^k||^2 \; (\because \Delta_t := \mathbb{E} ||\bar{w}_t - w^*||^2) \\&\leq (1 - \eta_t \mu) \Delta_t + \eta_t^2 \sum_{k=1}^N p_k^2 \sigma_k^2 + 6L \eta_t^2 \Gamma + 2 \mathbb{E} \sum_{k = 1}^N p_k ||\bar{w}_t - w_t^k||^2 \; (\because \text{Lemma 2}) \\&\leq (1 - \eta_t \mu) \Delta_t + \eta_t^2 \sum_{k=1}^N p_k^2 \sigma_k^2 + 6L \eta_t^2 \Gamma + 8 \eta_t^2 (E - 1)^2 G^2 \; (\because \text{Lemma 3}) \\&= (1 - \eta_t \mu) \Delta_t + \eta_t^2 \left( \sum_{k=1}^N p_k^2 \sigma_k^2 + 6L \Gamma + 8 (E - 1)^2 G^2 \right)  \\&= (1 - \eta_t \mu) \Delta_t + \eta_t^2 B     \end{align*}

이다. learning rate를 decaying하기 위해서, $\eta_t := \frac {\beta} {t + \gamma}$로 지정하자. (이때, $\beta > \frac {1} {\mu}$이며, $\gamma > 0$이다.) $\text{Lemma 1}$에서의 가정에 의하여 $\eta_t \leq \frac {1} {4L}$이므로, $\eta_1 = \frac {\beta} {1 + \gamma} \leq \min \{ \frac {1} \mu, \frac {1} {4L} \} = \frac {1} {4L}$이다. $(1)$ ($\text{Assumption 1, 2}$로부터 $\mu \leq L$임을 알 수 있다.) 또한, $\text{Lemma 3}$에서의 가정에 의하여, $\eta_t \leq 2 \eta_{t + E}$이다. $(2)$

---

$\text{Claim}$: $v := \max \{ \frac {\beta^2 B} {\beta \mu - 1}, (\gamma + 1) \Delta_1 \}$일 때, $\Delta_t \leq \frac {v} {\gamma + t}$이다. $(3)$

$\text{Proof}$

 $t$에 대하여 수학적 귀납법을 사용하여 증명하자.

 

$\text{(i)} \; t = 1$ 

 우선, $v = \max \{ \frac {\beta^2 B} {\beta \mu - 1}, (\gamma + 1) \Delta_1 \} = (\gamma + 1) \Delta_1$인 경우, $\frac {v} {\gamma + 1} = \frac {(\gamma + 1) \Delta_1} {\gamma + 1} = \Delta_1$이므로 자명하다. 다음으로, $v = \max \{ \frac {\beta^2 B} {\beta \mu - 1}, (\gamma + 1) \Delta_1 \} = \frac {\beta^2 B} {\beta \mu - 1}$인 경우, $\eta_t = \frac {\beta} {t + \gamma}$이므로 자명하다. 따라서, $\Delta_1 \leq \frac {v} {\gamma + 1}$이다.

 

$\text{(ii)} \; t$일 때 위 식이 성립한다고 가정하자. 그러면

\begin{align*} \Delta_{t+1} &\leq (1 - \eta_t \mu) \Delta_t + \eta_t^2 B = (1 - \frac {\beta} {t + \gamma} \mu) \Delta_t + (\frac {\beta} {t + \gamma})^2 B \; (\because \eta_t = \frac {\beta} {t + \gamma}) \\&\leq (1 - \frac {\beta \mu} {t + \gamma} ) \times \frac {v} {\gamma + t} + \frac {\beta^2 B} {(t + \gamma)^2} \; (\because \Delta_t \leq \frac {v} {\gamma + t}) \\&= \frac {t + \gamma - \beta \mu} {(t + \gamma)^2} v + \frac {\beta^2 B} {(t + \gamma)^2}  \\&= \frac {t + \gamma - 1 + 1 - \beta \mu} {(t + \gamma)^2} v + \frac {\beta^2 B} {(t + \gamma)^2} \\&= \frac {t + \gamma - 1} {(t + \gamma)^2} v + \frac {\beta^2 B - (\beta \mu - 1) v} {(t + \gamma)^2} \\&\leq \frac {t + \gamma - 1} {(t + \gamma)^2} v \; (\because \frac {\beta^2 B - (\beta \mu - 1) v} {(t + \gamma)^2} \leq 0) \\&\leq \frac {t + \gamma} {(t + \gamma + 1) (t + \gamma)} v = \frac {v} {t + \gamma + 1}      \end{align*}

이므로, 위 식이 $t + 1$일 때에도 성립한다는 것을 알 수 있다. 따라서, 모든 $t \in \mathbb{N}$에 대해서 해당 식이 성립한다. $\square$

---

 이때, $\text{Assumption 1}$에 의하여 $F_k(\bar{w}_t) - F_k^* = F_k(\bar{w}_t) - F_k(w^*) \leq (\bar{w}_t - w^*)^T \nabla F_k (w^*) + \frac {L} {2} ||\bar{w}_t - w^*||^2$이므로, 양변에 expectation을 취해주면 다음과 같이 $\mathbb{E} [F (\bar{w}_t)] - F^*$가 bound된다는 것을 알 수 있다. (이때, 마지막 부등호는 $(3)$에 의한 것이다.)

$$\mathbb{E} [F (\bar{w}_t)] - F^* =  \mathbb{E} [F (\bar{w}_t)] - \mathbb{E} [F_k^*] \leq \frac {L} {2} \mathbb{E} \left[ ||\bar{w}_t - w^*||^2 \right] = \frac {L} {2} \Delta_t \leq \frac{L} {2} \times \frac {v} {\gamma + t}$$

 여기에서, 해당 증명의 전제에서 주어진 바와 같이 $\kappa := \frac {L} {\mu}$, $\gamma := max  \{ 8 \kappa, E \}$, $\eta_t := \frac {2} {\mu (\gamma + t)}$로 지정할 경우(이때, $\beta = \frac {2} {\mu}$), $v = \max \{ \frac {\beta^2 B} {\beta \mu - 1}, (\gamma + 1) \Delta_1 \} \leq \frac {\beta^2 B} {\beta \mu - 1} + (\gamma + 1) \Delta_1 \leq \beta^2 B + (\gamma + 1) \Delta_1 = \frac {4B} {\mu^2} + (\gamma + 1) \Delta_1$이기 때문에,

\begin{align*} \mathbb{E} [F (\bar{w}_t)] - F^* &\leq \frac{L} {2} \times \frac {v} {\gamma + t} \leq \frac {L} {2} \times \frac {\frac {4B} {\mu^2} + (\gamma + 1) \Delta_1} {\gamma + t} \\&= \frac {\kappa \mu} {2} \times \frac {\frac {4B} {\mu^2} + (\gamma + 1) \Delta_1} {\gamma + t} \; (\because \kappa = \frac {L} {\mu}) \\&= \frac {\kappa} {\gamma + t} \times \frac {\frac {4B} {\mu} + \mu (\gamma + 1) \Delta_1} {2} \\&= \frac {\kappa} {\gamma + t} \left( \frac {2B} {\mu} + \frac{\mu ( \gamma + 1)} {2} \Delta_1 \right)  \end{align*}

로 $\mathbb{E} [F (\bar{w}_t)] - F^*$가 bound된다. $\square$

 

 다음 포스트에서는 이번에 증명한 $\text{Theorem 1}$에 기반하여 partial device participation case를 분석하도록 하겠습니다.