Personalized Federated Learning: A Meta-Learning Approach 个性化联邦学习PerFedAvg开山鼻祖之作

PerFedAvg，个性化联邦学习

元损失函数求导作为更新目标。在联邦学习中，我们的目标是跨多个计算单元(用户)训练模型，而用户只能与一个公共的中央服务器通信，而不交换他们的数据样本。目标是找到一个初始共享模型，当前或新用户可以通过对他们自己的数据执行一个或几个梯度下降步骤来轻松适应他们的本地数据集。这种方法保留了联邦学习体系结构的所有优点，并且根据结构，为每个用户提供了更加个性化的模型。

intro

开门见山介绍了FL的目标函数定义，以及为什么要用FML，最后提出了PerFedAvg

作者认为每个用户都可以在不与其他用户交换信息的情况下解决定义的局部问题;然而，所得到的模型可能不能很好地推广到新的样本，因为它已经在少数样本上进行了训练。如果用户相互合作，利用所有用户的可用数据，那么他们的局部模型可以获得更强的泛化保证。实现这一目标的传统方法是最小化局部函数的集合。在用户的底层数据分布不相同的异构设置中，通过最小化平均损失获得的全局模型一旦应用于每个用户的本地数据集可能会表现得非常差。在本文中，通过考虑包含个性化的联邦学习问题的修改公式来克服这个问题。在引入的模型不可知元学习(MAML)问题公式的基础上，新公式的目标是找到所有用户之间共享的初始点，该初始点在每个用户根据自己的损失函数更新后表现良好，可能是通过执行基于梯度的方法的几个步骤。这样，虽然初始模型是在所有用户中以分布式方式派生的，但每个用户基于自己的数据实现的最终模型与其他模型不同。提出了FedAvg算法的个性化变体，称为Per-FedAvg，旨在解决所提出的个性化FL问题(第3节)。

相关工作中，介绍了FL中隐私保护，降低通信开销，也有学习用户数据异构，但没有生成独立的模型。然后说了一下MAML和FL的联系以及自己这篇文章研究的新颖和不同。提到了ARUBA算法，MOCHA算法（解决数据和系统挑战）等等

正式介绍

简单回顾MAML。

给定一组从底层分布中提取的任务，在MAML中，与传统的监督学习设置相反，目标不是找到一个在预期的所有任务上都表现良好的模型。相反，在MAML中，假设在新任务到达后，有限的计算预算来更新模型，在这个新设置中，寻找一个初始化，它在相对于这个新任务更新后表现良好，可能是通过一个或几个梯度下降步骤。

这个公式的优势在于，保持FL的优势还捕获了用户之间的差异，无论是现有用户还是新用户都可以将这个新问题的解决方案作为初始点，并根据自己的数据稍微更新它。

对FedAvg的问题：简单的几步梯度下降不一定可以到达优化点，但是MAML设置就是如此。

具体算法如下。

注意更新过程，非常平常，但是里面有个二次梯度，这就是黑塞矩阵（Hessian Matrix）。对于普通的深度学习模型是不可能算出来的，因为占用O(n^2)空间。作者如何缓解这个问题呢？首先用一批数据做一次梯度下降来进行无偏估计；其次黑塞矩阵也类似。所以作者在这里没有解决超大占用的问题，他们用的是比较简单的小模型，根本不会遇到这个问题。具体解决要看后面其他人的工作。

为什么 Hessian 矩阵是 O(n^2)

Hessian 矩阵是损失函数相对于模型参数的二阶偏导数组成的矩阵。如果参数 w 是一个长度为 \(n\) 的向量，那么 Hessian 矩阵将是一个 n x n 的矩阵，因为你需要对每一对参数 w_i和 w_j计算二阶偏导 $\frac{\partial^2 f}{\partial w_i \partial w_j}$。

这个矩阵的大小随着参数数量的增加而呈平方增长，因此计算和存储成本都非常高。这是为什么在实际应用中，尤其是在参数数量非常多的深度学习模型中，直接计算和使用 Hessian 矩阵通常不可行。

Per-FedAvg和普通FedAvg一样，除了将损失函数改为元损失函数的平均值

于是更新就可以对这个函数进行求导，得到我们所需要的更新公式。这里有一堆弧线是因为它们都基于无偏估计。（样本抽取是随机的）

但是这里的计算结果是一个有偏估计，因为公式里有一个嵌套的随机梯度下降过程。

为什么有偏？因为公式最右侧的梯度下降是一个嵌套了梯度下降的梯度，而连续的梯度下降后其实是不满足假设的，即模型不变化。

理论分析

