0x00 Preface

Part 3 分析了 mini-batch SGD,Part 4 讨论同步和异步分布式更新。这一篇转向联邦学习:客户端不仅持有不同的数据,还会在一次通信之间连续更新多步。因此,本地随机梯度对本地目标无偏,并不意味着聚合后的方向对当前全局梯度无偏。

本篇从 FedAvg 的更新式出发,给出一个可以逐步核对的简化证明:先分析全参与、等权客户端、相同本地步数的模型,再推导均匀部分参与的扩展,并讨论不一致的聚合权重。非凸部分保证平均梯度范数变小;强凸部分额外使用 PL 不等式。两种结论不能混用。

FedAvg 的原始算法见 McMahan et al., 2017。下面的常数界是本文在明确假设下自行推导的保守界,不是照录某篇论文的定理,也不追求最优复杂度。

0x01 Objective and FedAvg Updates

Global Objective

设有 $N$ 个客户端,第 $i$ 个客户端的局部目标为

$$ f_i(w)=\mathbb E_{\xi\sim\mathcal D_i}[\ell(w;\xi)],\qquad f(w)=\sum_{i=1}^N p_i f_i(w),\qquad p_i>0,\quad\sum_{i=1}^Np_i=1. $$

如果希望每个样本具有相同权重,通常取 $p_i=n_i/\sum_j n_j$;如果希望每个客户端具有相同权重,则取 $p_i=1/N$。这是两个不同的优化目标,并非可以随意互换的实现细节。数据不上传也不等于自动具有差分隐私保证;本篇只讨论优化误差。

后面的主定理固定使用 $p_i=1/N$。加权目标和部分参与在后文讨论,不把它们默认包含在主定理中。

Local Steps and Communication Rounds

$t$ 表示通信轮,$k$ 表示轮内本地更新步。每轮所有客户端从同一个服务器模型开始,执行 $E\ge1$ 次本地 SGD:

$$ \begin{aligned} w_{i,t}^{(0)}&=w_t,\\ w_{i,t}^{(k+1)}&=w_{i,t}^{(k)}-\eta g_{i,t}^{(k)},\qquad k=0,\ldots,E-1,\\ w_{t+1}&=\frac1N\sum_{i=1}^Nw_{i,t}^{(E)}. \end{aligned} $$

$\eta$ 是固定的本地步长,服务器直接平均模型,不额外乘服务器学习率。定义一轮的有效步长 $\alpha=\eta E$,以及聚合方向

$$ v_t=\frac1{NE}\sum_{i=1}^N\sum_{k=0}^{E-1}g_{i,t}^{(k)}, \qquad w_{t+1}=w_t-\alpha v_t. $$

这里的 $E$ 是更新次数,不是本地 epoch 数。若各客户端的数据量不同,相同 epoch 数往往对应不同的更新次数,不能直接代入这一模型。总共 $T$ 轮时,每个客户端执行 $ET$ 次本地更新;全系统一共计算 $NET$ 个 batch 梯度。

0x02 Assumptions and Error Decomposition

Smoothness and Lower Bound

每个 $f_i$ 都是欧氏空间上的 $L$-光滑函数,因此平均目标 $f$ 也 $L$-光滑。假设 $f(w)\ge f_{\inf}>-\infty$,定义 $\Delta_0=f(w_0)-f_{\inf}$,并把 $w_0$ 视为确定的初值。非凸分析不要求 $f_{\inf}$ 被某个参数取到。

Conditional Gradient Noise

令 $\mathcal F_t$ 包含第 $t$ 轮开始前的历史;$\mathcal H_{t,k}$ 进一步包含该轮第 $k$ 步采样前所有客户端的参数和历史。写

$$ g_{i,t}^{(k)}=\nabla f_i(w_{i,t}^{(k)})+\varepsilon_{i,t}^{(k)}. $$ $$ \mathbb E[\varepsilon_{i,t}^{(k)}\mid\mathcal H_{t,k}]=0,\qquad \mathbb E[\|\varepsilon_{i,t}^{(k)}\|^2\mid\mathcal H_{t,k}] \le\frac{\sigma^2}{b}. $$

