SCAFFOLD: Stochastic Controlled Averaging for Federated Learning

背景

传统的联邦学习在数据异构(non-iid)的场景中很容易产生“客户漂移”(client-drift)的现象，这会导致系统的收敛不稳定或者缓慢

贡献

• 提出了考虑到client sampling和数据异构的一个更接近的收敛边界

• 证明即便没有client sampling，使用全批次梯度(full batch gradients)，传统的FedAvg依旧会因为client-drift而比SGD收敛速度更慢

• 提出Stochastic Controlled Averaging algorithm(SCAFFOLD)，目的便是为了解决client-drift的问题，并证明了SCAFFOLD算法在数据异构的情况下收敛速度至少和SGD一样快

• SCAFFOLD算法还可以利用client之间的相似度来减少通信开销

• 证明了SCAFFOLD算法不会被client sampling所影响，这使得SCAFFOLD算法更适合联邦学习

论文思想

传统联邦学习的方法FedAvg算法在异构数据集上表现不好的原因是有一些client会带偏整个系统的收敛结果，如下图所示：

在上图中，黑色点是全局模型，也就是每个训练轮次各个局部模型的“训练起点”，假设在某一轮训练中，服务器选择了client1和client2两个客户端来训练，然后client1是偏离整个系统的客户端，那么在客户端上训练三个轮次中，我们可以看到client1上的局部模型已经偏离了训练的方向($x^{\ast }$所在的方向)，然后聚合得到的server model也会稍微偏离$x^{\ast }$，使得系统向着偏离学习模型的方向上收敛。最终的结果不是造成整个系统的性能下降就是导致整个系统收敛缓慢

为了解决这个问题，论文使用一个“控制变量”(control variate)$c$来“纠正”系统训练的方向，在client对模型进行更新的时候，也会对该变量进行更新

算法

与传统的联邦学习类似，SCAFFOLD算法也分为三个主要的部分：

• 局部更新模型(local updates to the client model)
• 局部更新控制变量(local updates to the client control variate)
• 对局部的更新进行聚合

先给出算法的流程，后面再做出解释：

算法具体流程：

• 服务器输入：初始化的$x$和$c$，全局更新步长${\eta }_g$
• 客户端输入：控制变量$c_i$和局部更新步长${\eta }_l$
• 在每一个全局轮次
  • 对clients进行sample
  • 对于每一个client $i$
    • 发送全局模型$x$和控制变量$c$
    • 用全局模型$x$来初始化$y_i$
    • 在每一个局部轮次：
      • 计算mini-batch梯度$g_i(y_i)$
      • 用$y_i\leftarrow y_i-{\eta }_l(g_i(y_i)-c_i+c)$进行更新
    • 局部训练完毕之后，得到新的控制变量$c_i^{+}\leftarrow (i)g_i(x)or(ii)c_i-c+\frac{1}{K{\eta }_l}(x-y_i)$
    • 上传梯度$y_i-x$和$c_i^{+}-c_i$
    • 更新控制变量$c_i\leftarrow c_i^{+}$
  • 对上传的梯度进行聚合$(\Delta x,\Delta c)\leftarrow \frac{1}{|S|}{\sum }_{i\in S}(\Delta y_i,\Delta c_i)$
  • 全局更新：$x\leftarrow x+{\eta }_g\Delta x$和$c\leftarrow c+\frac{|S|}{N}\Delta c$

局部更新方式

SCAFFOLD算法在局部的更新方式是：
$$y_i\leftarrow y_i-{\eta }_l(g_i(y_i)-c_i+c)$$
其中控制变量$c$的作用很明显，便是用全局模型的知识去约束局部模型的训练，以防止其偏离系统的正确训练方向，如下图所示：

并且该控制变量也会更新，以下面的方式：

论文给给出了上面两种更新方式的选择，其中第一种是用局部的梯度来更新全局模型中的控制变量$c$，第二种复用了全局模型的知识，直观上理解是根据全局模型与局部模型的差异来更新$c$。论文中给出的上面两种选择的区别是第一种方法要更稳定，第二种方法更加取决于应用场景，但是第二种方法更容易计算并且在通常情况下也已经足够优秀

全局更新方式

对于模型的更新与传统联邦并无太大区别：
$$x\leftarrow x+{\eta }_g\Delta x$$
控制变量的更新：
$$c\leftarrow c+\frac{1}{N}\sum _{i\in S}(c_i^{+}-c_i)$$
控制变量的更新方法也是和模型的更新方法差不多，本质上都是将局部模型的知识更新传递到全局模型

实验

实验在EMNIST数据集上进行，结果证明了SCAFFOLD算法与FedAvg算法和FedProx算法相比是表现最好的，如下图所示：

上面的3幅图表示的是SGD和FedAvg的比较实验，可以看到当梯度差异(G)很小是，FebAvg在训练刚刚开始的时候要比SGD好，但是在当G比较大的时候，由于客户容易发生“客户漂移”现象，容易带偏系统的训练方向，因此收敛效果和速度都会变差。下面的3幅图表示的是论文提出的算法Scaffold与SGD的比较，可以看到Scaffold算法无论是收敛速度和效果都比SGD要好

总结

论文的一个基本思想本质上便是用全局模型的知识去限制局部模型的训练方向，以防止那些与全局模型相差较大的局部模型带偏整个系统的训练方向