OPT 2020 | FedAsync：异步联邦优化

前言

题目： Asynchronous Federated Optimization
研讨会： OPT 2020
论文地址：Asynchronous Federated Optimization

OPT 2020是NeurlPs 2020中的一个Workshops，其主题是机器学习优化。

在上一篇文章MLSys 2020 | FedProx：异质网络的联邦优化中我们了解到了一种用于缓解设备异质性和数据异质性的联邦优化算法FedProx。而本文作者为了提高联邦学习的灵活性和可扩展性，提出了一种新的异步联邦优化算法FedAsync，FedAsync同样可以缓解异质性，加速模型收敛。

FedAsync，简单来讲，客户端只要更新完毕就可以将本地模型上传到服务器，但此时客户端上传的模型可能是过时的，也就是说该客户端的计算能力和通信能力较差，没有按时完成模型的上传。比如说当前服务器已经进行到了第五轮聚合更新，而该设备才完成了第一次本地更新。​过时或者不过时的模型将按照时间顺序进行排队，服务器端通过陈旧性函数和混合超参数来控制本地模型在全局模型更新中所占的比重，对于过时程度较大的本地模型，其在全局模型更新中将占有较小的权重。

摘要

为了提高联邦学习的灵活性和可扩展性，提出了一种新的异步联邦优化算法FedAsync，该方法对强凸问题和一类受限非凸问题都具有近似线性的全局最优收敛性。

1. 引言

联邦学习有三个关键特性：

1. 不频繁的任务激活：对于弱边缘设备，学习任务仅在设备空闲、充电且连接到未测量网络时执行。
2. 不频繁的通信：边缘设备和远程服务器之间的连接可能经常不可用、速度慢或成本高（就通信成本或电池电量使用而言）。
3. 非独立同分布的训练数据。

联邦学习通常使用同步方式实现，也就是服务器同时给被选中的客户端发送最新的全局模型，然后在同一时间收集回所有的本地模型。这里存在的问题是显而易见的：​不同设备的计算和通信能力不同，所以可能会有很多弱设备无法按时完成本地更新任务。

异步训练广泛应用于传统的分布式随机梯度下降（SGD）中，用于离散和异构延迟。在本文中，作者利用异步训练的优势，并将其与联邦优化相结合，提出了一种新的联邦优化异步算法。其关键思想是：

1. 解决正则化局部问题以保证收敛。
2. 使用加权平均来更新全局模型，其中混合权重被自适应地设置为陈旧性的函数。

2. 问题定义

符号定义如下表所示：

3. 方法

系统框架：

第0步：调度器通过协调器触发训练。
第1、2步：服务器通过协调器将模型$x_{t-\tau }$发送给worker。
第3步：worker通过算法1计算本地更新。worker可以在Working和Idle两种状态间切换。
第4、5、6步：worker通过协调器将本地更新模型发送给服务器。其中第5步：协调器对收到的本地模型进行排队；第6步：按顺序将更新了的本地模型提交给更新程序Updater。
第7、8步：服务器更新全局模型，并且将更新后的模型传给协调器。

值得注意的是，第1步和第5步可以异步并行运行。

算法伪代码：

本文的重点就是详细解读FedAsync的过程。

3.1 Server进程

服务器进程的伪代码：

混合超参数$\alpha \in (0,1)$。首先初始化初始全局模型$x_0$以及混合超参数${\alpha }_t$（第$t$个epoch，每一个epoch表示一次通信迭代过程）。

然后并行异步运行Scheduler线程和Updater线程。

3.2 Scheduler线程

调度器线程的伪代码：

定期触发worker的训练任务，并发送带有时间戳的全局模型。

3.3 Updater线程

更新线程的伪代码：

服务器利用Updater线程进行全局模型的更新。

对每一次迭代，服务器从任意一个Worker处接收一个其更新好的本地模型$(x_{new},\tau )$，$\tau$为时间戳。

然后对全局模型进行更新：

首先需要更新混合超参数$\alpha$：
$${\alpha }_t\leftarrow \alpha \times s(t-\tau )$$
其中$s(.)$是staleness函数。

