对FedAvg中模型聚合过程的理解

问题

联邦学习原始论文中给出的FedAvg的算法框架为：

参数介绍：$K$表示客户端的个数，$B$表示每一次本地更新时的数据量，$E$表示本地更新的次数，$\eta$表示学习率。

首先是服务器执行以下步骤：

1. 初始化参数。
2. 对每一轮通信来说：首先计算出$m=max(C\cdot K,1)$，然后随机选择m个客户端，对这m个客户端做如下操作（所有客户端并行执行）：更新本地的$w_t^k$得到$w_{t+1}^k$。所有客户端更新结束后，将$w_{t+1}^k$传到服务器，服务器整合所有$w_{t+1}^k$得到最新的全局参数$w_{t+1}$。

对每一个本地客户端来说，要做的就是更新本地参数，具体来讲：

1. 把自己的数据集按照参数B分成若干个块，每一块大小都为B。
2. 对每一块数据，需要进行E轮更新：算出该块数据损失的梯度，然后进行梯度下降更新，得到新的本地$w$。
3. 更新完后$w$将被传送到中央服务器，服务器整合所有客户端计算出的$w$，得到最新的全局模型参数$w_{t+1}$
4. 客户端收到服务器发送的最新全局参数模型参数，进行下一次更新。

我们仔细观察server的最后一步：
$$w_{t+1}=\sum _{k=1}^K\frac{n_k}{n}{w}_{t+1}^k$$
也就是说，虽然我们只是对$m$个客户端进行本地训练更新得到了其对应的$w_{t+1}^k$，但最终我们却对所有$K$个客户端进行了聚合。

聚合

那么针对聚合，就有以下两种情况。

1. 聚合所有客户端

服务器端每次将新的全局模型发送给全部客户端，并且聚合全部客户端的模型参数。如果客户端未被选中，那么一轮通信结束后，该客户端的模型为一轮通信开始时从服务器获得的初始模型。

设当前全局模型为$w_t$，服务器选中了$m$个客户端（集合$V$），$m$个客户端本地更新完毕后，服务器端的聚合公式为：
$$w_{t+1}=\sum _{k\in V}\frac{n_k}{n}{w}_{t+1}^k+\sum _{k\notin V}\frac{n_k}{n}{w}_t$$
也就是说，每一次聚合时服务器端都将所有客户端的模型考虑在内。

2. 仅聚合被选中的客户端

服务器每次只是将当前新的参数传递给被选中的模型，并且只是聚合被选中客户端的模型参数。

设当前全局模型为$w_t$，服务器选中了$m$个客户端（集合$V$），然后将$w_t$只发送给这$m$个客户端。$m$个客户端训练完毕后，服务器端的聚合公式为：
$$w_{t+1}=\sum _{k\in V}\frac{n_k}{n}{w}_{t+1}^k$$

3. 选择

虽然原始论文中对所有$K$个客户端都进行了聚合，但在真正实现时，感觉用第二种会更好一点，因为如果客户端数量很庞大，每一次通信都会有不小的代价，用第二种会明显降低通信成本。