FedProx论文阅读笔记

Posted 2021-07-22 LeoJarvis

《FEDERATED OPTIMIZATION IN HETEROGENEOUS NETWORKS》这篇文章提出了一个基于FedAvg的优化框架Fedprox，主要在于解决FedAvg未曾处理的联邦学习中的异构性问题，并提供了更强的鲁棒性。

FedAvg

在FedAvg中，每个设备在模型训练过程中具有相同的学习率和迭代次数。在每一轮更新中，被选中参与本轮训练的设备K在本地运行E轮SGD，每个设备将更新结果返回给服务器，服务器对得到的模型更新进行平均从而获得全局更新，并完成本轮训练。伪代码如Algorithm 1所示。

FedProx对FedAvg的改进之处

首先，由于FedAvg本身的设置，本地迭代次数 E 的增大虽能减少通信成本，但迭代次数过多不仅可能使一些受限于算力的设备无法完成训练，还容易使设备的本地模型偏离全局模型，影响全局收敛。导致这一问题的原因主要是联邦学习中存在的异构性问题。
FedProx主要解决了两种异构问题：
1.设备异构：设备间通信和计算能力存在差异，可能导致不同设备间更新不同步等问题。
2.数据异构：主要是设备间的non-IID问题。

FedProx

FedProx的伪代码如上所示，主要做了以下两点改进：
一、Tolerating partial work
相比于FedAvg对所有设备在每轮 global update都使用相同的local epoch， FedProx允许根据设备的可用系统资源在不同设备和不同的local epoch中使用不同的local epoch，对local subproblem进行粗略的minimize(而不是直接drop那些算力不足的设备)。（缓和设备异构）

二、Proximal term
1、在local subproblem上增加了一个proximal term（式二）。其中正则化项减去的是上一轮的全局模型，使得本地更新不会过分偏离global model。（缓和Non-IId）

FedProx的特点

1、可以应用于Non-IID场景。
2、可以选择任意的local server，而不单单局限于SGD。
3、允许某些设备使用inexact update。
4、不要求所有设备都参与每轮global update。
5、当μ=0，local solver选择SGD，不同设备在每轮global update都使用相同的local epoch，FedProx就变成了FedAvg，所以实际上，FedProx是FedAvg的一种泛化形式。

实验结果

1、不同程度的设备异构性对 FedAvg 和 FedProx 的影响

2、不同程度的数据异构性对 FedAvg 和 FedProx 的影响

总结

文章提出了一个基于FedAvg的优化框架Fedprox，用于解决联邦学习中的数据异构和设备异构问题。
通过加入 proximal term，缓解了数据异构性，提高了全局模型收敛的稳定性。
通过Tolerating partial work，缓解了设备异构性。