这是本文章的重点。大致概括一下。这里都是为了证明和推导一些性质的，不需要完全看懂，很多都是基础且常用的假设。

定义4.1 定义了一个不动点，即此处梯度平方小于某个值

• 假设1. 损失函数有下界

• 假设2. 所有的损失函数都是二次连续可分，L平滑并且其梯度有非负上界

• 假设3. 所有黑塞矩阵都是利普西斯连续的

• 假设4. 计算出来的黑塞矩阵和梯度都是有上界的

• 假设5. 损失函数f和本地损失函数fi的梯度和黑塞矩阵之间的差距都是有上界的

其实如果假设2成立，那么假设5也在特定情况下成立。作者分开是为了强调函数的相似程度在收敛性分析中的作用。

任务相似度和分布之间距离的关系，作者在两个距离上进行分析。

• TV距离，将两个分布之间的距离投影到可数集合Z上进行判断。如果更进一步强化假设2在任意z属于Z的情况下成立，梯度和黑塞就都有上界B和L，从而假设5中右侧的两个数就可以确认了。这就表明了假设5中的参数可以表示用户之间不同的分布。

• 1-Wasserstein距离，也叫推土机距离。简单解释就是将一个概率分布变换为另一个概率分布所需的"最小工作量"。其中的1指的是在距离度量中使用的是一次范数（或相应的度量的一次幂）。同样可以确认假设5

收敛性分析：为了证明算法可以到达一个不动点。

• 引理4.2 如果假设2和3成立，那么元损失函数是平滑的，平均损失函数也平滑

• 引理4.3 如果假设2、3、4成立，那么估计的元损失函数的梯度和真实梯度之间的差距是有上界的，同样其之间的差距平方也有上界。特别地如果真的去计算真是的梯度和黑塞，这个上界就是0即完全相同。作者特意提到这个上界是紧到常数级别的。

• 引理4.4 本地元损失函数和全局元损失函数的差距有上界，如果假设2、3、5满足。同样强调了紧致

定理4.5 考虑目标函数F，如果假设1-4成立且各参数按照4.2到4.4定义，则本文定义的算法可以到达一个不动点。

只要batch size足够大，就能让上界足够小。作者之后开始寻找合适的更新次数，通信次数和外层学习率。

• 推论4.6 确定了本地更新等参数的数量级

• 评论4.7 从4.5和4.6中可以找到上界，但是我们只能获取每个通信初期的模型，所以可以让服务器随机挑选一个本地更新次数让客户端将最终模型和这个本地更新次数的模型一起上传从而获得一个相同的模型。

• 评论4.8 学习率逐渐减小也可以达到相同的收敛界限（用于实际训练）

数值试验

数据集：

• MNIST和CIFAR10，前者a=196，后者a=68

分配：分给n个用户。一半客户端拥有前五个类别的a张图片，剩下的只拥有前五个类别的一种中的a/2张图片，还有剩下的五个类别的一种的2a张。测试集合也是一样的分布

baseline：FedAvg，之后在test集合上做一步梯度下降后测试性能。之后使用了两种一阶估计版本的PerFedAvg来进行评估。这两种方法，第一种直接忽略了Hessian矩阵，第二种采用了在参数上添加前后扰动然后近似估计hessian的方法。

在小学习率情况下两种版本FirstOrderMAML和HessianFreeMAML表现都很好，但是在比较大时HFMAML更好（中心化设置）。所有模型都在测试前梯度下降一次。

使用规模比较小的神经网络（这解释了为什么能计算hessian矩阵），ELU激活函数，50个用户，1000轮，每次20%参与，batchsize=40，beta=0.001，汇报95%置信区间。

结果：

• MNIST上两种PerFedAvg都取得了微小的提升（相比FedAvg）

• CIFAR10上HF效果更好（很合理，因为至少采用了某种方式计算Hessian）

• alpha 0.001，tau=10，FO和HF表现相似且好于FedAvg。tau越小三种方法效果都很差

• alpha从0.001到0.01，tau=4，HF表现提升，即模型在测试时表现更好，但是alpha变大后FO效果显著变差。

• 把拥有前五个类别图片中一类的a/2客户端的这一类图片移除来提升异构性，可以看见HF效果显著更好，但是FO还是出现性能下降。因此推荐采用HF

总结

这是一篇非常好的文章。虽然实验部分不多，但是理论依据充足，数学功底扎实，非常有想法且充分，比绝大多数国内跟风的FL论文好上太多，显得他们都是垃圾。

仍然有些值得思考的问题：在test集合上梯度下降，是否不公平？能否进一步减少开销？这就得看后面的Reptile算法了。