个性化联邦学习PFedMe详细解读（NeurIPS 2020）

本文介绍一篇 NeurIPS 2020 的论文『Personalized Federated Learning with Moreau Envelopes』，对个性化联邦学习 PFedMe 进行详细解读。

详细信息如下：

• 论文链接：https://arxiv.org/pdf/2006.08848.pdf
• 项目链接：https://github.com/CharlieDinh/pFedMe

01

背景

经典的机器学习方法基于样本数据（库）训练得到适用于不同任务和场景的机器学习模型。这些样本数据一般通过从不同用户、终端、系统中收集并集中存储而得到。在实际应用场景中，这种收集样本数据的方式面临很多问题。一方面，这种方法损害了数据的隐私性和安全性。在一些应用场景中，例如金融行业、政府行业等，受限于数据隐私和安全的要求，根本无法实现对数据的集中存储；另一方面，这种方法会增加通信开销。在物联网等一些大量依赖于移动终端的应用中，这种数据汇聚的通信开销成本是非常巨大的。

联邦学习允许多个用户（称为客户机）协作训练共享的全局模型，而无需分享本地设备中的数据。由中央服务器协调完成多轮联邦学习以得到最终的全局模型。其中在每一轮开始时，中央服务器将当前的全局模型发送给参与联邦学习的客户机。每个客户机根据其本地数据训练所接收到的全局模型，训练完毕后将更新后的模型返回中央服务器。中央服务器收集到所有客户机返回的更新后，对全局模型进行一次更新，进而结束本轮更新。通过上述多轮学习和通信的方法，联邦学习消除了在单个设备上聚合所有数据的需要，克服了机器学习任务中的隐私和通信挑战，允许机器学习模型学习分散在各个用户（客户机）上存储的数据。

联邦学习自提出以来获得了广泛的关注，并在一些场景中得以应用。联邦学习解决了数据汇聚的问题，使得一些跨机构、跨部门的机器学习模型、算法的设计和训练成为了可能。特别地，对于移动设备中的机器学习模型应用，联邦学习表现出了良好的性能和鲁棒性。此外，对于一些没有足够的私人数据来开发精确的本地模型的用户（客户机）来说，通过联邦学习能够大大改进机器学习模型和算法的性能。但是，由于联邦学习侧重于通过分布式学习所有参与客户机（设备）的本地数据来获得高质量的全局模型，因此它无法捕获每个设备的个人信息，从而导致推理或分类的性能下降。此外，传统的联邦学习需要所有参与设备就协作训练的共同模型达成一致，这在实际复杂的物联网应用中是不现实的。研究人员将联邦学习在实际应用中面临的问题总结如下：

1. 各个客户机（设备）在存储、计算和通信能力方面存在异构性；
2. 各个客户机（设备）本地数据的非独立同分布（Non-Idependently and Identically Distributed，Non-IID）所导致的数据异构性问题；
3. 各个客户机根据其应用场景所需要的模型异构性问题。

为了解决这些异构性挑战，一种有效的方法是在设备、数据和模型级别上进行个性化处理，以减轻异构性并为每个设备获得高质量的个性化模型，即个性化联邦学习（Personalized Federated Learning）。

02

引言

参考个性化模型在医疗保健、金融和人工智能服务等领域中应用的模式，本文提出了一种个性化联邦学习方案，该方案引入了基于客户端损失函数的 Moreau envelopes 优化。通过该方案，客户端不仅可以像经典联邦学习一样构建全局模型，而且可以利用全局模型来优化其个性化模型。从几何的角度分析，该方案中的全局模型可以看作是所有客户端一致同意的中心点，而个性化模型是客户端根据其异构数据分布来构建的遵循不同方向的点。这项工作中的主要贡献总结如下：

1. 首先，利用Moreau Envelopes作为正则化的损失函数，提出了一个新的面向个性化FL的双层优化问题pFedMe。pFedMe的双层结构有一个关键的优点：将个性化模型的优化过程与全局模型的学习过程解耦。因此，pFedMe以类似于FedAvg等标准FL算法的方式更新全局模型，同时以较低的复杂度并行优化个性化模型。
2. 利用Moreau Envelopes的保凸性和光滑性来促进pFedMe的收敛性分析，pFedMe同时表征了客户采样和客户漂移误差，通过调整超参数，pFedMe可以获得最先进的二次加速比。

03

方法

作者首先回顾了经典联邦学习算法（如下图所示）：服务器端通过聚合客户端模型参数来进行全局模型的更新迭代。

不同于经典的联邦学习，本文pFedMe对每个客户端使用 L2 范数的正则化损失函数，如下图所示：

其中，θ_i 表示客户端 i 的个性化模型，λ表示控制个性化模型的ω强度的正则化参数。虽然较大的λ（保证θ_i和ω尽可能相似）可以从丰富的数据聚合中受益于不可靠的数据，但是较小的λ可以帮助拥有足够多有用数据的客户端优先进行个性化设置。总之，本文方法的目的是允许客户端沿不同的方向更新本地模型，同时不会偏离每个客户端都贡献所得到的参考点ω。

