[ICLR 2020] Convergence of FedAvg - (4)

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

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

 

 지난 두 개의 포스트를 통해, 우리는 full decive participation case에서 FedAvg의 convergence를 확인하였습니다. (이전 글 보기) 이번 포스트부터는 partial device participation case를 살펴볼 것입니다. 마찬가지로 증명 과정이 다소 길기 때문에 두세 편에 나누어서 작성될 예정입니다.

 

5. Assumtions - (2)

 

 이전 포스트에서 살펴보았듯이,  Sampling 과정에서 replacement가 허용되는지(FedProx 논문에서 제시한 Sampling 기법) 혹은 그렇지 않은지(해당 논문에서 제시한 Sampling 기법)에 따라 Averaging 기법이 상이하게 나타납니다. 해당 논문에서는 전자를 $\text{Scheme I}$, 후자를 $\text{Scheme II}$로 표기하며, 아래 두 assumption들은 각각의 Sampling과 Averaging에 관한 것입니다. 여기에서, $[N]$은 서로 다른 $N$개의 client의 집합을 의미하며, 정의에 따라 $\mathcal{S_t} \subset [N]$입니다. (partial case이기 때문에 $\mathcal{S_t} = [N]$은 고려하지 않습니다.)

 

$\text{Assumption 5}$ [$\text{Scheme I}$]

 $[N]$으로부터 $K$개의 device를 sampling probability $p_1, p_2, \cdots, p_k$에 따라 replacement를 허용하여 sampling한 것을 $\mathcal{S}_t$라고 하자. 이때, FedAvg 알고리즘의 aggregation step은 $w_t \leftarrow \frac {1} {K} \sum_{k \in \mathcal{S}_t} w_t^k$로 진행된다.

 

$\text{Assumption 6}$ [$\text{Scheme II}$]

 $[N]$으로부터 $K$개의 device를 replacement 없이 uniform하게 sampling한 것을 $\mathcal{S}_t$라고 하자. 그리고 $p_1 = p_2 = \cdots = p_k = \frac {1} {N}$이라는 의미에서 data가 balanced라고 가정하자. 이때, FedAvg 알고리즘의 aggregation step은 $w_t \leftarrow \frac {N} {K} \sum_{k \in \mathcal{S}_t} p_k w_t^k$로 진행된다.

 

 $\text{Assumption 6}$에는 balanced data라는 가정이 추가되어 있는데, 이것은 data가 서로 IID하다는 의미가 아니라, data의 개수가 고르게 분포한다는 의미입니다. 물론, 이러한 가정 역시 현실과 다소 동떨어진 것입니다. 현실에서는 각 device가 서로 다른 개수의 data를 가지고 있는 것이 일반적이죠. 그럼에도 왜 이러한 비현실적인 가정을 감수하면서 $\text{Scheme II}$를 고려해야 하는지, 그리고 이러한 가정을 정당화하거나 완화할 수는 없는지에 대한 논의는 잠시 뒤로 미루도록 하겠습니다.

 

6. Key Lemmas for $\text{Theorem 2}$ and $3$ 

 

 full device participation case에서는 $t + 1 \in \mathcal{I}_E$인 경우 $t + 1$기에 자동으로 aggregation step을 진행하였습니다. 하지만 partial device participation case에서는 $t + 1$기에 적절한 scheme을 사용하여 $N$개의 device 중 $K$개를 적절히 sampling한 뒤, 이 $K$개의 device에 대해서만 aggregation step을 진행해야 합니다. 문제는, 먼저 $\mathcal{S_t}$를 구성한 뒤, 이에 해당하는 client에 대해서만 update를 진행하는 것이 (communication cost 등을 고려할 때) 현실적이지만, $\mathcal{S_t}$는 $t$에 따라 계속해서 변하기 때문에 이러한 과정을 그대로 반영할 경우 convergence analysis에 어려움이 발생한다는 것입니다. 따라서, 저자들은 우선 모든 device가 local update를 수행한 후, sampling된 $K$개의 device에게서만 update 결과를 받아 global update를 하는 상황을 가정하고 analysis를 진행하였습니다. 결론적으로는 (local에서의 computation cost는 다소 증가하겠지만, convergence의 관점에서는) 동일한 과정이기 때문에 이는 큰 문제가 없는 가정이며, 이러한 가정 하에 $t$기에서 $t + 1$기로의 update 과정은 다음과 같습니다.