$b$ 是每步的 batch 大小,$\sigma^2$ 是单样本噪声方差的统一上界。同一步不同客户端的噪声还要求条件互不相关。同一客户端不同步之间不要求无条件独立,条件无偏使噪声增量构成鞅差序列。

batch 内的 $1/b$ 缩减需要独立采样或相应的协方差条件。直接遍历同一个随机排列、根据损失选择样本,或采样过程受到掉线机制影响时,都需要重新验证上述条件;不能仅凭“用了 SGD”就认为条件成立。

Bounded Gradient Heterogeneity

采用以下统一异质性上界:

$$ \frac1N\sum_{i=1}^N\|\nabla f_i(w)-\nabla f(w)\|^2\le\zeta^2, \qquad \forall w. $$

$\sigma^2$ 描述一个客户端内部的随机噪声,$\zeta^2$ 描述客户端之间的目标差异。增大 batch 能减少前者,却不能自动消除后者。由 $\sum_i(\nabla f_i-\nabla f)=0$,有

$$ \frac1N\sum_i\|\nabla f_i(w)\|^2 =\|\nabla f(w)\|^2+ \frac1N\sum_i\|\nabla f_i(w)-\nabla f(w)\|^2 \le\|\nabla f(w)\|^2+\zeta^2. $$

这是一个较强的假设,不由“非 IID 数据”这个描述自动推出。它要求差异在所有参数上有统一界,而不仅是最优点附近有界。若只在一个区域成立,还必须证明所有相关迭代留在该区域。本篇不把“全局梯度一致有界”作为额外假设。

Noise and Local Drift

令 $h_t=\nabla f(w_t)$,把聚合方向准确地拆成

$$ \begin{aligned} v_t&=h_t+r_t+e_t,\\ r_t&=\frac1{NE}\sum_{i,k} \left[\nabla f_i(w_{i,t}^{(k)})-\nabla f_i(w_t)\right],\\ e_t&=\frac1{NE}\sum_{i,k}\varepsilon_{i,t}^{(k)}. \end{aligned} $$

$r_t$ 是本地路径带来的漂移,$e_t$ 是平均随机噪声。条件无偏和互不相关性给出

$$ \mathbb E[e_t\mid\mathcal F_t]=0,\qquad \mathbb E[\|e_t\|^2\mid\mathcal F_t]\le \frac{\sigma^2}{bNE}. $$

不同步的噪声交叉项通过迭代条件期望消失。但 $r_t$ 取决于本地参数,而本地参数又取决于之前的噪声,所以一般不能断言 $\mathbb E\langle r_t,e_t\rangle=0$。后面的证明会保留这个相关性,通过范数不等式控制它。

0x03 Bounding Local Drift

这一节给出后面定理需要的关键估计。令 $q=E-1$,并定义给定轮初历史下的平均位移:

$$ D_{t,k}=\frac1N\sum_i \mathbb E[\|w_{i,t}^{(k)}-w_t\|^2\mid\mathcal F_t], \qquad M_t=\max_{0\le k\le q}D_{t,k}. $$

这里只需要控制 $k=0,\ldots,E-1$ 的参数,因为它们是计算梯度的位置。$E=1$ 时 $q=0$,因此 $M_t=r_t=0$。

展开本地更新,将梯度均值与噪声分开。使用 $|a+c|^2\le2|a|^2+2|c|^2$、$|\sum_{s<k}a_s|^2\le k\sum_{s<k}|a_s|^2$,以及噪声的鞅差性质,得到

$$ D_{t,k}\le \frac{2\eta^2k}{N}\sum_i\sum_{s<k} \mathbb E[\|\nabla f_i(w_{i,t}^{(s)})\|^2\mid\mathcal F_t] +\frac{2\eta^2k\sigma^2}{b}. $$

再用光滑性和上一节的异质性界:

$$ \frac1N\sum_i \mathbb E[\|\nabla f_i(w_{i,t}^{(s)})\|^2\mid\mathcal F_t] \le2(\|h_t\|^2+\zeta^2)+2L^2D_{t,s}. $$ $$ D_{t,k}\le4\eta^2k^2(\|h_t\|^2+\zeta^2) +4\eta^2kL^2\sum_{s<k}D_{t,s} +\frac{2\eta^2k\sigma^2}{b}. $$

对 $0\le k\le q$ 取最大值,得到

