Personalized Federated Learning with Exact Stochastic Gradient Descent 论文阅读

论文地址点这里

一. 介绍

为了实现个性化，作者任务除了由一组主干神经网络层获得的全局或共享模型参数外，对于每个客户机，还存在一组具有客户机特定参数的客户机特定输出层。每个此类参数集都是客户独有的，其目的是允许对客户特定的数据分布或任务进行更灵活的建模。因此提出了一种新的个性化学习方法，称之为基于精确梯度优化的个性化联邦学习（PFLEGO）。其主要优点为实现了训练损失函数的精确的随机梯度下降，也就是所有数据集中在一个地方训练的等价。具体如下图所示：

二. 相关工作

通过微调实现个性化： 该方法允许客户使用局部适应或元学习技术微调共享的全局模型，与本方法不太一样，因为微调需要传递完整的参数集。
通过特征转移视线个性化学习： 作者的方法主要涉及通过使用多任务特征转移模型在FL中实现个性化的方法，类似于传统的非分布式多任务体系结构。
梯度返回的联邦： 最后，存在优化一组全局参数的非个性化FL算法，其中梯度返回到服务器。然而，只有在特殊情况下才能保证收敛，即客户机执行一次迭代，从而使其算法简化为SGD优化。

三. 构建框架

3.1 个性化联邦学习的设定

考虑一个监督的FL设置，其中有一个服务器和$I$个客户端，每一个客户端都有一个本地数据集$D_i=(X_i,Y_i)$，其中 X i = { x i , j } j = 1 N i X_i=\{x_{i,j}\}^{N_i}_{j=1} Xi​={xi,j​}j=1Ni​​作为输入数据的采样， Y i = { x i , j } j = 1 N i Y_i=\{x_{i,j}\}^{N_i}_{j=1} Yi​={xi,j​}j=1Ni​​作为相应的目标输出。联邦的目标是通过利用客户端本地存储的所有数据集，优化共享全局模型或个性化地更新特定于客户端的参数。作为共享的模型，假设其为一个深度神经网络，由多个具有总体参数的公共层组成$\theta$。公共层的输出数，也就是特征向量的大小为$M$。具体来说，特征层根据输入$x$，在其最终输出中构造一个特征向量$\varphi (x;\theta )\in {\mathbb{R}}^M$。每个客户端都有一个与此共享表示对应的相同网络体系结构的副本，此外每一个客户端i都有一个自身连接共享参数的个性层$W_i$。

3.1.1 训练目标分类

学习目标是通过调整全局参数$\theta$和个性化参数 { W i } i = 1 I \{W_i\}^I_{i=1} {Wi​}i=1I​。所有的参数集合由 ψ = { θ , W 1 , . . . , W I } \psi=\{\theta,W_1,...,W_I\} ψ={θ,W1​,...,WI​}。学习需要最小化以下函数：
$$\mathcal{L}(\psi )=\sum _{i=1}^I{\alpha }_i{l}_i(W_i,\theta )\text{\hspace{1em}(1)}$$
其中${\alpha }_i=\frac{N_i}{{\sum }_{j=1}{N}_j}$量化了不同客户端的数据比例，$l_i(W_i,\theta )$是客户端的损失：
$$l_i(W_i,\theta )=\frac{1}{N_i}\sum _{j=1}^{N_i}l(y_{i,j},x_{i,j};W_i,\theta )\text{\hspace{1em}(2)}$$
在多分类的情况下，每个客户端包括$K_i$个类，跨客户端的类可以相互排斥，也可以部分重叠。个性化权重集$W_i$变为$K_i\ast M$的矩阵，允许计算表春交叉熵损失的对数：
l ( y i , j , x i , j ; W i , θ ) = − log ⁡ { Pr ⁡ ( y i , j ∣ x i , j ; W i , θ ) } (3) l(y_{i,j},x_{i,j};W_i,\theta)=-\log\{\Pr(y_{i,j}|x_{i,j};W_i,\theta)\} \tag3 l(yi,j​,xi,j​;Wi​,θ)=−log{Pr(yi,j​∣xi,j​;Wi​,θ)}(3)
其中类概率由softmax函数建模。

3.2 提出算法

个性化的联邦是去最小化（1）中的损失对于参数集合$\psi$。为此，作者提出一种分布式的优化算法，该算法结合了精确随机梯度下降的步骤。随机性是由服务端每轮随机选取的客户端产生的，如果没有随机性，也就是所有客户端每轮都参与训练，那么就变成了精确的梯度下降。

3.2.1 客户端参与的流程