$$v_{t+1}^k = w_t^k - \eta_t \nabla F_k(w_t^k, \xi_t^k)$$

$$\begin{equation} w_{t+1}^k = \begin{cases} v_{t+1}^k & \text{if $t+1 \notin \mathcal{I}_E$}\\ \text{samples } \mathcal{S}_{t+1} \text{ and average } \{ v_{t+1}^k \}_{k \in \mathcal{S}_t} & \text{if $t+1 \in \mathcal{I}_E$} \end{cases} \end{equation}$$

 여기에서 Sampling은 어떠한 scheme을 사용하는지에 따라서 달라질 것이며, 이 Sampling 때문에 우리는 더 이상 $\bar{w}_{t+1} = \bar{v}_{t+1}$을 보장할 수 없게 되었습니다. 평소에는 해당 식이 성립하겠지만, 문제가 되는 경우는 $t+1 \in \mathcal{I}_E$일 때입니다. $\bar{w}_{t+1} \neq \bar{v}_{t+1}$라는 것은 곧 특정 device(들)에 biased된 학습을 하고 있다는 것이므로, 우리는 평균적으로라도 $\bar{w}_{t+1}$과 $\bar{v}_{t+1}$이 같다는 보장이 필요합니다. 다행히도, 이를 $\text{Lemma 4}$가 해결해줍니다.

 

$\text{Lemma 4}$ [Unbiased Sampling Scheme]

 $t + 1 \in \mathcal{I}_E$일 때, $\text{Scheme I}$과 $\text{II}$ 모두 $\mathbb{E}_{S_t} [\bar{w}_{t+1}] = \bar{v}_{t+1}$을 만족한다.

 

$\text{Proof}$

  $t$기에서 $t+1$기로 넘어가는 과정에서 $\mathcal{S}_t := \{ i_1, i_2, \cdots, i_K \} \subset [N]$를 sampling할 때, 어떠한 (이미 정해져서 변하지 않는) 수열 $\{ x_i \}_{i=1}^{N}$에 대하여 $x_k$를 $q_k$의 확률로 함께 sampling한다고 가정하자. (이때, scheme에 따라서 특정 $i_k$들이 서로 동일할 수도 있으며, $\text{Scheme I}$의 경우 $q_k = p_k$이고, $\text{Scheme II}$의 경우 $q_k = \frac {1} {N}$이다.) 이때, $N - K$개의 선택되지 못한 $x_k$들은 global update에 반영되지 않으므로, 다음과 식이 성립한다. $(*)$

