Personalized Federated Learning with Moreau Envelopes论文阅读+代码解析

（好久没更新文章啦，现在开学继续肝）
论文地址点这里

一. 介绍

尽管FL具有数据隐私和减少通信的优势，但它面临着影响其性能和收敛速度的主要挑战：统计多样性，这意味着客户之间的数据分布是不同的（即非i.i.d.）。因此，使用这些非i.i.d.数据训练的全局模型很难在每个客户的数据上得到很好的推广。因此，个性化联邦学习在改变传统追求全局一个较好模型做到了平衡——全局与局部的调整，以适应本地数据集。

二. 相关工作

个性化联邦学习有：混合模型，情景化，元学习和多任务学习。（论文介绍了很多相关论文，在这里不加以赘述）

三. pFedMe

3.1 问题构想

在传统的联邦学习中，在N个客户端时需要解决以下问题：
min ⁡ w ∈ R d { f ( w ) : = 1 N ∑ i = 1 N f i ( w ) } (1) \min_{w\in\mathbb{R}^d}\{f(w):=\frac{1}{N}\sum^N_{i=1}f_i(w)\} \tag{1} w∈Rdmin​{f(w):=N1​i=1∑N​fi​(w)}(1)
也就是找到一个全局的模型w，其中$f_i:{\mathbb{R}}^d\to \mathbb{R},i=1,...,N$，表示客户端i的损失函数：
$$f_i(w)={\mathbb{E}}_{{\xi }_i}[{\tilde{f}}_i(w;{\xi }_i)]$$
其中${\xi }_i$表示从客户端i上随意选取一个样本。
在pFedMe中，对损失函数添加一个$l_2$正则项：
$$f_i({\theta }_i)+\frac{\lambda }{2}\parallel {\theta }_i-w{\parallel }^2\text{\hspace{1em}(2)}$$
其中${\theta }_i$表示为客户端i的个性化模型，$\lambda$作为超参数控制个性化程度。$\lambda$越大个性化程序越低。基于此，个性化联邦学习被定义为下面的问题：
pFedMe : min ⁡ w ∈ R d F ( w ) : = 1 N ∑ i = 1 N F i ( w ) , 其中 F i ( w ) = min ⁡ θ i ∈ R d { f i ( θ i ) + λ 2 ∥ θ i − w ∥ 2 } \text{pFedMe}:\min_{w\in\mathbb{R}^d}{F(w):=\frac{1}{N}\sum^N_{i=1}F_i(w)},\text{其中}F_i(w)=\min_{\theta_i\in\mathbb{R}^d}\{f_i(\theta_i)+\frac{\lambda}{2}\parallel\theta_i-w \parallel^2\} pFedMe:w∈Rdmin​F(w):=N1​i=1∑N​Fi​(w),其中Fi​(w)=θi​∈Rdmin​{fi​(θi​)+2λ​∥θi​−w∥2}
在pFedMe中，w为各个客户端参数的聚合(外循环)，${\theta }_i$通过本地数据以及与w的距离进行更新（内循环）。最优个性化模型是pFedME的唯一解决方案，也被称为近端算子，定义如下：
θ i ^ ( w ) : = p r o x f i / λ ( w ) = arg min ⁡ θ i ∈ R d { f i ( θ i ) + λ 2 ∥ θ i − w ∥ 2 + μ 2 θ i 2 } (3) \hat{\theta_i}(w):=prox_{f_i/\lambda}(w)=\argmin_{\theta_i\in\mathbb{R}^d}\{f_i(\theta_i)+\frac{\lambda}{2}\parallel\theta_i-w \parallel^2+\frac{\mu}2\theta_i^2\}\tag3 θi​^​(w):=proxfi​/λ​(w)=θi​∈Rdargmin​{fi​(θi​)+2λ​∥θi​−w∥2+2μ​θi2​}(3)
（根据代码这里我添加了$\frac{\mu }{2}{\theta }_i^2$，作者这里没写出来）
为了比较，作者将Per-FedAvg算法也进行了定义：
$$\underset{w\in {\mathbb{R}}^d}{\min }F(w):=\frac{1}{N}\sum _{i=1}^N{F}_i(w),\text{其中}{\theta }_i(w)=w-\alpha \nabla f_i(w)\text{\hspace{1em}(4)}$$
Per-Fedavg算法我们之前也提过，其借鉴MAML思想，让客户端获取初始化参数w，再通过w更新${\theta }_i$（详细的可以看我的主页）。
作者表明，本方法和元学习类似，但不同于MAML找到好的初始化参数，pFedMe听过解决一个双层问题追求个性化和全局模型。其有几个优势。一是Ped-FedAvg针对个性化优化了梯度更新，但pFedMe对内部优化器是不可知的，也就是可以多次更新(3)。二是该方法直接优化了$f$。三是Per-Fedavg需要计算Hessian矩阵而pFedMe只需要使用一阶方法进行梯度计算。

