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$ 个客户端的局部目标为
如果希望每个样本具有相同权重,通常取 $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:
$\eta$ 是固定的本地步长,服务器直接平均模型,不额外乘服务器学习率。定义一轮的有效步长 $\alpha=\eta E$,以及聚合方向
这里的 $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$ 步采样前所有客户端的参数和历史。写
$b$ 是每步的 batch 大小,$\sigma^2$ 是单样本噪声方差的统一上界。同一步不同客户端的噪声还要求条件互不相关。同一客户端不同步之间不要求无条件独立,条件无偏使噪声增量构成鞅差序列。
batch 内的 $1/b$ 缩减需要独立采样或相应的协方差条件。直接遍历同一个随机排列、根据损失选择样本,或采样过程受到掉线机制影响时,都需要重新验证上述条件;不能仅凭“用了 SGD”就认为条件成立。
Bounded Gradient Heterogeneity
采用以下统一异质性上界:
$\sigma^2$ 描述一个客户端内部的随机噪声,$\zeta^2$ 描述客户端之间的目标差异。增大 batch 能减少前者,却不能自动消除后者。由 $\sum_i(\nabla f_i-\nabla f)=0$,有
这是一个较强的假设,不由“非 IID 数据”这个描述自动推出。它要求差异在所有参数上有统一界,而不仅是最优点附近有界。若只在一个区域成立,还必须证明所有相关迭代留在该区域。本篇不把“全局梯度一致有界”作为额外假设。
Noise and Local Drift
令 $h_t=\nabla f(w_t)$,把聚合方向准确地拆成
$r_t$ 是本地路径带来的漂移,$e_t$ 是平均随机噪声。条件无偏和互不相关性给出
不同步的噪声交叉项通过迭代条件期望消失。但 $r_t$ 取决于本地参数,而本地参数又取决于之前的噪声,所以一般不能断言 $\mathbb E\langle r_t,e_t\rangle=0$。后面的证明会保留这个相关性,通过范数不等式控制它。
0x03 Bounding Local Drift
这一节给出后面定理需要的关键估计。令 $q=E-1$,并定义给定轮初历史下的平均位移:
这里只需要控制 $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$,以及噪声的鞅差性质,得到
再用光滑性和上一节的异质性界:
对 $0\le k\le q$ 取最大值,得到
以后统一使用步长条件 $0<\eta LE\le1/8$。于是 $4\eta^2L^2q^2\le1/16$,移项并将常数放宽,得到
由 Jensen 不等式和局部目标的光滑性:
这也解释了为什么“把一轮的 $NEb$ 个样本看成一个大 batch”不够:噪声平均确实有 $1/(NE)$ 缩减,但梯度是在不同的本地参数上计算的,还必须控制 $r_t$。
0x04 A Non-convex Convergence Bound
One-round Descent
由 $f$ 的光滑性和 $w_{t+1}=w_t-\alpha v_t$,
先给定 $\mathcal F_t$ 取期望。$h_t$ 可测,$\mathbb E[e_t\mid\mathcal F_t]=0$,因此 $h_t$ 与噪声的内积消失。对于漂移项使用
这里没有假设 $r_t$ 与 $e_t$ 独立。记 $R_t=\mathbb E[|r_t|^2\mid\mathcal F_t]$,便有
因为 $L\alpha=L\eta E\le1/8$,第一项的下降系数至少为 $1/2$,且 $1+2L\alpha\le5/4$。代入漂移界后,梯度项新增的系数至多为
因此可以保守地留下 $\alpha/4$ 的下降量。定义
得到整个证明的核心递推:
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 不等式
核心递推变为
由于 $\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$,足够的轮数为
当 $A=0$ 时直接使用 $\rho^T\mathcal E_0$;当目标低于残差上界时,上面的轮数公式不能提供保证。关于递减步长的更细致强凸分析,可以参阅 Li et al., ICLR 2020,但该文的假设、时间索引与聚合方案需要逐项对齐,不能只拿一个 $O(1/T)$ 结论替换这里的定理。
0x06 A Deterministic Drift Example
前面的残差只是上界。下面用一个无噪声例子说明,本地漂移也可能真实存在,而不是证明技巧凭空产生的。
取两个等权客户端,$0<a<1/\sqrt2$:
两个局部目标都 $L$-光滑且强凸,可以取 $L=1+\sqrt2a$、局部强凸参数 $1-\sqrt2a$。异质性界在全空间成立:
从全局最优点 $w_t=0$ 开始,每个客户端做两步确定性 GD。第一步分别到达 $-\eta a$ 和 $\eta a$;把第二步的模型平均后,恰好得到
对于足够小的正步长,这不等于零:FedAvg 甚至会从全局最优点离开。相反,$E=1$ 时两个局部梯度在 $w=0$ 完全抵消。
如果 $0<\eta\le1/L$,每个局部 GD 映射都是收缩映射;两步映射的平均仍收缩。因此固定步长的这一 FedAvg 映射有唯一吸引不动点,而由于它不把零映射到零,不动点也不等于全局最优解。这个例子满足主定理的异质性假设;再取 $\eta\le1/(16L)$,也满足 $E=2$ 时的主定理步长限制。
下面的 Python 代码可以验证一轮公式与不动点,直接运行不需要第三方库:
| |
这是解析例子的数值核对,不是用一组实验替代一般收敛证明。
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)$,则
第二项是有限总体抽样方差,与客户端内部的随机梯度噪声不同。下面把这个系数推出来。给定 $\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$。均匀无放回抽样满足
将平方展开,对角项与非对角项分别取期望:
样本梯度噪声在给定 $S_t$ 后均值为零,因此与这里的客户端采样偏差的交叉项消失,得到前述总方差界。$m=N$ 时采样项为零;有放回采样时,异质性系数变为 $1/m$,不能沿用无放回的有限总体修正。
于是对 $0<\eta\le1/L$,直接套用 Part 3 的一步下降,有
$N=1$ 时只能取 $m=1$,不计算含 $N-1$ 的公式。对 $E>1$,不能只把主定理的 $N$ 改成 $m$:还需要同时控制本地漂移与客户端选择。
Extending the Bound to Multiple Local Steps
仍使用每轮独立于新样本噪声的均匀无放回选择,所有入选客户端从 $w_t$ 开始、执行相同的 $E$ 步,服务器等权平均这 $m$ 个本地模型。定义
$u_t$ 是在共同轮初参数上的客户端采样误差,不能归入本地样本噪声。把 $z_t=u_t+e_t^{(m)}$ 看作一轮的合成噪声,由上一小节的抽样方差与鞅差性质,有
这里用到了 $\mathbb E[e_t^{(m)}\mid\mathcal F_t,S_t]=0$,从而 $u_t$ 与 $e_t^{(m)}$ 的交叉项为零。漂移 $r_t^{(m)}$ 与 $z_t$ 的交叉项仍不假设为零。
还需要检查漂移界是否保留。定义入选客户端的期望平均位移
轮初的局部梯度不依赖 $S_t$,且每个客户端入选概率都是 $m/N$,所以
因此,对入选客户端先展开路径、再对选择和噪声取期望,$D_{t,k}^{(m)}$ 满足与 0x03 相同的递推。由 Jensen 得到
现在才可以重复 0x04 的下降证明:用 $z_t$ 替代 $e_t$,保留其与漂移的相关性,将噪声方差替换为 $\chi_m\zeta^2+\sigma^2/(bmE)$。定义
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 形式为
这要求选择机制与新样本噪声满足相应条件。随机地除以“入选客户端权重之和”得到的比率估计,一般不再严格无偏;重要性校正也可能增大方差。实际的 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$ 的运行,部分参与的非凸界具有以下结构:
若 $c_0>0$ 且 $c_1+c_2>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 在一轮内对本地目标加入近端项:
它为本地参数偏离服务器参数增加代价。近端项不能仅凭形式就保证所有问题中精确消除漂移;其分析还涉及目标条件与本地求解精度。它也不是本文未经修改的 FedAvg 更新式。
SCAFFOLD 使用控制变量修正方向:
如果在轮初理想地取 $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
- McMahan, B. et al. Communication-Efficient Learning of Deep Networks from Decentralized Data, AISTATS 2017. FedAvg 的原始算法与实验。
- Li, X., Huang, K., Yang, W., Wang, S. and Zhang, Z. On the Convergence of FedAvg on Non-IID Data, ICLR 2020. 非 IID FedAvg 的强凸分析与固定步长问题。
- Karimireddy, S. P. et al. SCAFFOLD: Stochastic Controlled Averaging for Federated Learning, ICML 2020, especially Sections 2–4. 梯度异质性、本地漂移与控制变量方法;该文的更一般条件和更紧界不等于本文的简化界。
- Li, T. et al. Federated Optimization in Heterogeneous Networks, MLSys 2020. FedProx 与统计、系统异质性。
- Woodworth, B., Patel, K. K. and Srebro, N. Minibatch vs Local SGD for Heterogeneous Distributed Learning, NeurIPS 2020. 异质目标下的 local SGD 与 mini-batch SGD 比较。