$$\mathbb{E}_{\mathcal{S}_t} \left[ \sum_{k \in \mathcal{S}_t} x_k \right] = \mathbb{E}_{\mathcal{S}_t} \left[ \sum_{k=1}^K x_{i_k} \right] = K \sum_{k=1}^N q_k x_k$$

 

 따라서, $\mathbb{E}_{\mathcal{S}_t} [\bar{w}_{t+1}] = \mathbb{E}_{\mathcal{S}_t} \left[ \frac {1} {K} \sum_{k \in \mathcal{S}_t} v_{t+1}^k \right] = K \sum_{k=1}^N \frac {1} {K} q_k v_{t+1}^k = \sum_{k=1}^N q_k v_{t+1}^k = \bar{v}_{t+1}$ 이다. $\square$

 

 저자들이 FedAvg 논문에서 제안되었던 기존의 Sampling 방법은 고려하지 않았는데, 그 이유는 해당 방법을 사용할 경우 $\text{Lemma 4}$가 성립하지 않기 때문입니다. 물론, 이 부분이 기존의 FedAvg가 heterogeneous한 구성에서 적절하게 동작하지 못한 근본적인 원인이라고 단정지을 수는 없지만, 적어도 Sampling과  Averaging 기법을 잘 선택하는 것이 convergence에 직접적으로 영향을 준다는 것은 알 수 있으며, 이는 저자들도 강조하는 부분입니다. 그리고 이러한 이유 때문에, 기존의 방법론의 경우 앞서 살펴본 표에 convergence rate가 기입되어 있지 않습니다.

 

 다음으로, $\text{Lemma 5}$는 $\mathbb{E}_{\mathcal{S}_t} [||\bar{v}_{t+1} - \bar{w}_{t+1}||^2]$이 bounded되어있음을 보장해줍니다. 이는 $\text{Lemma 4}$에 의하여 $\mathbb{E}_{S_t} [\bar{w}_{t+1}] = \bar{v}_{t+1}$이므로 $\text{Var} [\bar{w}_{t+1}]$이 bounded되어있다는 것과 동치입니다.

 

$\text{Lemma 5}$ [Bounding the variance of $\bar{w}_t$]

 임의의 t에 대해서 $\eta_t$가 non-ibcreasing하고, $\eta_t \leq 2 \eta_{t + E}$일 때, 다음이 성립한다.

 $(i)$ 만약 $\text{Scheme I}$을 채택하였다면, $\mathbb{E}_{\mathcal{S}_t} [||\bar{v}_{t+1} - \bar{w}_{t+1}||^2] \leq \frac {4} {K} \eta_t^2 E^2 G^2$이다.

 $(ii)$ 만약 $\text{Scheme II}$를 채택하였다면, $\mathbb{E}_{\mathcal{S}_t} [||\bar{v}_{t+1} - \bar{w}_{t+1}||^2] \leq \frac {N - K} {N - 1} \frac {4} {K} \eta_t^2 E^2 G^2$이다.

        (단, 이때 $\text{Assumption 6}$에 의하여 $p_1 = p_2 = \cdots = p_N = \frac {1} {N}$이라고 가정한다.)

 

$\text{Proof}$

 증명에 앞서, $t_0 := t - E + 1$로 정의한다. (즉, $t_0$는 바로 직전의 communication round를 의미한다.)

 $(i)$ $\bar{w}_{t+1} = \frac {1} {K} \sum_{k \in \mathcal{S}_t} v_{t+1}^k = \frac {1} {K} \sum_{l=1}^K v_{t+1}^{i_l}$이므로, $\mathbb{E}_{\mathcal{S}_t} [||\bar{w}_{t+1} - \bar{v}_{t+1}||^2]$는 다음과 같이 bounded된다.

 