$$ M_t\le4\eta^2q^2(\|h_t\|^2+\zeta^2) +4\eta^2L^2q^2M_t+\frac{2\eta^2q\sigma^2}{b}. $$

以后统一使用步长条件 $0<\eta LE\le1/8$。于是 $4\eta^2L^2q^2\le1/16$,移项并将常数放宽,得到

$$ M_t\le8\eta^2q^2(\|h_t\|^2+\zeta^2) +\frac{4\eta^2q\sigma^2}{b}. $$

由 Jensen 不等式和局部目标的光滑性:

$$ \begin{aligned} \mathbb E[\|r_t\|^2\mid\mathcal F_t] &\le\frac{L^2}{E}\sum_{k=0}^{E-1}D_{t,k}\\ &\le8L^2\eta^2q^2(\|h_t\|^2+\zeta^2) +\frac{4L^2\eta^2q\sigma^2}{b}. \end{aligned} $$

这也解释了为什么“把一轮的 $NEb$ 个样本看成一个大 batch”不够:噪声平均确实有 $1/(NE)$ 缩减,但梯度是在不同的本地参数上计算的,还必须控制 $r_t$。

0x04 A Non-convex Convergence Bound

One-round Descent

由 $f$ 的光滑性和 $w_{t+1}=w_t-\alpha v_t$,

$$ f(w_{t+1})\le f(w_t)-\alpha\langle h_t,v_t\rangle +\frac{L\alpha^2}{2}\|v_t\|^2. $$

先给定 $\mathcal F_t$ 取期望。$h_t$ 可测,$\mathbb E[e_t\mid\mathcal F_t]=0$,因此 $h_t$ 与噪声的内积消失。对于漂移项使用

$$ |\langle h_t,r_t\rangle|\le\frac14\|h_t\|^2+\|r_t\|^2, \qquad \|h_t+r_t+e_t\|^2\le4\|h_t\|^2+4\|r_t\|^2+2\|e_t\|^2. $$

这里没有假设 $r_t$ 与 $e_t$ 独立。记 $R_t=\mathbb E[|r_t|^2\mid\mathcal F_t]$,便有

$$ \mathbb E[f(w_{t+1})\mid\mathcal F_t] \le f(w_t) -\alpha\left(\frac34-2L\alpha\right)\|h_t\|^2 +\alpha(1+2L\alpha)R_t +\frac{L\alpha^2\sigma^2}{bNE}. $$

因为 $L\alpha=L\eta E\le1/8$,第一项的下降系数至少为 $1/2$,且 $1+2L\alpha\le5/4$。代入漂移界后,梯度项新增的系数至多为

$$ 10L^2\eta^2q^2\le\frac{10}{64}<\frac14. $$

因此可以保守地留下 $\alpha/4$ 的下降量。定义

$$ A=\frac{L\eta\sigma^2}{bN} +\frac{5L^2\eta^2(E-1)\sigma^2}{b} +10L^2\eta^2(E-1)^2\zeta^2. $$

得到整个证明的核心递推:

$$ \mathbb E[f(w_{t+1})\mid\mathcal F_t] \le f(w_t)-\frac\alpha4\|\nabla f(w_t)\|^2+\alpha A. $$

Stationarity Guarantee

defination

在前述假设、全参与等权聚合及 $0<\eta\le1/(8LE)$ 下,运行 $T\ge1$ 个通信轮,有
$$ \begin{aligned} \frac1T\sum_{t=0}^{T-1}\mathbb E\|\nabla f(w_t)\|^2 &\le\frac{4\Delta_0}{\eta ET}+4A\\ &=\frac{4\Delta_0}{\eta ET} +\frac{4L\eta\sigma^2}{bN} +\frac{20L^2\eta^2(E-1)\sigma^2}{b} +40L^2\eta^2(E-1)^2\zeta^2. \end{aligned} $$

证明只需对核心递推取总期望并在 $t=0,\ldots,T-1$ 上求和:函数值项相消,再用 $\mathbb E[f(w_T)]\ge f_{\inf}$。若独立地从这 $T$ 个轮初参数中均匀选一个 $w_R$,则左侧恰好等于 $\mathbb E|\nabla f(w_R)|^2$。

这是平均驻点界,不是最后一轮 $w_T$ 的保证,也不是到全局最优解的保证。梯度小也不等于测试准确率高。

