[ICLR 2020] Convergence of FedAvg - (5)

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

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

 

 지난 포스트에 이어서 partial device participation case의 convergence를 증명하도록 하겠습니다. (이전 글 보기) $C$ term만 다르기 때문에, $\text{Scheme I}$과 $\text{II}$ 구분 없이 동시에 증명합니다.

 

7. Partial Device Participation Case

 

$\text{Theorem 2}$

 $\text{Assumtion 1 ~ 4}$가 성립하고, $\text{Scheme I}$을 사용한다고 가정하자. (즉, $\text{Assumption 5}$가 성립한다고 가정하자.) 그러면 partial device participation case의 FedAvg는 다음을 만족한다. (단, $L, \mu, \sigma_k, G$는 $\text{Assumption 1 ~ 4}$에서 정의된 것과 동일하며, $\kappa. \gamma, \eta_t, B$는 $\text{Theorem 1}$에서 정의된 것과 동일하다. 그리고 $C:= \frac {4} {K} E^2 G^2$이다.)

 

$\text{Theorem 3}$

 $\text{Assumtion 1 ~ 4}$가 성립하고, $\text{Scheme II}$를 사용한다고 가정하자. (즉, $\text{Assumption 6}$가 성립한다고 가정하자.) 그러면 partial device participation case의 FedAvg는 다음을 만족한다. ((단, $L, \mu, \sigma_k, G$는 $\text{Assumption 1 ~ 4}$에서 정의된 것과 동일하며, $\kappa. \gamma, \eta_t, B$는 $\text{Theorem 1}$에서 정의된 것과 동일하다. 그리고 $C:= \frac {N - K} {N - 1} \frac {4} {K} E^2 G^2$이다.)

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

$\text{Proof}$

$$||\bar{w}_{t+1} - w^*||^2 = ||\bar{w}_{t+1} - \bar{v}_{t+1} + \bar{v}_{t+1} - w^*||^2 = \underbrace{||\bar{w}_{t+1} - \bar{v}_{t+1}||^2}_{A_1} + \underbrace{||\bar{v}_{t+1} - w^*||^2}_{A_2} + \underbrace{2 \langle \bar{w}_{t+1} - \bar{v}_{t+1}, \; \bar{v}_{t+1} - w^* \rangle}_{A_3}$$

 에서, $A_3$는 $\mathcal{S}_t$에 대한 expectation을 취할 경우 $0$이 되기 때문에, 우리는 $A_1$과 $A_2$만 정리하면 된다. 우선, $t+1 \notin \mathcal{I}_E$일 경우, $\bar{w}_{t+1} = \bar{v}_{t+1}$이므로 $A_1 = 0$이다. 따라서, 다음이 성립한다. $(1)$

$$\begin{align*} \mathbb{E} ||\bar{w}_{t+1} - w^*||^2 &= \mathbb{E} [A_2] = \mathbb{E} [||\bar{v}_{t+1} - w^*||^2] \\&\leq (1 - \eta_t \mu) \mathbb{E} [||\bar{w}_t - 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}) \\&\leq (1 - \eta_t \mu) \mathbb{E} [||\bar{w}_t - w^*||^2] + \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) \mathbb{E} [||\bar{w}_t - w^*||^2] + \eta_t^2 \sum_{k=1}^N p_k^2 \sigma_k^2 + 6L \eta_t^2 \Gamma + 2 \times 4\eta_t^2 (E - 1)^2 G^2 \; (\because \text{Lemma 3}) \\&= (1 - \eta_t \mu) \mathbb{E} [||\bar{w}_t - w^*||^2] + \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) \mathbb{E} [||\bar{w}_t - w^*||^2] + \eta_t^2 B \end{align*}$$

 다음으로, $t+1 \in \mathcal{I}_E$일 경우, $\text{Lemma 5}$에 의하여 $\mathbb{E}_{\mathcal{S}_t} [A_1] \leq ||\bar{w}_{t+1} - \bar{v}_{t+1}||^2 \leq \eta_t^2 C$이다. $(2)$ 따라서, 다음이 성립한다.

$$\begin{align*} \mathbb{E} ||\bar{w}_{t+1} - w^*||^2 &= \mathbb{E} [A_1] + \mathbb{E} [A_2] = \mathbb{E} [||\bar{w}_{t+1} - \bar{v}_{t+1}||^2] + \mathbb{E} [||\bar{v}_{t+1} - w^*||^2] \\&\leq \mathbb{E} [||\bar{w}_{t+1} - \bar{v}_{t+1}||^2] + (1 - \eta_t \mu) \mathbb{E} [||\bar{w}_t - w^*||^2] + \eta_t^2 B \; (\because (1)) \\&\leq \eta_t^2 C + (1 - \eta_t \mu) \mathbb{E} [||\bar{w}_t - w^*||^2] + \eta_t^2 B \; (\because (2)) \\&= (1 - \eta_t \mu) \mathbb{E} [||\bar{w}_t - w^*||^2] + \eta_t^2 (B+C) \end{align*}$$

 $\text{Theorem 1}$과의 차이점은 $C$ term 뿐이므로, $v = \max \{ \frac {\beta^2 B} {\beta \mu - 1}, (\gamma + 1) \Delta_1 \} = \max \{ \frac {\beta^2 B} {\beta \mu - 1}, (\gamma + 1) \mathbb{E} ||\bar{w}_1 - w^*||^2 \}$로 정의하였던 것을 $v = \max \{ \frac {\beta^2 (B+C)} {\beta \mu - 1}, (\gamma + 1) \mathbb{E} ||\bar{w}_1 - w^*||^2 \}$로 바꾸면 동일하게 성립한다. $\square$

 