$$\begin{align*} \mathbb{E}_{\mathcal{S}_t} [||\bar{w}_{t+1} - \bar{v}_{t+1}||^2] &= \mathbb{E}_{\mathcal{S}_t} \left[ \left| \left| \frac {1} {K} \sum_{l=1}^K v_{t+1}^{i_l} - \bar{v}_{t+1} \right| \right|^2 \right] = \mathbb{E}_{\mathcal{S}_t} \left[ \frac {1} {K^2} \sum_{l=1}^K || v_{t+1}^{i_l} - \bar{v}_{t+1} ||^2 \right] \\&= \frac {1} {K^2} \times K \sum_{k=1}^N p_k ||  v_{t+1}^{k} - \bar{v}_{t+1} ||^2 \; (\because (*)) \\&= \frac {1} {K} \sum_{k=1}^N p_k || v_{t+1}^{k} - \bar{v}_{t+1} ||^2 \\&= \frac {1} {K} \sum_{k=1}^N p_k || (v_{t+1}^{k} - \bar{w}_{t_0}) - (\bar{v}_{t+1} - \bar{w}_{t_0}) ||^2 \\&\leq \frac {1} {K} \sum_{k=1}^N p_k || v_{t+1}^{k} - \bar{w}_{t_0} ||^2 \\&\quad (\because \sum_{k=1}^N p_k(v_{t+1}^k - \bar{w}_{t_0}) = \bar{v}_{t+1} - \bar{w}_{t_0} \text{ and } \mathbb{E} [ ||X - \mathbb{E} [X] ||^2 ] \leq \mathbb{E} [||X||^2]) \\&= \frac {1} {K} \sum_{k=1}^N p_k \mathbb{E} || v_{t+1}^{k} - \bar{w}_{t_0} ||^2 \; (\because \mathbb{E} [X] = \mathbb{E} [\mathbb{E} [X | Y]]) \\&\leq \frac {1} {K} \sum_{k=1}^N p_k \mathbb{E} || v_{t+1}^{k} - \bar{w}_{t_0}^k ||^2 \\&\leq \frac {1} {K} \sum_{k=1}^N p_k E \sum_{i={t_0}}^t \mathbb{E} || \eta_i \nabla F_k (w_i^k, \xi_i^k) ||^2 \; (\because v_{t+1}^k = w_t^k - \eta_t \nabla F_k(w_t^k, \xi_t^k)) \\&\leq \frac {1} {K} E^2 \eta_{t_0}^2 G^2 \; (\because \text{Assumption 4}) \\&\leq \frac {4} {K} \eta_t^2 E^2 G^2 \; (\because \eta_t \text{ is non-increasing and } \eta_{t_0} \leq 2 \eta_t) \end{align*}$$

 

 $(ii)$ 가정에 의하여 $p_1 = p_2 = \cdots = p_N = \frac {1} {N}$이므로, $(i)$과 동일하게 $\bar{w}_{t+1} = \frac {1} {K} \sum_{k \in \mathcal{S}_t} v_{t+1}^k = \frac {1} {K} \sum_{l=1}^K v_{t+1}^{i_l}$이다. 따라서,  $\mathbb{E}_{\mathcal{S}_t} [||\bar{w}_{t+1} - \bar{v}_{t+1}||^2]$는 다음과 같이 bounded된다. (여기에서, $\mathbb{I}$는 support이다.)

 