Reading the Four Terms

项含义从本界能看出的影响
$4\Delta_0/(\eta ET)$有限轮数的优化项在其他条件固定时,更多轮数让这一项减小
$4L\eta\sigma^2/(bN)$聚合后的梯度噪声更多独立客户端或更大 batch 可以减小此项
$20L^2\eta^2(E-1)\sigma^2/b$噪声经过本地路径累积的漂移不能仅靠聚合中的 $1/N$ 缩减来消除
$40L^2\eta^2(E-1)^2\zeta^2$异质目标导致的漂移多做本地步可能增大此项

固定步长下,后三项不会因为增加 $T$ 而自动消失。它们是上界中的残差,不是算法实际误差的精确下界。即使 $\zeta=0$,随机噪声仍可能通过不同的本地路径产生漂移;即使 $\sigma=0$,异质性仍可能造成偏差。

若固定 $N,E,b$,并为长度为 $T$ 的一次运行选择 $\eta=c/\sqrt T$,其中 $0<c\le1/(8LE)$,本界给出 $O(T^{-1/2})+O(T^{-1})$。这里每次运行内部仍使用固定步长;这不是对在线变化的 $\eta_t$ 的证明。$E$ 或 $N$ 随 $T$ 增长时,必须重新代入全部项,不能继续把它们藏在常数中。

$E=1$ 时漂移项完全消失,算法恢复全参与的同步 mini-batch SGD。上述通用界的常数较松;直接使用 Part 3 的无偏梯度证明,可以得到更好的常数。若目标梯度平方范数为 $\epsilon$,本界只能在 $\epsilon>4A$ 时通过增大 $T$ 给出保证。

0x05 Strongly Convex Objectives

若再假设 $f$ 为 $\mu$-强凸且存在最优解 $w^\ast$,令 $f^\ast=f(w^\ast)$、$\mathcal E_t=\mathbb E[f(w_t)]-f^\ast$。利用 PL 不等式

$$ \|\nabla f(w)\|^2\ge2\mu(f(w)-f^\ast), $$

核心递推变为

$$ \mathcal E_{t+1}\le \left(1-\frac{\alpha\mu}{2}\right)\mathcal E_t+\alpha A. $$

由于 $\mu\le L$,步长条件保证 $\rho=1-\alpha\mu/2\in(0,1)$。展开几何级数:

defination

$$ \mathcal E_T\le\rho^T\mathcal E_0 +\frac{2A}{\mu}(1-\rho^T). $$

固定本地步长时,瞬态项几何收缩,但本界通常只保证 $\limsup_T\mathcal E_T\le2A/\mu$。只有残差消失,才从这个递推得到到精确最优值的几何收敛。

若 $0<2A/\mu<\epsilon<\mathcal E_0$,足够的轮数为

$$ T\ge\left\lceil \frac{\log\left((\mathcal E_0-2A/\mu)/(\epsilon-2A/\mu)\right)}{-\log\rho} \right\rceil. $$

当 $A=0$ 时直接使用 $\rho^T\mathcal E_0$;当目标低于残差上界时,上面的轮数公式不能提供保证。关于递减步长的更细致强凸分析,可以参阅 Li et al., ICLR 2020,但该文的假设、时间索引与聚合方案需要逐项对齐,不能只拿一个 $O(1/T)$ 结论替换这里的定理。

0x06 A Deterministic Drift Example

前面的残差只是上界。下面用一个无噪声例子说明,本地漂移也可能真实存在,而不是证明技巧凭空产生的。

取两个等权客户端,$0<a<1/\sqrt2$:

$$ f_1(w)=\frac12w^2+a(\sin w+\cos w),\qquad f_2(w)=\frac12w^2-a(\sin w+\cos w). $$ $$ f(w)=\frac12w^2,\qquad w^\ast=0,\qquad \nabla f_{1,2}(w)=w\pm a(\cos w-\sin w). $$

两个局部目标都 $L$-光滑且强凸,可以取 $L=1+\sqrt2a$、局部强凸参数 $1-\sqrt2a$。异质性界在全空间成立:

$$ \frac12\sum_{i=1}^2|\nabla f_i(w)-\nabla f(w)|^2 =a^2(\cos w-\sin w)^2\le2a^2. $$

