On the convergence of FedAvg on non-iid data

在这篇blog中我们一起来阅读一下 On the convergence of FedAvg on non-iid data 这篇 ICLR 2020 的paper.

主要目的

本文的主要目的是证明联邦学习算法的收敛性。与之前其他工作中的证明不同，本文的证明更贴近于实际联邦学习的场景。特别的，

1. 所有用户的数据non-iid分布；
2. 每次只有一部分用户参与FedAvg.

系统模型

考虑一个联邦学习系统 with $N$ 用户和一个PS. 每用户有一些local data，训练发生在用户处，每隔一段时间用户上传自己学习的模型来做FedAvg.

将第 $k$ 个用户的数据记为 x = { x k , 1 , x k , 2 , x k , 3 , . . . , x k , n k } \bm{x}=\{x_{k,1},x_{k,2},x_{k,3},...,x_{k,n_k}\} x={ xk,1​,xk,2​,xk,3​,...,xk,nk​​}, 每个人都有一个学习目标，即最小化 loss 函数
$$F_k(\mathbf{w})=\sum _{j=1}^{n_k}{\ell }_k(\mathbf{w},x_{k,j})\text{\hspace{1em}(1)}$$

其中 ${\ell }_k(\mathbf{w},x_{k,j})$ 是每个训练数据的 loss. $F_k(\mathbf{w})$ 相当于是每个人所有数据上的loss，如果仅仅做local training, 那最终每个用户会 arrive at
$$\text{Local minimum}:F_k^{\ast }=\underset{\mathbf{w}}{\min }{F}_k$$

而FL考虑的是一种分布式的优化，即我们要minimize的目标函数
$$\text{Global minimum}:F^{\ast }=\underset{\mathbf{w}}{\min }\sum _{k=1}^N{p}_k{F}_k(\mathbf{w})$$

其中 $p_k$ 是一个distribution用来表示每个用户所占的权重。换句话说，我们最终想找到一个共同的 $\mathbf{w}$ 来最小化每个用户 loss 的一个加权和。

To this end, 本文考虑FedAvg, 并证明其能收敛到 global optimum.

FedAvg 的具体步骤描述如下：首先，我们按单次SGD为一个时间刻度把时间轴分为离散的slot $t=1,2,3,...,T$, 即总共进行 $T$ 次 local SGD, 每次 SGD每个用户从自己的数据集中随机均匀的采样出一个数据来进行训练。特别的，每隔 $E$ slots, 所有 active users 把自己的本地参数发送给PS进行 FedAvg，之后PS会把avg后的参数发还给各个用户。以上模型用数学语言可以写为以下两步：

Local training

每个用户在第 $t$ 个时刻基于 ${\mathbf{w}}_t^k$ 进行 SGD, 得到
$${\mathbf{v}}_{t+1}^k={\mathbf{w}}_t^k-{\eta }_t\nabla {\ell }_k({\mathbf{w}}_t^k,{\xi }_t^k)\text{\hspace{1em}(2)}$$

其中 ${\xi }_t^k$ 是从本地数据中随机采样出的一个sample。注意，这样单步SGD得到的${\mathbf{v}}_{t+1}^k$ 只是一个中间变量而不是下一时刻的 ${\mathbf{w}}_{t+1}^k$，因为我们还有可能做 FedAvg。 更具体地说，在 $E$ 的非整数倍slot上，
w t + 1 k = v t + 1 k ,     if   t + 1 ∉ J E = { n E : n = 1 , 2 , . . . } . \bm{w}^k_{t+1}=\bm{v}^k_{t+1},~~~~\text{if}~~t+1\notin\mathcal{J}_E=\{nE:n=1,2,...\}. wt+1k​=vt+1k​,    if  t+1∈/​JE​={ nE:n=1,2,...}.

而在 $E$ 的整数倍slot上，我们还得额外做 FedAvg.

FedAvg

若下一时刻是$E$的整数倍周期，即 t + 1 ∈ J E = { n E : n = 1 , 2 , . . . } t+1\in\mathcal{J}_E=\{nE:n=1,2,...\} t+1∈JE​={ nE:n=1,2,...}，我们进行FedAvg，此时
$${\mathbf{w}}_{t+1}^k=\sum _{k=1}^N{p}_k{\mathbf{v}}_{t+1}^k\text{\hspace{1em}(3)}$$

注意，这里面我们假设每个人都参与更新，稍后我们会release这个条件允许PS按照某种分布采样一部分人进行更新。

小结

如果我们从每个用户的角度看，它的参数变化可以用下图归纳 ($E=3$)。

几个假设

本文的推导基于以下假设。

Assumption 1 ($L$-smoothness). 所有用户的 loss 函数 { F k : k = 1 , 2 , . . . , N } \{F^k:k=1,2,...,N\} { Fk:k=1,2,...,N} 都是 L-smooth.
$$F^k({\mathbf{x}}_2)-F^k({\mathbf{x}}_1)\le \nabla f({\mathbf{x}}_1{)}^{\top }({\mathbf{x}}_2-{\mathbf{x}}_1)+\frac{L}{2}\Vert {\mathbf{x}}_2-{\mathbf{x}}_1{\Vert }^2$$

Assumption 2 ($\mu$-strongly convex). 所有用户的 loss 函数 { F k : k = 1 , 2 , . . . , N } \{F^k:k=1,2,...,N\} { Fk