8. 어떠한 $\text{Scheme}$을 사용하여야 할까?

 

 $\text{Scheme II}$에는 balanced dataset이라는 비현실적인 가정이 포함되어 있기 때문에, 이론적으로는 $\text{Scheme I}$을 사용하는 것이 좋습니다. 특히 $p_1, p_2, \cdots, p_N$이 매우 non-uniform할 경우에 이러한 replacement를 허용하는 sampling 기법이 더욱 효과적으로 작동합니다. 문제는, (Cross-Silo case라면 이러한 걱정이 덜 하겠지만) 일반적으로 연합학습 과정에는 불특정 다수의 기기가 참여하기 때문에, 우리가 $N$개의 device 중 $K$개를 언제나 임의로 sampling할 수 없다는 점입니다. 즉, straggler를 고려하여야 한다는 것이죠. 따라서, $K$개를 확정적으로 sampling할 수 없는 경우에는, 그저 먼저 반환되는 $K$개의 parameter들을 가지고 update를 하는 것이 현실적인 방법일 수 있습니다. 이러한 상황에서는 $K$개의 device가 uniformly distributed되어 있다고 가정하고 $\text{Theorem 3}$을 이용하여 수렴성을 보장하면 됩니다. 다시 말해, 비록 다소 비현실적인 가정을 갖고 있더라도, $\text{Scheme II}$가 $\text{Scheme I}$보다 더 현실에 들어맞는 이야기일 수도 있다는 의미입니다.

 그렇다면, balanced dataset이라는 가정을 조금 더 완화할 수 있는 방법은 없을까요? $\widetilde{F}_k := p_k N F_k(w)$로 정의할 때, global objective function $F(w)$를 다음과 같이 정리할 수 있고, 따라서 uniform distribution을 보장할 수 있습니다.

$$F(w) = \sum_{k=1}^N p_k F_k(w) = \frac{1} {N} \sum_{k=1}^N \widetilde{F}_k (w)$$

그리고 이때에도 $\text{Theorem 3}$는 동일하게 성립하는데, $F_k$를 $p_k N F_k$로 scaling한 만큼 $L, \mu, \sigma_k, G$도 scaling을 해주어야 합니다. 정확히는, $\nu := N \cdot \max_k p_k$, $\varsigma := N \cdot \min_k p_k$일 때, $L, \mu, \sigma_k, G$를 각각 $\widetilde{L} := \nu L,$ $\widetilde{\mu}:= \varsigma \mu,$ $\widetilde{\sigma}_k := \sqrt{\nu} \sigma,$ $\widetilde{G} := \sqrt{\varsigma}G$로 바꾸면  convergence가 동일하게 보장됩니다. 이러한 방법을 $\text{Scheme T}$라고 이야기하겠습니다.

 물론, uniform distribution을 가정하더라도, 만약$p_1, p_2, \cdots, p_N$이 매우 non-uniform하다면, $\nu$는 매우 커지고 $\varsigma$는 매우 작아질 것입니다. 이는 FedAvg 알고리즘의 수렴 속도를 낮추는 요인이 되며, (이에 관한 이야기는 다음 포스트에서 다루겠습니다.) 따라서 이러한 경우에는 앞서 언급한 대로 $\text{Scheme I}$이 더 좋은 성능을 보일 것입니다. 하지만 앞서 이야기하였듯이 $\text{Scheme I}$은 이상적인 case에서만 성립합니다. 따라서 $\text{Scheme I}$과 $\text{Scheme II}$ 중 어떠한 방법이 더 낫다고 이야기하기는 어려우며, 우리는 상황에 따라 유연하게 둘 중 하나를 선택하거나, 혹은 이 둘의 중간자적인 성격의 $\text{Scheme}$을 사용하여야 합니다. $\text{Scheme T}$가 이에 해당하는 셈인데, 실제로 experiments를 보면 $\text{Scheme I}$가 제일 우수한 성능을 보여주기는 하나, $\text{T}$도 $\text{I}$에 버금가는 정도의 performance를 보여주고 있습니다. 즉, $\text{Scheme T}$가 적절한 대안이라는 것을 확인할 수 있습니다.

 

 다음 포스트에서는 hyperparameter $K$와 $E$를 어떻게 정해야 할지에 관하여 살펴본 후, experiments를 확인하도록 하겠습니다.