联邦平均算法(Federated Averaging Algorithm,FedAvg)
设一共有\(K\)个客户机,
中心服务器初始化模型参数,执行若干轮(round),每轮选取至少1个至多\(K\)个客户机参与训练,接下来每个被选中的客户机同时在自己的本地根据服务器下发的本轮(\(t\)轮)模型\(w_t\)用自己的数据训练自己的模型\(w^k_{t+1}\),上传回服务器。服务器将收集来的各客户机的模型根据各方样本数量用加权平均的方式进行聚合,得到下一轮的模型\(w_{t+1}\):
\[\begin{aligned}
& \qquad w_{t+1} \leftarrow \sum^K_{k=1} \frac{n_k}{n} w^k_{t+1} \qquad\qquad //n_k为客户机k上的样本数量,n为所有被选中客户机的总样本数量\\
\end{aligned}
\]
【伪代码】
\[\begin{aligned}
& 算法1:Federated\ Averaging算法(FedAvg)。 \\
& K个客户端编号为k;B,E,\eta分别代表本地的minibatch\ size,epochs,学习率learning\ rate \\
& \\
& 服务器执行:\\
& \quad 初始化w_0 \\
& \quad for \ 每轮t=1,2,...,do \\
& \qquad m \leftarrow max(C \cdot K,1) \qquad\qquad //C为比例系数 \\
& \qquad S_t \leftarrow (随机选取m个客户端) \\
& \qquad for \ 每个客户端k \in S_t 同时\ do \\
& \qquad \qquad w^k_{t+1} \leftarrow 客户端更新(k,w_t) \\
& \qquad w_{t+1} \leftarrow \sum^K_{k=1} \frac{n_k}{n} w^k_{t+1} \qquad\qquad //n_k为客户机k上的样本数量,n为所有被选中客户机的总样本数量\\
& \\
& 客户端更新(k,w): \qquad \triangleright 在客户端k上运行 \\
& \quad \beta \leftarrow (将P_k分成若干大小为B的batch) \qquad\qquad //P_k为客户机k上数据点的索引集,P_k大小为n_k \\
& \quad for\ 每个本地的epoch\ i(1\sim E) \ do \\
& \qquad for\ batch\ b \in \beta \ do \\
& \qquad \qquad w \leftarrow w-\eta \triangledown l(w;b) \qquad\qquad //\triangledown 为计算梯度,l(w;b)为损失函数\\
& \quad 返回w给服务器
\end{aligned}
\]
为了增加客户机计算量,可以在中心服务器做聚合(加权平均)操作前在每个客户机上多迭代更新几次。计算量由三个参数决定:
- \(C\),每一轮(round)参与计算的客户机比例。
- \(E(epochs)\),每一轮每个客户机投入其全部本地数据训练一遍的次数。
- \(B(batch size)\),用于客户机更新的batch大小。\(B=\infty\)表示batch为全部样本,此时就是full-batch梯度下降了。
当\(E=1\ B=\infty\)时,对应的就是FedSGD,即每一轮客户机一次性将所有本地数据投入训练,更新模型参数。
对于一个有着\(n_k\)个本地样本的客户机\(k\)来说,每轮的本地更新次数为\(u_k=E\cdot \frac{n_k}{B}\)。
参考文献:
- H. B. McMahan, E. Moore, D. Ramage, S. Hampson, and B. A. Y. Arcas, “Communication-efficient learning of deep networks from decentralized data,” in Proc. AISTATS, 2016, pp. 1273–1282.