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={ xk,1,xk,2,xk,3,...,xk,nk}\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上，
wt+1k=vt+1k,    if  t+1∉JE={ nE: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∈JE={ nE: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 函数 { Fk: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 函数 { Fk:k=1,2,...,N}\{F^k:k=1,2,...,N\}{ Fk:k=1,2,...,N} 都是 $\mu$-strongly convex.
Fk(x2)−Fk(x1)≥∇f(x1)⊤(x2−x1)+μ2∥x2−x1∥2F^k(\bm{x_2})-F^k(\bm{x_1})\geq \nabla f(\bm{x_1}