联邦学习FedAvg记录

Notation

符号	含义
$F(w)$	总目标函数
$w$	待优化参数
$F_k(w)$	第$k$个client的目标函数
$p_k$	第$k$个client在总表函数中的占比
$E$	每个local update的次数
$T$	总迭代次数，即通讯次数为$T/E$
${\eta }_t$	$t$时刻学习率
$L$	函数$F_k$为$L-smooth$，即${\mathrm{\nabla }}^2{F}_k(w)\le L$
$\mu$	函数$F_k$为$\mu$convex，即${\mathrm{\nabla }}^2{F}_k(w)\ge u$
${\sigma }_k$	$E(\Vert \mathrm{\nabla }{F}_k(w_t^k,{\xi }_t^k)-\mathrm{\nabla }{F}_k(w_t^k)\Vert )\le {\sigma }_k^2$
$G$	$E(\Vert \mathrm{\nabla }{F}_k(w_t^k,{\xi }_t^k)\Vert )\le G^2$
$F^{\ast }$	最优函数值
$F_k^{\ast }$	单独对第$k$个client优化得到的最优函数值
$\mathrm{\Gamma }$	$F^{\ast }-\sum _k{p}_k{F}_k^{\ast }$，度量异质性
${\eta }_t$	$t$时刻的学习率
$\kappa$	$\frac{L}{\mu }$，可近似看为条件数
$w^{\ast }$	最优参数
${\xi }_t^k$	client $k$在$t$时刻进行随机梯度下降选出的样本

假设

1. 函数$F_k$为$L-smooth$，对于所有的$k$
2. 函数$F_k$为$\mu -convex$，对于所有的$k$
3. client$k$每次计算的随机梯度的方差是${\sigma }_k^2$有界的
4. 所有计算的随机梯度的范数是$G$有界的

全参与下的收敛性证明

引入两个变量，$v_t^k$和$w_t^k$，