3.2 算法

与联邦学习类似，客户端训练后将参数传给服务器，作者这里用了一个超参数$\beta$表示对聚合前服务器参数的利用程度，当$\beta =1$则为普通的联邦平均。
训练时，我们利用式(3)更新我们的参数$\theta$(更新K轮)，之后去找$\delta$-approximate去更新w，再将w传入服务端即可。（这里文字叙述有点复杂，建议直接看图）

不懂得小伙伴可以看我之后对应的代码讲解。

四. 代码详解

作者的代码点这里
为了讲解整个过程，我们从t-1轮客户端训练完传给服务器参数开始讲到第t轮。
t-1轮服务器获得一部分客户端的参数，利用$\beta$参数进行聚合：

def persionalized_aggregate_parameters(self):
    assert (self.users is not None and len(self.users) > 0)

    # store previous parameters
    previous_param = copy.deepcopy(list(self.model.parameters()))
    for param in self.model.parameters():
        param.data = torch.zeros_like(param.data)
    total_train = 0
    #if(self.num_users = self.to)
    for user in self.selected_users:
        total_train += user.train_samples

    for user in self.selected_users:
        self.add_parameters(user, user.train_samples / total_train)
        #self.add_parameters(user, 1 / len(self.selected_users))

    # aaggregate avergage model with previous model using parameter beta 
    for pre_param, param in zip(previous_param, self.model.parameters()):
        param.data = (1 - self.beta)*pre_param.data + self.beta*param.data

这里有一个函数self.add_parameters，其为聚合客户端的参数，具体如下：

def add_parameters(self, user, ratio):
    model = self.model.parameters()
    for server_param, user_param in zip(self.model.parameters(), user.get_parameters()):
        server_param.data = server_param.data + user_param.data.clone() * ratio

这里写个公式给大家$G^{t-1}$表示t-1轮服务器的参数，$C_i$表示客户端的参数，当有l个客户端参与聚合时，如下：
$$G^t=(1-\beta )G^{t-1}+\beta J(C_1,...,C_l)$$
J ( C 1 , . . . , C l ) = G t − 1 + ∑ i = 1 l { 1 l C i } J(C_1,...,C_l)=G^{t-1}+\sum_{i=1}^l\{\frac{1}{l}C_i\} J(C1​,...,Cl​)=Gt−1+i=1∑l​{l1​Ci​}
之后将聚合好的参数传给客户端，开始客户端训练：

def train(self, epochs):
    LOSS = 0
    self.model.train()
    for epoch in range(1, self.local_epochs + 1):  # local update
        
        self.model.train()
        X, y = self.get_next_train_batch()

        # K = 30 # K is number of personalized steps
        for i in range(self.K):
            self.optimizer.zero_grad()
            output = self.model(X)
            loss = self.loss(output, y)
            loss.backward()
            self.persionalized_model_bar, _ = self.optimizer.step(self.local_model)

        # update local weight after finding aproximate theta
        for new_param, localweight in zip(self.persionalized_model_bar, self.local_model):
            localweight.data = localweight.data - self.lamda* self.learning_rate * (localweight.data - new_param.data)

    #update local model as local_weight_upated
    #self.clone_model_paramenter(self.local_weight_updated, self.local_model)
    self.update_parameters(self.local_model)

    return LOSS

其中self.model.weight对应的是$\theta$，self.local_model对应的是$w$
先更新$\theta$，更新的方法为式(3)，对应的代码如下：

def step(self, local_weight_updated, closure=None):
    loss = None
    if closure is not None:
        loss = closure
    weight_update = local_weight_updated.copy()
    for group in self.param_groups:
        for p, localweight in zip( group['params'], weight_update):
            p.data = p.data - group['lr'] * (p.grad.data + group['lamda'] * (p.data - localweight.data) + group['mu']*p.data)
    return  group['params'], loss

之后像上图的算法更新w即可。最后再将w传给服务器即结束，我这里给大家画了个图。

到这里就结束啦，希望大家能看懂。