从全局最优点 $w_t=0$ 开始,每个客户端做两步确定性 GD。第一步分别到达 $-\eta a$ 和 $\eta a$;把第二步的模型平均后,恰好得到

$$ w_{t+1}=-\eta a\sin(\eta a). $$

对于足够小的正步长,这不等于零:FedAvg 甚至会从全局最优点离开。相反,$E=1$ 时两个局部梯度在 $w=0$ 完全抵消。

如果 $0<\eta\le1/L$,每个局部 GD 映射都是收缩映射;两步映射的平均仍收缩。因此固定步长的这一 FedAvg 映射有唯一吸引不动点,而由于它不把零映射到零,不动点也不等于全局最优解。这个例子满足主定理的异质性假设;再取 $\eta\le1/(16L)$,也满足 $E=2$ 时的主定理步长限制。

下面的 Python 代码可以验证一轮公式与不动点,直接运行不需要第三方库:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
from math import cos, sin, sqrt

a, eta = 0.25, 0.04
assert eta * 2 * (1 + sqrt(2) * a) <= 1 / 8


def fedavg_round(w, steps=2):
    local = []
    for sign in (1, -1):
        x = w
        for _ in range(steps):
            x -= eta * (x + sign * a * (cos(x) - sin(x)))
        local.append(x)
    return sum(local) / 2


assert abs(fedavg_round(0) + eta * a * sin(eta * a)) < 1e-12
assert abs(fedavg_round(0, steps=1)) < 1e-12
w = 0.0
for _ in range(2000):
    w = fedavg_round(w)
assert abs(w) > 1e-6
assert abs(fedavg_round(w) - w) < 1e-12
print(f"FedAvg fixed point: {w:.8f}; global optimum: 0")

这是解析例子的数值核对,不是用一组实验替代一般收敛证明。

0x07 Partial Participation and Aggregation Weights

Client Sampling Adds Another Variance Term

现在只看 $E=1$,仍优化等权目标,每轮从 $N>1$ 个客户端中均匀、无放回选 $m$ 个。选择发生在 batch 采样之前,且与该轮新样本噪声独立。若聚合方向为 $\hat g_t=m^{-1}\sum_{i\in S_t}g_i(w_t)$,则

$$ \mathbb E[\hat g_t\mid\mathcal F_t]=\nabla f(w_t), $$ $$ \begin{aligned} \mathbb E[\|\hat g_t-\nabla f(w_t)\|^2\mid\mathcal F_t] &\le\frac{\sigma^2}{bm} +\frac{N-m}{m(N-1)}\cdot\frac1N\sum_i \|\nabla f_i(w_t)-\nabla f(w_t)\|^2\\ &\le\frac{\sigma^2}{bm}+\frac{N-m}{m(N-1)}\zeta^2. \end{aligned} $$

第二项是有限总体抽样方差,与客户端内部的随机梯度噪声不同。下面把这个系数推出来。给定 $\mathcal F_t$,记 $a_i=\nabla f_i(w_t)-\nabla f(w_t)$、$I_i=\mathbf 1_{{i\in S_t}}$,此时 $\sum_i a_i=0$。均匀无放回抽样满足

$$ \mathbb E[I_i\mid\mathcal F_t]=\frac mN,\qquad \mathbb E[I_iI_j\mid\mathcal F_t]=\frac{m(m-1)}{N(N-1)},\quad i\ne j. $$ $$ \sum_{i\ne j}\langle a_i,a_j\rangle=-\sum_i\|a_i\|^2. $$

将平方展开,对角项与非对角项分别取期望:

$$ \begin{aligned} \mathbb E\left[\left\|\frac1m\sum_i I_i a_i\right\|^2\middle|\mathcal F_t\right] &=\frac1{m^2}\left(\frac mN-\frac{m(m-1)}{N(N-1)}\right)\sum_i\|a_i\|^2\\ &=\frac{N-m}{m(N-1)}\cdot\frac1N\sum_i\|a_i\|^2. \end{aligned} $$

样本梯度噪声在给定 $S_t$ 后均值为零,因此与这里的客户端采样偏差的交叉项消失,得到前述总方差界。$m=N$ 时采样项为零;有放回采样时,异质性系数变为 $1/m$,不能沿用无放回的有限总体修正。