$$\begin{align*} \mathbb{E}_{\mathcal{S}_t} [||\bar{w}_{t+1} - \bar{v}_{t+1}||^2] &= \mathbb{E}_{\mathcal{S}_t} \left[ \left| \left| \frac {1} {K} \sum_{i \in \mathcal{S}_t} v_{t+1}^i - \bar{v}_{t+1} \right| \right| ^2 \right] = \frac{1} {K^2} \mathbb{E}_{\mathcal{S}_t} \left[ \sum_{i \in \mathcal{S}_t} \left| \left| v_{t+1}^i - \bar{v}_{t+1} \right| \right| ^2 \right] \\&= \frac{1} {K^2} \mathbb{E}_{\mathcal{S}_t} \left[ \left| \left| \sum_{i=1}^N \mathbb{I}_{\{ i \in \mathcal{S}_t \}} (v_{t+1}^i - \bar{v}_{t+1}) \right| \right| ^2 \right] = \frac{1} {K^2} \left| \left| \sum_{i=1}^N \mathbb{P} (i \in \mathcal{S}_t) (v_{t+1}^i - \bar{v}_{t+1}) \right| \right| ^2 \\&= \frac{1} {K^2} \left\langle \sum_{i=1}^N \mathbb{P} (i \in \mathcal{S}_t) (v_{t+1}^i - \bar{v}_{t+1}), \; \sum_{i=1}^N \mathbb{P} (i \in \mathcal{S}_t) (v_{t+1}^i - \bar{v}_{t+1}) \right\rangle \\&= \frac {1} {K^2} \left[ \sum_{i \in [N]} \mathbb{P} (i \in \mathcal{S}_t) \left| \left| v_{t+1}^i - \bar{v}_{t+1} \right| \right| ^2 + \sum_{ \; \: i \neq j \\ i, j \in [N]} \mathbb{P} (i, j \in \mathcal{S}_t) \langle v_{t+1}^i - \bar{v}_{t+1}, \; v_{t+1}^j - \bar{v}_{t+1} \rangle \right] \\&= \frac {1} {K^2} \sum_{i=1}^N \frac {K} {N} ||v_{t+1}^i - \bar{v}_{t+1}||^2 + \sum_{ \; \: i \neq j \\ i, j \in [N]} \frac {K - 1} {KN (N-1)} \langle v_{t+1}^i - \bar{v}_{t+1}, \; v_{t+1}^j - \bar{v}_{t+1} \rangle \\&\quad (\because \mathbb{P} (i \in \mathcal{S}_t) = \frac {K} {N} \text{ and } \mathbb{P} (i, j \in \mathcal{S}_t) = \frac{K} {N} \times \frac{K-1} {N-1}) \\&= \frac {1} {KN} \sum_{i=1}^N ||v_{t+1}^i - \bar{v}_{t+1}||^2 + \frac {K - 1} {KN (N-1)} \sum_{ \; \: i \neq j \\ i, j \in [N]} \langle v_{t+1}^i - \bar{v}_{t+1}, \; v_{t+1}^j - \bar{v}_{t+1} \rangle \\&= \frac {1} {KN} \sum_{i=1}^N ||v_{t+1}^i - \bar{v}_{t+1}||^2 + \frac {K - 1} {KN (N-1)} \times \left(- \sum_{i=1}^N ||v_{t+1}^i - \bar{v}_{t+1}||^2 \right) \\&= \left( \frac {1} {KN} - \frac {K - 1} {KN (N-1)} \right) \sum_{i=1}^N ||v_{t+1}^i - \bar{v}_{t+1}||^2 \\&= \frac {1} {KN} \left( 1 - \frac {K - 1} {N-1} \right) \sum_{i=1}^N ||v_{t+1}^i - \bar{v}_{t+1}||^2 \\&= \frac {1} {KN} \left(\frac {N-K} {N-1} \right) \sum_{i=1}^N ||v_{t+1}^i - \bar{v}_{t+1}||^2 \\&= \frac {1} {K(N-1)} \left(\frac {N-K} {N} \right) \sum_{i=1}^N ||v_{t+1}^i - \bar{v}_{t+1}||^2 \\&= \frac {1} {K(N-1)} \left(1 - \frac {K} {N} \right) \sum_{i=1}^N ||v_{t+1}^i - \bar{v}_{t+1}||^2 \\&= \frac {N} {K(N-1)} \left(1 - \frac {K} {N} \right) \frac {1} {N} \sum_{i=1}^N ||v_{t+1}^i - \bar{v}_{t+1}||^2 \\&= \frac {N} {K(N-1)} \left(1 - \frac {K} {N} \right) \mathbb{E} \left[ \frac {1} {N} ||v_{t+1}^i - \bar{v}_{t+1}||^2 \right]  \; (\because \mathbb{E} [X] = \mathbb{E} [\mathbb{E} [X | Y]]) \\&\leq \frac {N} {K(N-1)} \left(1 - \frac {K} {N} \right) 4 \eta_t^2 E^2 G^2 \; (\because \text{the same with } (i)) \\&= \frac {N - K} {N - 1} \frac {4} {K} \eta_t^2 E^2 G^2 \end{align*}$$

 

 따라서, $(i)$과 $(ii)$ 모두 성립한다. $\square$

 

 다음 포스트에서는 $\text{Lemma 4, 5}$와 $\text{Theorem 1}$을 이용하여 $\text{Theorem 2, 3}$을 증명하도록 하겠습니다.