$$\begin{array}{rl}{v}_{t+1}^k& =v_t^k-{\eta }_t\mathrm{\nabla }{F}_k(w_t^k,{\xi }_t^K)\\ w_{t+1}^k& =\{\begin{array}{cc}{v}_{t+1}& \text{for }t+1\notin I_E\\ \sum _k^N{p}_k{v}_{k+1}^k& \text{for }t+1\in I_E\end{array}\end{array}$$

定义

$$\begin{array}{rl}{\overline{v}}_t& ={\sum }_k{p}_k{v}_t^k\\ {\overline{w}}_t& ={\sum }_k{p}_k{w}_t^k\\ {\overline{g}}_t& ={\sum }_k{p}_k\mathrm{\nabla }{F}_k(w_t^k)\\ g_t& ={\sum }_k{p}_k\mathrm{\nabla }{F}_k(w_t^k,{\xi }_k^t)\end{array}$$

由于$t\in I_E$，在能够交换的迭代轮次才能获取参数的更新$w$，变量$v$用来表示在不能进行交换数据的轮次的参数。由于全局参与，${\overline{v}}_t={\overline{w}}_t$对于所有的$t$，而且${\overline{v}}_{t+1}={\overline{w}}_t-{\eta }_t{g}_t$。

个人理解：

要证明收敛性需要证明，参数是收敛的，由于参数${\overline{w}}_t$是根据梯度下降求出来的，所以需要证明，

$$E\Vert {\overline{w}}_{t+1}-w^{\ast }\Vert \le l({\overline{w}}_t-w)$$

即当前迭代的参数和最优点的$w^{\ast }$的距离是小于上一次迭代参数与最优参数的距离，而且$l$函数是可以递推的。也就是说，当前迭代的参数和最优点的$w^{\ast }$的距离的上界是逐渐减小的。

文章没有选择$\overline{w}-w^{\ast }$而是选择了$\overline{v}-w^{\ast }$，因为$\overline{v}$是对应所有client的，在部分参与的场景下，$\overline{w}$是偏差的。

$$\begin{array}{}\text{(1)}& \begin{array}{rl}E\Vert {\overline{v}}_{t-1}-w^{\ast }{\Vert }^2& =E\Vert {\overline{w}}_t-{\eta }_t{g}_t-w^{\ast }-{\eta }_t{\overline{g}}_t+{\eta }_t{\overline{g}}_t{\Vert }^2\\ & =E(\Vert {\overline{w}}_t-{\eta }_t{\overline{g}}_t-w^{\ast }{\Vert }^2+2{\eta }_t<{\overline{w}}_t-{\eta }_t{\overline{g}}_t-w^{\ast },-g_t+{\overline{g}}_t>+{\eta }_t^2\Vert {\overline{g}}_t-g_t{\Vert }^2)\end{array}\end{array}$$

上面之所以拆分出$-{\eta }_t{\overline{g}}_t+{\eta }_t{\overline{g}}_t$的原意是想利用$E({\eta }_t{g}_t-{\eta }_t{\overline{g}}_t)=0$来对$\text{1}$进行拆分。

对于$E(\Vert {\overline{w}}_t-{\eta }_t{\overline{g}}_t-w^{\ast }{\Vert }^2)$继续进行计算

$$\begin{array}{rl}\Vert {\overline{w}}_t-{\eta }_t{\overline{g}}_t-w^{\ast }{\Vert }^2& =\Vert {\overline{w}}_t-w^{\ast }{\Vert }^2-2{\eta }_t<{\overline{w}}_t-w^{\ast },{\overline{g}}_t>+{\eta }_t^2\Vert {\overline{g}}_t{\Vert }^2\end{array}$$

根据$L-smooth$^[1]

$$\begin{array}{}\text{(3)}& \Vert \mathrm{\nabla }{F}_k(w_t^k){\Vert }^2\le 2L(F_k(w_t^k-F_k^{\ast }))\end{array}$$

因为二范数为凸函数再结合$\text{3}$，得到

$$\begin{array}{rl}{\eta }_t^2\Vert {\overline{g}}_t{\Vert }^2& \le {\eta }_t^2\sum p_k\Vert \mathrm{\nabla }{F}_k(w_t^k){\Vert }^2\\ & \le 2L{\eta }_t^2(F_k(w_t^k-F_k^{\ast }))\end{array}$$

对于$2{\eta }_t<{\overline{w}}_t-w^{\ast },{\overline{g}}_t>$展开

$$\begin{array}{rl}-2{\eta }_t<{\overline{w}}_t-w^{\ast },{\overline{g}}_t>& =-2{\eta }_t\sum p_k<{\overline{w}}_t-w^{\ast },\mathrm{\nabla }{F}_k(w_t^k)>\\ & =-2{\eta }_t\sum p_k<{\overline{w}}_t-w_t^k,\mathrm{\nabla }{F}_k(w_t^k)>-2{\eta }_t\sum p_k<w_t^k-w_t^k,\mathrm{\nabla }{F}_k(w_t^k)>\end{array}$$

根据可惜施瓦茨不等式和矩阵不等式得到

$$\begin{array}{r}-2<{\overline{w}}_t-w_t^k,\mathrm{\nabla }{F}_k(w_t^k)>\le \frac{1}{{\eta }_t}\Vert {\overline{w}}_t-w_t^k{\Vert }^2+{\eta }_k\Vert \mathrm{\nabla }{F}_k(w_t^k){\Vert }^2\end{array}$$

根据$\mu -convex$得到

$$\begin{array}{r}-<w_t^k-w^{\ast },\mathrm{\nabla }{F}_k(w_t^k)>\le -(F_k(w_t^k)-F_k(w^{\ast }))-\frac{\mu }{2}\Vert w_t^k-w^{\ast }{\Vert }^2\end{array}$$

因此$\text{???}$写为

$$\begin{array}{rl}\Vert {\overline{w}}_t-{\eta }_t{\overline{g}}_t-w^{\ast }{\Vert }^2& =\Vert {\overline{w}}_t-w^{\ast }{\Vert }^2-2{\eta }_t<{\overline{w}}_t-w^{\ast },{\overline{g}}_t>+{\eta }_t^2\Vert {\overline{g}}_t{\Vert }^2\\ & \le \Vert {\overline{w}}_t-w^{\ast }{\Vert }^2+2L{\eta }_t^2(F_k(w_t^k-F_k^{\ast }))+\\ & {\eta }_t\sum p_k(\frac{1}{{\eta }_t}\Vert {\overline{w}}_t-w_t^k{\Vert }^2+{\eta }_k\Vert \mathrm{\nabla }{F}_k(w_t^k){\Vert }^2)-\\ & 2{\eta }_t\sum p_k(F_k(w_t^k)-F_k(w^{\ast })-\frac{\mu }{2}\Vert w_t^k-w^{\ast }{\Vert }^2)\\ & \le (1-\mu {\eta }_t)\Vert {\overline{w}}_t-w^{\ast }{\Vert }^2+\sum p_k\Vert {\overline{w}}_t-w_t^k{\Vert }^2+\\ & 4L{\eta }_t^2\sum p_k(F_k(w_t^k)-F_k^{\ast })-2{\eta }_t\sum p_k(F_k(w_t^k)-F_k(w^{\ast }))\end{array}$$

拿出后面的$4L{\eta }_t^2\sum p_k(F_k(w_t^k)-F_k^{\ast })-2{\eta }_t\sum p_k(F_k(w_t^k)-F_k(w^{\ast }))$一项，定义${\gamma }_t=2{\eta }_t(1-2L{\eta }_t)$，同时假定${\eta }_t$随时间非增（也就是学习率是衰减的）且${\eta }_t\le \frac{1}{4L}$。可以得到定义的${\eta }_t\le {\gamma }_t\le 2{\eta }_t$。整理得到

$$\begin{array}{rl}& 4L{\eta }_t^2\sum p_k(F_k(w_t^k)-F_k^{\ast })-2{\eta }_t\sum p_k(F_k(w_t^k)-F_k(w^{\ast }))\\ & =-2{\eta }_t(1-2L{\eta }_t)\sum p_k(F_k(w_t^k)-F_k^{\ast })+2{\eta }_t\sum p_k(F_k(w^{\ast })-F_k\ast )\\ & =-{\gamma }_t\sum p_k(F_k(w_t^k)-F^{\ast })+(2{\eta }_t-{\gamma }_t)\sum p_k(F^{\ast }-F_k^{\ast })\\ & =-{\gamma }_t\sum p_k(F_k(w_t^k)-F^{\ast })+4L{\eta }_t^2\mathrm{\Gamma }\end{array}$$

第三个等号将$F_k(w_t^k-F_k^{\ast })$拆分为$F_k(w_t^k)-F^{\ast }+F^{\ast }-F_k^{\ast }$

$$\begin{array}{rl}\sum p_k(F_k(w_t^k)-F^{\ast })& =\sum p_k(F_k(w_t^k)-F_k({\overline{w}}_t))+\sum p_k(F_k({\overline{w}}_t)-F^{\ast })\\ & \ge \sum p_k<\mathrm{\nabla }{F}_k({\overline{w}}_t),w_t^k-{\overline{w}}_t>+F_k({\overline{w}}_t)-F^{\ast }\\ & \ge -\frac{1}{2}\sum p_k[{\eta }_t\Vert \mathrm{\nabla }{F}_k({\overline{w}}_t){\Vert }^2+\frac{1}{\eta }\Vert w_t^k-{\overline{w}}_t{\Vert }^2]+(F({\overline{w}}_t)-F^{\ast })\\ & \ge -\sum p_k[{\eta }_{t}L(F_k({\overline{w}}_t)-F_k^{\ast })+\frac{1}{2\eta }\Vert w_t^k-{\overline{w}}_t{\Vert }^2]+(F({\overline{w}}_t)-F^{\ast })\end{array}$$

综上所述，

$$\begin{array}{rl}& 4L{\eta }_t^2\sum p_k(F_k(w_t^k)-F_k^{\ast })-2{\eta }_t\sum p_k(F_k(w_t^k)-F_k(w^{\ast }))\\ & \le {\gamma }_t\sum p_k[{\eta }_{t}L(F_k({\overline{w}}_t)-F_k^{\ast })+\frac{1}{2\eta }\Vert w_t^k-{\overline{w}}_t{\Vert }^2]-\gamma (F({\overline{w}}_t)-F^{\ast })++4L{\eta }_t\mathrm{\Gamma }\\ & ={\gamma }_t({\eta }_{t}L-1)\sum p_k(F_k({\overline{w}}_t)-F^{\ast })+(4L{\eta }_t^2+{\gamma }_t{\eta }_{t}L)+\frac{\gamma }{2{\eta }_t}\sum p_k\Vert w_t^k-{\overline{w}}_t{\Vert }^2\\ & \le 6L{\eta }_t^2\mathrm{\Gamma }+\sum p_k\Vert w_t^k-{\overline{w}}_t{\Vert }^2\end{array}$$

最后一个不等式取得的原因是$({\eta }_{t}L-1)\le -\frac{3}{4}$，$\sum p_k(F_k({\overline{w}}_t)-F^{\ast })\ge 0$

因此

$$\Vert {\overline{w}}_t-{\eta }_t{\overline{g}}_t-w^{\ast }{\Vert }^2\le (1-\mu {\eta }_t)\Vert {\overline{w}}_t-w^{\ast }{\Vert }^2+6L{\eta }_t^2\mathrm{\Gamma }+\sum p_k\Vert w_t^k-{\overline{w}}_t{\Vert }^2$$

加上梯度方差有限假设，即

$$\begin{array}{rl}E\Vert g_t-{\overline{g}}_t{\Vert }^2& =E\Vert \sum p_k(\mathrm{\nabla }{F}_k(w_k,{\xi }_t^k)-\mathrm{\nabla }{F}_k(w_t^k))\Vert \\ & \le \sum p_k^2{\sigma }_k^2\end{array}$$

$$\begin{array}{rl}E\sum p_k\Vert {\overline{w}}_t-w_t^k{\Vert }^2& =E\sum p_k\Vert w_t^k-{\overline{w}}_{t_0}-({\overline{w}}_t-{\overline{w}}_{t_0}){\Vert }^2\\ & \le E\sum p_k\Vert w_t^k-{\overline{w}}_{t_0}{\Vert }^2\\ & \le \sum p_{k}E{\sum }_{t=t_0}^{t-1}(E-1){\eta }_t^2\Vert \mathrm{\nabla }{F}_k(w_t^k,{\xi }_t^k)\Vert \\ & \le 4{\eta }_t^2(E-1{)}^2{G}^2\end{array}$$

因此最终，令${\mathrm{\Delta }}_t=E\Vert {\overline{w}}_{t+1}-w^{\ast }\Vert$

$$\begin{array}{r}{\mathrm{\Delta }}_{t+1}\le (1-\mu {\eta }_t){\mathrm{\Delta }}_t+{\eta }_t^{2}B\end{array}$$

$B=\sum p_k^2{\sigma }_k^2+6L\mathrm{\Gamma }+8(E-1{)}^2{G}^2$

参考资料

1. Francis Bach, Statistical machine learning and convex optimization