于是对 $0<\eta\le1/L$,直接套用 Part 3 的一步下降,有

$$ \frac1T\sum_{t<T}\mathbb E\|\nabla f(w_t)\|^2 \le\frac{2\Delta_0}{\eta T} +L\eta\left(\frac{\sigma^2}{bm} +\frac{N-m}{m(N-1)}\zeta^2\right). $$

$N=1$ 时只能取 $m=1$,不计算含 $N-1$ 的公式。对 $E>1$,不能只把主定理的 $N$ 改成 $m$:还需要同时控制本地漂移与客户端选择。

Extending the Bound to Multiple Local Steps

仍使用每轮独立于新样本噪声的均匀无放回选择,所有入选客户端从 $w_t$ 开始、执行相同的 $E$ 步,服务器等权平均这 $m$ 个本地模型。定义

$$ \chi_m=\frac{N-m}{m(N-1)},\qquad u_t=\frac1m\sum_{i\in S_t}\nabla f_i(w_t)-h_t, $$ $$ \begin{aligned} r_t^{(m)}&=\frac1{mE}\sum_{i\in S_t}\sum_{k=0}^{E-1} \bigl[\nabla f_i(w_{i,t}^{(k)})-\nabla f_i(w_t)\bigr],\\ e_t^{(m)}&=\frac1{mE}\sum_{i\in S_t}\sum_{k=0}^{E-1}\varepsilon_{i,t}^{(k)},\\ v_t^{(m)}&=h_t+u_t+r_t^{(m)}+e_t^{(m)}. \end{aligned} $$

$u_t$ 是在共同轮初参数上的客户端采样误差,不能归入本地样本噪声。把 $z_t=u_t+e_t^{(m)}$ 看作一轮的合成噪声,由上一小节的抽样方差与鞅差性质,有

$$ \mathbb E[z_t\mid\mathcal F_t]=0,\qquad \mathbb E[\|z_t\|^2\mid\mathcal F_t] \le\chi_m\zeta^2+\frac{\sigma^2}{bmE}. $$

这里用到了 $\mathbb E[e_t^{(m)}\mid\mathcal F_t,S_t]=0$,从而 $u_t$ 与 $e_t^{(m)}$ 的交叉项为零。漂移 $r_t^{(m)}$ 与 $z_t$ 的交叉项仍不假设为零。

还需要检查漂移界是否保留。定义入选客户端的期望平均位移

$$ D_{t,k}^{(m)}=\mathbb E\left[\frac1m\sum_{i\in S_t} \|w_{i,t}^{(k)}-w_t\|^2\middle|\mathcal F_t\right]. $$

轮初的局部梯度不依赖 $S_t$,且每个客户端入选概率都是 $m/N$,所以

$$ \mathbb E\left[\frac1m\sum_{i\in S_t}\|\nabla f_i(w_t)\|^2 \middle|\mathcal F_t\right] =\frac1N\sum_i\|\nabla f_i(w_t)\|^2 \le\|h_t\|^2+\zeta^2. $$

因此,对入选客户端先展开路径、再对选择和噪声取期望,$D_{t,k}^{(m)}$ 满足与 0x03 相同的递推。由 Jensen 得到

$$ \mathbb E[\|r_t^{(m)}\|^2\mid\mathcal F_t] \le8L^2\eta^2(E-1)^2(\|h_t\|^2+\zeta^2) +\frac{4L^2\eta^2(E-1)\sigma^2}{b}. $$

现在才可以重复 0x04 的下降证明:用 $z_t$ 替代 $e_t$,保留其与漂移的相关性,将噪声方差替换为 $\chi_m\zeta^2+\sigma^2/(bmE)$。定义

$$ A_m=\frac{L\eta\sigma^2}{bm} +L\eta E\chi_m\zeta^2 +\frac{5L^2\eta^2(E-1)\sigma^2}{b} +10L^2\eta^2(E-1)^2\zeta^2. $$

defination

在上述均匀无放回部分参与模型与 $0<\eta\le1/(8LE)$ 下,
$$ \frac1T\sum_{t<T}\mathbb E\|\nabla f(w_t)\|^2 \le\frac{4\Delta_0}{\eta ET}+4A_m. $$
若再满足强凸条件,则同一个 $\rho=1-\eta E\mu/2$ 给出
$$ \mathcal E_T\le\rho^T\mathcal E_0+\frac{2A_m}{\mu}(1-\rho^T). $$

