本文介绍的是一篇收录于 AAAI 2023 的论文，论文由位于上海交通大学的上海市可扩展计算与系统重点实验室、贝尔法斯特女王大学的华扬老师和路易斯安那州立大学的王灏老师共同完成。
- 论文链接：https://arxiv.org/abs/2212.01197
- 代码链接（包含ALA模块的使用说明）：https://github.com/TsingZ0/FedALA
该论文提出了一种用于联邦学习的自适应本地聚合方法，通过从全局模型中自动捕获客户机所需信息的方式来应对联邦学习中的统计异质性问题。作者对比了 11 个 SOTA 模型，并取得了超越最优方法 3.27% 的优异表现。作者将其中的自适应本地聚合模块应用到其他联邦学习方法上取得了最多 24.19% 的提升。
1 介绍
联邦学习（FL）通过将用户隐私数据保留在本地而不进行传播的方式，帮助人们在保护隐私的情况下，充分发掘用户数据中蕴含的价值。但由于客户机之间数据不可见，数据的统计异质性（数据非独立同分布（non-IID）和数据量不平衡现象）便成了 FL 的巨大挑战之一。数据的统计异质性使得传统联邦学习方法（如 FedAvg 等）很难通过 FL 过程训练得到适用于每个客户机的单一全局模型。
近年来，个性化联邦学习（pFL）方法因其应对数据统计异质性的能力受到了越来越多的关注。与寻求高质量全局模型的传统 FL 不同，pFL 方法的目标是借助联邦学习的协同计算能力为每个客户机训练适用于自身的个性化模型。现有的在服务器上聚合模型的 pFL 研究可以分为以下三类：
(1) 学习单个全局模型并对其进行微调的方法，包括 Per-FedAvg 和 FedRep；
(2) 学习额外个性化模型的方法，包括 pFedMe 和 Ditto；
(3) 通过个性化聚合（或本地聚合）学习本地模型的方法，包括 FedAMP、FedPHP、FedFomo、APPLE 和 PartialFed。
类别（1）和（2）中的 pFL 方法将全局模型中的所有信息用于本地初始化（指在每次迭代的局部训练之前初始化局部模型）。然而，在全局模型中，只有提高本地模型质量的信息（符合本地训练目标的客户机所需信息）才对客户机有益。全局模型的泛化能力较差是因为其中同时存在对于单一客户机来说需要和不需要的信息。因此，研究者们提出类别（3）中的 pFL 方法，通过个性化聚合捕获全局模型中每个客户机所需的信息。但是，类别（3）中的 pFL 方法依旧存在（a）没有考虑客户机本地训练目标（如 FedAMP 和 FedPHP）、（b）计算代价和通讯代价较高（如 FedFomo 和 APPLE）、（c）隐私泄露（如 FedFomo 和 APPLE）和（d）个性化聚合与本地训练目标不匹配（如 PartialFed）等问题。此外，由于这些方法对 FL 过程做了大量修改，它们使用的个性化聚合方法并不能被直接用于大多数现有 FL 方法。
为了从全局模型中精确地捕获客户机所需信息，且相比于 FedAvg 不增加每一轮迭代中的通讯代价，作者提出了一种用于联邦学习的自适应本地聚合方法（FedALA）。如图 1 所示，FedALA 在每次本地训练之前，通过自适应本地聚合（ALA）模块将全局模型与本地模型进行聚合的方式，捕获全局模型中的所需信息。由于 FedALA 相比于 FedAvg 仅使用 ALA 修改了每一轮迭代中的本地模型初始化过程，而没有改动其他 FL 过程，因此 ALA 可被直接应用于大多数现有的其他 FL 方法，以提升它们的个性化表现。
图 1：在第次迭代中客户机上的本地学习过程
2 方法
2.1 自适应本地聚合（ALA）
图 2：自适应本地聚合（ALA）过程
自适应本地聚合（ALA）过程如图 2 所示。相比于传统联邦学习中直接将下载的全局模型覆盖本地模型得到本地初始化模型的方式（即），FedALA 通过为每个参数学习本地聚合权重，进行自适应本地聚合。
其中，作者把
这一项称之为“更新”。此外，作者通过逐元素权重剪枝方法
实现正则化并将
中的值限制在 [0,1] 中。
因为深度神经网络（DNN）的较低层网络相比于较高层倾向于学习相对更通用的信息，而通用信息是各个本地模型所需信息，所以全局模型中较低层网络中的大部分信息与本地模型中较低层网络所需信息一致。为了降低学习本地聚合权重所需的计算代价，作者引入一个超参数 p 来控制 ALA 的作用范围，使得全局模型中较低层网络参数直接覆盖本地模型中的较低层网络，而只在较高层启用 ALA。
其中，
表示
中的神经网络层数（或神经网络块数），
与
中的低层网络形状一致，
则与
中剩下的 p 层高层网络形状一致。
作者将
中的值全部初始化为 1，且在每一轮本地初始化过程中基于旧的
更新
。为了进一步降低计算代价，作者采用随机采样 s
其中，
是更新
的学习率。作者在学习
的过程中，将除
之外的其他可训练参数冻结。
图 3：在 MNIST 和 Cifar10 数据集上 8 号客户机的学习曲线
通过选择较小的 p 值，在几乎不影响 FedALA 表现的情况下，大幅度地降低 ALA 中训练所需的参数。此外，如图 3，作者观察到：一旦在第一次训练
将其训练到收敛，即使在后续迭代中训练
，其对本地模型质量也没有很大影响。也就是说，每个客户机可以复用旧的
实现对其所需信息的捕获。作者采取在后续迭代中微调
的方式，降低计算代价。
2.2 ALA 分析
在不影响分析的情况下，为了简便起见，作者忽略
并设
。根据上述公式可以得到
，其中
代表
。作者可以把在 ALA 中更新
看成更新
。
梯度项
在每一轮中都做了逐元素地缩放。不同于本地模型训练（或 fine-tuning）方法，上述对
的更新过程能够感知到全局模型中的通用信息。在不同的迭代轮数之间，动态变化的
为 ALA 模块引入了动态信息，使其 FedALA 容易适应复杂的环境。
3 实验
作者在实际（practical）数据异质环境下的 Tiny-ImageNet 数据集上用 ResNet-18 进行了对超参数 s 和 p 的对 FedALA 影响的研究，如表 1 所示。对于 s 来说，采用越多的随机采样的本地训练数据用于 ALA 模块学习可以使个性化模型表现更好，但也同时增加了计算代价。在使用 ALA 的过程中，可以根据每个客户机的计算能力调整 s 的大小。从表中可以得知，即使使用极小的 s（如 s=5），FedALA 依旧具有杰出的表现。对于 p 来说，不同的 p 值对个性化模型的表现几乎没有影响，在计算代价方面却有着巨大的差别。这一现象也从一个侧面展示了 FedRep 等方法，将模型分割后保留靠近输出的神经网络层在客户机不上传的做法的有效性。使用 ALA 时，我们可以采用较小且合适的 p 值，在保证个性化模型表现能力的情况下，进一步降低计算代价。
表 1：关于超参数和对 FedALA 影响的研究
作者在病态（pathological）数据异质环境和实际（practical）数据异质环境下，将 FedALA 与 11 个 SOTA 方法进行了对比和详细分析。如表 2 所示，数据显示 FedALA 在这些情况下的表现都超越了这 11 个 SOTA 方法，其中 “TINY” 表示在 Tiny-ImageNet 上使用 4-layer CNN。例如，FedALA 在 TINY 情况下比最优基线（baseline）高了 3.27%。
表 2：病态和实际数据异质环境下的实验结果
此外，作者也在不同异质性环境和客户机总量情况下评估了 FedALA 的表现。如表 3 所示，FedALA 在这些情况下依旧保持着优异的表现。
表 3：其他实验结果
根据表 3 的实验结果，将 ALA 模块应用于其他方法可以获得最多 24.19% 的提升。
最后，作者还在 MNIST 上可视化了 ALA 模块的加入对原本 FL 过程中模型训练的影响，如图 4 所示。不激活 ALA 时，模型训练轨迹与使用 FedAvg 一致。一旦 ALA 被激活，模型便可以通过全局模型中捕获的其训练所需信息径直朝着最优目标优化。
图 4：4 号客户机上模型训练轨迹的可视化图