更进一步而言，个性化联邦学习可以表述为如下一个双层问题：

通过在外部层利用来自多个客户端的数据聚合来确定ω，在内部层针对客户端 i 的数据分布优化θ_i，并使其与ω保持一定距离。F_i(ω)定义为 Moreau envelope。最优个性化模型是解决pFedMe内部层问题的唯一解决方案，在文献中也被称为邻近算子，定义如下：

为了进行比较，作者讨论了Per-FedAvg，这可以说是最接近pFedMe的公式：

Per-FedAvg是一个联邦元学习方法，基于经典元学习的与模型无关的元学习MAML框架，Per-FedAvg的目标是找到一个全局模型ω，可以用它作为初始化全局模型，进一步对损失函数执行梯度更新（步长为 α）来得到它的个性化模型θ_i(ω)。正如上篇文章所述，个人感觉就是全局模型ω在本地数据上再进行一次/若干次微调。

与Per-FedAvg相比，本文的问题具有类似于元模型的考虑，但是没有使用ω作为初始化，而是通过解决一个双层问题来并行地求得个性化和全局模型，这种方式有几个好处：首先，虽然Per-FedAvg针对其个性化模型进行了一步梯度更新的优化，但pFedMe对内部优化器是不可知的，这意味着公式(3)可以使用任何具有多步更新的迭代方法来求解。其次，可以将Per-FedAvg的个性化模型更新重写为：

接着，如下图1所示，便是PFedMe完整算法伪代码：首先，在内部层每个客户端 i 求解公式（3）以获得其个性化模型，其中 w^t_(i,r)表示客户端 i 在全局轮次 t 和局部轮次 r 的局部模型。与 FedAvg 类似，本地模型的目的是帮助构建全局模型，减少客户端和服务器之间的通信轮数。

图1：PFedMe伪代码

04

实验分析

本文实验考虑了一个使用真实MNIST和合成数据集的分类问题。MNIST 是一个手写数字数据集，包含 10 个标签和 70000 个实例。由于 MNIST 数据量的限制，作者将完整的 MNIST 数据集分发给N=20个客户端。为了根据本地数据大小和类别对异构设置进行建模，每个客户端都被分配了一个不同的本地数据大小范围为[1165,3834]，并且只有 10 个标签中的 2 个。对于合成数据（Federated Optimization in Heterogeneous Networks,），作者采用数据生成和分布过程，使用两个参数 α=0.5 和β=0.5 来控制每个客户端的本地模型和数据集的差异。具体来说，数据集使用 60 维实值数据为 10 类分类器提供服务。每个客户端的数据大小在[250,25810] 范围内。最后，将数据分发给 N=100 个客户端。

作者对pFedMe、FedAvg和Per-FedAvg进行了比较。MNIST 数据集中的实验结果如下图2所示，pFedMe的个性化模型在强凸设置下的准确率分别比其全局模型 Per-FedAvg 和 FedAvg 高 1.1%、1.3% 和 1.5%。非凸设置下的相应数据为 0.9%、0.9% 和 1.3%。

图2：PFedMe in MNIST 实验结果

对于合成数据集，利用相同参数和微调参数的比较结果见下图。尽管 pFedMe 的全局模型在测试准确率和训练损失方面表现不如其他模型，但pFedMe的个性化模型仍然显示出它的优势，因为它获得了最高的测试准确率和最小的训练损失。下图3显示，pFedMe 的个性化模型比其全局模型 Per-FedAvg 和 FedAvg 的准确率分别高出分别为 6.1%、3.8% 和 5.2%。

图3：PFedMe in 合成数据集实验结果

从实验结果来看，当客户端之间的数据是非独立同分布（Non-IID）时，pFedMe和Per-FedAvg都获得了比FedAvg更高的测试准确度，因为这两种方法允许全局模型针对特定客户端进行个性化处理。通过多次梯度更新近似优化个性化模型从而避免计算Hessian矩阵，pFedMe的个性化模型在收敛速度和计算复杂度方面比Per-FedAvg 更具优势。

05

总结

该论文提出了一种个性化联邦学习方法pFedMe。pFedMe利用了Moreau envelope函数，该函数有助于将个性化模型优化从全局模型学习中分解出来，从而使得 pFedMe可以类似于 FedAvg 更新全局模型，但又能根据T时刻每个客户端的本地数据分布并行优化个性化模型。理论结果表明，pFedMe 可以达到最快的收敛加速率。实验结果表明，在凸和非凸环境下，使用真实和合成数据集，pFedMe 的性能都优于经典 FedAvg 和基于元学习的个性化联邦学习算法 Per-FedAvg。

参考文献

T Dinh, Canh, Nguyen Tran, and Josh Nguyen. "Personalized federated learning with moreau envelopes." Advances in Neural Information Processing Systems 33 (2020): 21394-21405.