$m=N$ 时 $\chi_m=0$、$A_m=A$,恢复全参与结果。$E=1$ 时也恢复“噪声加采样误差”的结构,但这里的步长限制和常数更保守,应优先使用前一小节的精确无漂移分析。$N=m=1$ 时单独定义 $\chi_m=0$。

多出的 $L\eta E\chi_m\zeta^2$ 是客户端采样项,与原来的 $\eta^2(E-1)^2\zeta^2$ 本地漂移项具有不同的尺度。增大 $m$ 可以减少采样误差,却不能据此宣称本地漂移也按 $1/m$ 消失。依赖训练状态的非均匀选择、同一批客户端连续多轮参与、加权聚合和不同的本地步数,都不直接被这个扩展覆盖。

Weighting Can Change the Target

如果目标是 $f=\sum_i p_i f_i$,但每轮均匀选择一个客户端并完全采用它的更新,那么 $E=1$ 时的期望梯度是 $N^{-1}\sum_i\nabla f_i$,而不一定是 $\nabla f$。

例如 $p_1=0.9,p_2=0.1$,$f_1(w)=w^2/2$、$f_2(w)=(w-2)^2/2$。加权目标的最优点为 $0.2$,等权目标的最优点为 $1$。这里的问题不是“收敛得慢”,而是优化了另一个目标。

若客户端 $i$ 的入选概率为 $\pi_i>0$,在同一个 $w$ 上计算梯度时,无偏的 Horvitz–Thompson 形式为

$$ \hat g(w)=\sum_{i\in S}\frac{p_i}{\pi_i}g_i(w), \qquad\mathbb E[\hat g(w)]=\sum_i p_i\nabla f_i(w). $$

这要求选择机制与新样本噪声满足相应条件。随机地除以“入选客户端权重之和”得到的比率估计,一般不再严格无偏;重要性校正也可能增大方差。实际的 FedAvg 聚合方案需要具体分析,而不是默认上述估计器与所有实现等价。

按在线率、完成速度或损失选择客户端,同样可能改变目标和噪声条件。Part 4 对“选择最快的客户端”的提醒在这里仍适用。

0x08 Communication, Drift, and Corrections

Local Work Is Not Free Progress

在总本地更新次数 $H=ET$ 固定时,主界的优化项是 $4\Delta_0/(\eta H)$。增加 $E$ 可以减少通信轮数,但漂移项增加,而且步长还必须满足 $\eta LE\le1/8$。因此,比较不同 $E$ 时不能只比较通信轮数或只保留界中的第一项。

若把有效步长 $\alpha=\eta E$ 固定,则代入 $\eta=\alpha/E$ 后,异质性项中的 $\eta^2(E-1)^2$ 趋近 $\alpha^2$,也不会因为本地步数无限增加而自动消失。对于 wall-clock time,还要计算下载、上传、本地训练以及等待的耗时;可与 效率分析 中的思路结合,但该文的同质 worker 时间模型不能直接代表移动设备。

Choosing Parameters from the Bound

步长、本地步数和参与数量不是三个可以独立无限增大的参数。把一次实验看作固定 $T,E,m,b$ 的运行,部分参与的非凸界具有以下结构:

$$ B(\eta)=\frac{c_0}{\eta}+c_1\eta+c_2\eta^2, \qquad 0<\eta\le\frac1{8LE}, $$ $$ \begin{aligned} c_0&=\frac{4\Delta_0}{ET},\\ c_1&=\frac{4L\sigma^2}{bm}+4LE\chi_m\zeta^2,\\ c_2&=\frac{20L^2(E-1)\sigma^2}{b} +40L^2(E-1)^2\zeta^2. \end{aligned} $$

若 $c_0>0$ 且 $c_1+c_2>0$,无约束最小值由下式的唯一正根决定:

$$ -c_0+c_1\eta^2+2c_2\eta^3=0. $$