然后更新全局模型（可选）：
$$x_t\leftarrow (1-{\alpha }_t)x_{t-1}+{\alpha }_t{x}_{new}$$
也就是说，全局模型参数的更新与未更新前的全局模型参数和本地模型（更新后）有关，二者间的占比通过混合超参数$\alpha$来衡量，并且$\alpha$通过staleness函数来动态调节。

需要注意的是，在服务器端，Scheduler()和Updater()异步并行运行。调度器定期触发训练任务并且控制陈旧性（Updater()中的$t-\tau$）。

直观地说，更新全局模型时，更大的陈旧性意味着更大的错误！

所谓具有较大陈旧性的本地模型，也就是没有按时完成更新任务的模型（过时的模型），比如说当全局模型已经进行到第5轮更新时，突然收到了其中一个客户端第一轮更新后的本地模型，那么这个本地模型我们就可以说是具有较大的陈旧性。

对于具有较大陈旧性的本地模型，我们可以通过减小混合超参数$\alpha$来减轻其过时对全局模型更新的影响。

一般来讲，如果$t=\tau$，那么此时$s(t-\tau )=1$，也就是更新后混合超参数不变，此时服务器认为该模型就是最新的本地模型。

如果$t-\tau$增大，也就是模型比较陈旧，那么陈旧性函数$s(t-\tau )$应该减小，然后混合超参数也就减小，在进行模型更新时我们就会更少的考虑这个过时的本地模型$x_{new}$（权重为$\alpha$）。

当然有很多函数满足上述条件，比如最简单的：
$$s_a(t-\tau )=\frac{1}{t-\tau +1}$$

3.4 Worker进程

Worker进程的伪代码：

客户端通过Worker进程进行本地模型的更新。

具体对每一个客户端来讲：如果其被调度器触发，就会接收服务器发来的全局模型$(x_t,t)$，然后利用本地数据进行模型更新，需要优化的损失函数为：
$$g_{x_t}(x;z)=f(x;z)+\frac{\rho }{2}||x-x_t|{|}^2$$
这里$\rho >\mu$。

这里其实就是FedProx的思想，有关FedProx可见：MLSys 2020 | FedProx：异质网络的联邦优化。

我们再来捋一捋FedSync的具体流程：

1. 服务器初始化参数，然后并行异步运行Scheduler线程和Updater线程。
2. Scheduler定期触发worker的训练任务，并发送带有时间戳的全局模型。
3. worker通过协调器将本地更新模型发送给服务器，这些模型中可能会有过时的模型，协调器将对这些模型进行排队，然后按顺序将这些本地模型提交给更新程序Updater。
4. 更新程序Updater读取队列中的本地模型，对于过时的模型，Updater会通过陈旧性函数和混合超参数调整过时模型在更新全局模型时所占的权重，距离当下越久的本地模型将被考虑的更少以减小error。
5. 更新后的全局模型被发送给协调器，然后再次发送给客户端。

4. 收敛性分析

不详细讨论。

5. 实验

5.1 数据集

CIFAR-10和WikiText-2，mini-batch size分别为50和20，它们被划分到100个客户端中。

5.2 评估设置

baseline有两个：一个是FedAvg（每次随机选择$k=10$个客户端进行通信），另一个是单线程SGD。

FedAsync中定义了三种陈旧性函数$s(t-\tau )$：

其中$a,b>0$。

5.3 实验结果

讨论：

1. 一般来说，FedAsync的收敛速度介于单线程SGD和FedAvg之间。
2. 较大的$\alpha$和较小的陈旧性使FedAsync更接近单线程SGD。
3. 较小的$\alpha$和较大的陈旧性使FedAsync更接近FedAvg。
4. FedAsync对超参数不敏感，在陈旧性较大时我们可以调高$\alpha$来加速收敛。
5. 对于FedAsync，最佳的陈旧函数为Hinge。
6. 当过时性很小时，FedAsync的收敛速度比FedAvg快得多。当过时性较大时，FedAsync仍能获得与FedAvg类似的性能。

6. 总结

本文提出了一种基于非IID训练数据的异步联邦优化算法FedAsync，并且证明了一类受限非凸问题的收敛性。针对FedAsync，作者指出，一个比较好的研究方向是混合超参数$\alpha$的自适应策略。