假设（1）中的全局损失优化是在不同的区域执行的，其中每一轮都涉及服务器与一些客户机的通信。更具体的说，我们假设在$t$轮开始的时候，一部分客户端 I t ⊂ { 1 , . . . , I } I_t\sub\{1,...,I\} It​⊂{1,...,I}被随机选择参与。这就导致两个选择:一是客户端的数量$r_t:=|I_t|$遵循二项分布$B(I,\rho )$，也就是每个客户端被概率$\rho$独立地选择。二是总是一定数量$0<r\le I$的客户端被一致地选择，也就是$|I_t|=r$对于任意时候$t$。对于这两种情况，每一个客户端以如下概率参与到每一轮：
$$\mathrm{Pr}(i\in {\mathcal{I}}_t)=\frac{r}{I}$$
其中对于情况一来说$r=I\rho$会是一个浮点数，而对于情况二来说$r$是一个整数。

3.2.2 客户端和服务端的更新

客户端： 在第$t$轮，被训啊的客户端${\mathcal{I}}_t$参与训练，服务器将全局参数$\theta$传给客户端。然后，每一个客户端$i\in {\mathcal{I}}_t$在本地训练特定的任务参数$W_i$,迭代$\tau$次。在前$\tau -1$次的训练中，服务器传来的参数$\theta$不变，只计算${\nabla }_{W_i}{l}_i(W_i,\theta )$。这$\tau -1$次地步骤表述为：
$$W_i\leftarrow W_i-\beta {\nabla }_{W_i}{l}_i(W_i,\theta )$$
其中$\beta$为学习率。相比之下，对于最后一次训练$\tau$，客户端将计算两个梯度$({\nabla }_{W_i}{l}_i,{\nabla }_{{\theta }_i}{l}_i)$，使用下式的规则更想你参数$W_i$:
$$W_i\leftarrow W_i-{\rho }_t\frac{I}{r}{\nabla }_{W_i}{l}_i(W_i,\theta )\text{\hspace{1em}(4)}$$
其中${\rho }_t$为第$t$轮的学习率，乘法标量$\frac{1}{\mathrm{Pr}(i\in {\mathcal{I}}_t)}=\frac{I}{r}$确保所有参数$\psi$进行无偏更新。
服务端： 服务端获得所有来自客户端的梯度进行更新，如下：
$$\theta \leftarrow \theta -{\rho }_t\frac{I}{r}\sum _{i\in {\mathcal{I}}_t}{\alpha }_i{\nabla }_{\theta }{l}_i(W_i,\theta )\text{\hspace{1em}(5)}$$
具体看下图：

3.3 精确的随机梯度下降优化

这一节是对上面的解释，也就是为什么按照这种更新方式是为精确的梯度更新呢？这里作者写的比较简单，我就不按照原文进行书写。首先在3.2.1中说到参与的客户端是有限制的，因此，如果单单只依靠参与的客户端来更新参数的话，就不能做到精确梯度(所有客户端更新)。既然做不到精确梯度更新，那么我们就想方法对这个梯度做无偏估计。现在我们只需要证明上述的梯度为精确梯度的无偏估计。下面是证明过程：
在t轮参与的客户端的数量为$r$，总的客户端为$I$，参数集合 ψ = { θ , W 1 , . . . , W I } \psi=\{\theta,W_1,...,W_I\} ψ={θ,W1​,...,WI​}，其参数梯度表示为：
∇ ψ S = { ∇ θ S L , ∇ W 1 S L , . . . , ∇ W I S L } \nabla_\psi^S=\{\nabla_\theta^S\mathcal{L},\nabla_{W_1}^S\mathcal{L},...,\nabla_{W_I}^S\mathcal{L}\} ∇ψS​={∇θS​L,∇W1​S​L,...,∇WI​S​L}
S表示为随机梯度，其中
$${\nabla }_{W_i}^S\mathcal{L}=P(i)\ast \frac{I}{r}{\nabla }_{W_i}l(W_i,\theta )$$
$${\nabla }_{\theta }^S\mathcal{L}=\frac{I}{r}\sum _{i=1}^{I}P(i)\ast {\alpha }_i{\nabla }_{\theta }{l}_i(W_i,\theta )$$
$$P(i)=\{\begin{array}{ll}1& \text{if }i\in {\mathcal{I}}_t\\ 0& \text{if }i\notin {\mathcal{I}}_t\end{array}$$
我们来计算均值：
$$\begin{array}{rl}\mathbb{E}({\nabla }_{W_i}^S\mathcal{L})& =\frac{P(i)\ast \frac{I}{r}{\nabla }_{W_i}l(W_i,\theta )}{I}\\ & =\frac{{\sum }_{i\in {\mathcal{I}}_t}{\nabla }_{W_i}l(W_i,\theta )}{r}\end{array}$$
可以发现其就是对r个客户端的梯度值的平均，同理对$\theta$也是，因此可以作为对精确梯度的求解。

（本次代码没有找到，如果作者之后补充了代码，我会添加进来的）