左侧在正半轴严格递增,所以可以用二分法求根,再和 $1/(8LE)$ 取较小值。$c_2=0,c_1>0$ 时得到 $\eta=\sqrt{c_0/c_1}$;$c_1=c_2=0$ 时,界随步长增加而下降,在允许区间内取上限即可。$c_0=0$ 时不能套用唯一正根的说法,需单独检查从最优值出发时的残差。

这解释了为什么只按 $\eta\propto1/\sqrt T$ 调参不总是充分:异质性和本地步数大时,二次项不能忽略。上述公式用于理解上界,不是一个可以直接照抄的实际最优学习率;$L,\Delta_0,\sigma^2,\zeta^2$ 通常未知,而且理论允许的步长往往较保守。

对实际训练,可以让单步本地更新作为基线,再增加 $E$ 并检查训练曲线。比较时应固定一种预算,例如总本地更新次数或总耗时,并同时报告通信轮数;如果一组实验既增加本地计算量又增加通信轮数,只比较最终准确率就无法说明本地更新是否更有效。

Distinguishing Error Sources in Experiments

观察或修改本模型中主要影响的量不能据此直接推出的结论
增大每步 batch客户端内部噪声 $\sigma^2/b$客户端分布变得更一致
增大每轮参与数 $m$聚合噪声与采样系数 $\chi_m$本地模型漂移被消除
增大本地步数 $E$通信频率、有效步长与漂移项收敛速度一定提高
减小本地学习率 $\eta$噪声残差和漂移项,同时减慢优化项下降在固定有限预算下必然更准确
调整聚合权重被优化的目标 $f$只是一个不影响结论的实现细节

另外,标签分布不一致、样本量不一致和最优点不一致并不是同一件事。$\zeta^2$ 度量的是参数位置上的梯度差异,而不是某个数据划分参数。用 Dirichlet 分布划分数据时,浓度参数可以控制划分方式,却不能无条件写成 $\zeta^2$ 的解析表达式。

可以在服务器参数 $w_t$ 上估计各客户端的梯度,再观察梯度之间的差异;但用一个有限 batch 计算出的差异还包含样本噪声,不能直接当作真实异质性。重复采样、更大的诊断 batch 或显式估计噪声,都有助于区分两者。在真实跨设备训练中,这类额外诊断还会带来通信和数据访问成本。

FedProx and SCAFFOLD

FedProx 在一轮内对本地目标加入近端项:

$$ f_i(w)+\frac\lambda2\|w-w_t\|^2,\qquad\lambda\ge0. $$

它为本地参数偏离服务器参数增加代价。近端项不能仅凭形式就保证所有问题中精确消除漂移;其分析还涉及目标条件与本地求解精度。它也不是本文未经修改的 FedAvg 更新式。

SCAFFOLD 使用控制变量修正方向:

$$ w_i\leftarrow w_i-\eta\bigl(g_i(w_i)-c_i+c\bigr). $$

如果在轮初理想地取 $c_i=\nabla f_i(w_t)$、$c=\nabla f(w_t)$,修正后的梯度均值在 $w_i=w_t$ 时恰好等于全局梯度;本地参数移动后,仍有由位置变化产生的误差。实际算法维护的是控制变量估计,还要分析估计误差、更新方式与采样。这里解释修正机制,不把它当作已证明的 SCAFFOLD 收敛定理。

0x09 What the Analysis Guarantees

本文的主结论建立在全参与、等权、同步、相同本地步数、条件无偏噪声和统一异质性界之上。非凸结论是平均驻点保证;强凸结论是固定步长下向误差上界平台几何收缩。部分参与的扩展覆盖每轮均匀无放回抽样、相同本地步数的等权模型;加权采样部分讨论目标一致性,不冒充所有 FedAvg 实现的全覆盖定理。

训练不同的神经网络时,还需要检查光滑性、采样、batch 相关性、参与机制和优化器是否符合模型。Adam、动量、本地不同步数、异步服务器和非光滑激活都不会因为算法也叫 FedAvg 就自动满足本篇假设。验证准确率、泛化、公平性和隐私也不由这里的优化界直接保证。

对一份实际实验,我会先记录目标权重、$N,m,E,b,\eta$,再分别观察客户端采样、本地漂移和通信耗时。把这些量区分清楚,才知道是在改进优化方向、减少随机噪声,还是用更多本地计算换取更少的通信。

Reference