联邦学习中的基础算法介绍

提示

联邦学习（Federated Learning）是一种分布式机器学习方法，通过将模型训练分散到多个客户端上，从而保护用户隐私。在联邦学习中，各客户端在本地训练模型，然后将更新结果传回服务器，由服务器聚合这些更新。本文将介绍联邦学习中一些重要的算法：FedAvg、FedProx、FedBN、SCAFFOLD 和 MOON。

1. FedAvg

Federated Averaging (FedAvg) 是联邦学习中的经典算法之一，也是许多后续算法的基础。FedAvg 的核心思想是：各客户端在本地数据上独立训练模型一段时间（多轮次迭代），然后将模型更新发送到服务器端，由服务器进行平均化聚合，从而生成全局模型。FedAvg 具有简单、高效的优点，但在客户端数据存在非独立同分布（non-IID）问题时，FedAvg 的效果会有所下降。

• 原论文链接：Communication-Efficient Learning of Deep Networks from Decentralized Data (opens new window)

给定参数：$K$个客户端，客户端索引为$k$；本地小批量大小为$B$，本地训练轮数为$E$，学习率为$\eta$。

---

服务器执行：

1. 初始化$w_0$
2. 对于每一轮$t=1,2,\dots$做：

   1. $m\leftarrow max(C\cdot K,1)$

   2. $S_t\leftarrow$随机选择的$m$个客户端

   3. 对于每个客户端$k\in S_t$并行地执行：

      $w_{t+1}^k\leftarrow \text{ClientUpdate}(k,w_t)$

   4. 计算聚合更新：

      $$w_{t+1}\leftarrow \sum _{k\in S_t}\frac{n_k}{n_{S_t}}{w}_{t+1}^k$$

      其中$n_k$是客户端$k$的样本数，$n_{S_t}$是所有参与客户端的总样本数。

客户端更新过程 ClientUpdate(k, w)： // 在客户端$k$上运行

1. 将客户端$k$的数据$P_k$划分为批次，每批大小为$B$

2. 对于本地轮次$i=1$到$E$做：

   对于每个批次$b\in B$进行梯度下降更新：

   $$w\leftarrow w-\eta \mathrm{\nabla }\ell (w;b)$$

3. 返回更新后的$w$给服务器

2. FedProx

为了解决 FedAvg 在非IID数据分布时的表现问题，Federated Proximal (FedProx) 在 FedAvg 的基础上引入了一个正则项（Proximal Term） $\frac{\mu }{2}\parallel w-w_t{\parallel }^2$ 旨在控制当前模型参数 $w$ 与上一轮全局模型参数 $w_t$ 之间的距离。通过限制模型更新的幅度，FedProx 可以防止客户端在训练过程中发生过度偏移，尤其是当数据分布差异较大时。这种控制能够确保客户端更新不至于偏离全局模型太远，从而保持全局模型的稳定性。

• 原论文链接：Federated Optimization in Heterogeneous Networks (opens new window)

输入： 客户端总数 $K$，总轮数 $T$，正则参数 $\mu$，学习率 $\gamma$，初始模型参数 $w_0$，客户端数 $N$，每个客户端被选中的概率 $p_k$（$k=1,\cdots ,N$）

---

For 对于每一轮 $t=0,\cdots ,T-1$：

1. 服务器随机选择一个大小为 $S_t$的客户端子集，其中每个客户端 $k$被选中的概率为 $p_k$

2. 服务器将当前模型参数 $w_t$发送给所有被选中的客户端

3. 每个被选中的客户端 $k\in S_t$计算一个 ${\gamma }_t$-近似解 $w_{t+1}^k$，使得：

   $$w_{t+1}^k\approx \arg \underset{w}{min}{h}_k(w;w_t)=F_k(w)+\frac{\mu }{2}\parallel w-w_t{\parallel }^2$$

4. 每个客户端 $k\in S_t$将更新后的 $w_{t+1}^k$发送回服务器

5. 服务器聚合更新的 $w$值，得到新的全局模型：

   $$w_{t+1}=\frac{1}{K}\sum _{k\in S_t}{w}_{t+1}^k$$

End 结束循环

3. FedBN

Federated Batch Normalization (FedBN) 主要针对包含不同特征分布的多域数据，例如在医疗影像等场景中，各客户端的数据分布存在显著差异。FedBN 保持每个客户端的批归一化（Batch Normalization）层的独立性，仅聚合其他层的参数。这种方法能够很好地适应每个客户端的本地数据分布，同时保证全局模型的通用性。

• 原论文链接：FedBN: Federated Learning on Non-IID Features via Local Batch Normalization (opens new window)

4. SCAFFOLD

SCAFFOLD 算法在 FedAvg 的基础上，通过引入控制变量来减少本地梯度更新中的差异，克服了在非IID数据下的偏差问题。每个客户端和服务器都维护一个控制变量，用于校正梯度更新中的偏差，从而加速收敛。

• 原论文链接：SCAFFOLD: Stochastic Controlled Averaging for Federated Learning (opens new window)

5. MOON

Model Contrastive Federated Learning (MOON) 是一种对比学习方法，旨在优化联邦学习中的局部模型和全局模型的相似度。MOON 通过对比学习的损失函数，使本地模型更新更接近全局模型的方向，避免了各客户端的更新偏差。

• 原论文链接：MOON: Model Contrastive Federated Learning (opens new window)