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} wRdmin{f(w):=N1i=1Nfi(w)}(1)
也就是找到一个全局的模型w,其中 f i : R d → R , i = 1 , . . . , N f_i:\mathbb{R}^d\rightarrow\mathbb{R},i=1,...,N fi:RdR,i=1,...,N,表示客户端i的损失函数:
f i ( w ) = E ξ i [ f ~ i ( w ; ξ i ) ] f_i(w)=\mathbb{E}_{\xi_i}[\tilde{f}_i(w;\xi_i)] fi(w)=Eξi[f~i(w;ξi)]
其中 ξ i \xi_i ξi表示从客户端i上随意选取一个样本。
在pFedMe中,对损失函数添加一个 l 2 l_2 l2正则项:
f i ( θ i ) + λ 2 ∥ θ i − w ∥ 2 (2) f_i(\theta_i)+\frac{\lambda}{2}\parallel\theta_i-w \parallel^2 \tag{2} fi(θi)+2λθiw2(2)
其中 θ i \theta_i θ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:wRdminF(w):=N1i=1NFi(w),其中Fi(w)=θiRdmin{fi(θi)+2λθiw2}
在pFedMe中,w为各个客户端参数的聚合(外循环), θ i \theta_i θ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)=θiRdargmin{fi(θi)+2λθiw2+2μθi2}(3)
(根据代码这里我添加了 μ 2 θ i 2 \frac{\mu}2\theta_i^2 2μθi2,作者这里没写出来)
为了比较,作者将Per-FedAvg算法也进行了定义:
min ⁡ w ∈ R d F ( w ) : = 1 N ∑ i = 1 N F i ( w ) , 其中 θ i ( w ) = w − α ∇ f i ( w ) (4) \min_{w\in\mathbb{R}^d}{F(w):=\frac{1}{N}\sum^N_{i=1}F_i(w)},\text{其中}\theta_i(w)=w-\alpha\nabla f_i(w)\tag4 wRdminF(w):=N1i=1NFi(w),其中θi(w)=wαfi(w)(4)
Per-Fedavg算法我们之前也提过,其借鉴MAML思想,让客户端获取初始化参数w,再通过w更新 θ i \theta_i θi(详细的可以看我的主页)。
作者表明,本方法和元学习类似,但不同于MAML找到好的初始化参数,pFedMe听过解决一个双层问题追求个性化和全局模型。其有几个优势。一是Ped-FedAvg针对个性化优化了梯度更新,但pFedMe对内部优化器是不可知的,也就是可以多次更新(3)。二是该方法直接优化了 f f f。三是Per-Fedavg需要计算Hessian矩阵而pFedMe只需要使用一阶方法进行梯度计算。

3.2 算法

与联邦学习类似,客户端训练后将参数传给服务器,作者这里用了一个超参数 β \beta β表示对聚合前服务器参数的利用程度,当 β = 1 \beta=1 β=1则为普通的联邦平均。
训练时,我们利用式(3)更新我们的参数 θ \theta θ(更新K轮),之后去找 δ \delta δ-approximate去更新w,再将w传入服务端即可。(这里文字叙述有点复杂,建议直接看图)
Personalized Federated Learning with Moreau Envelopes论文阅读+代码解析_第1张图片
不懂得小伙伴可以看我之后对应的代码讲解。

四. 代码详解

作者的代码点这里
为了讲解整个过程,我们从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 G^{t-1} Gt1表示t-1轮服务器的参数, C i C_i Ci表示客户端的参数,当有l个客户端参与聚合时,如下:
G t = ( 1 − β ) G t − 1 + β J ( C 1 , . . . , C l ) G^t=(1-\beta)G^{t-1}+\beta J(C_1,...,C_l) Gt=(1β)Gt1+βJ(C1,...,Cl)
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)=Gt1+i=1l{l1Ci}
之后将聚合好的参数传给客户端,开始客户端训练:

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 w 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传给服务器即结束,我这里给大家画了个图。
Personalized Federated Learning with Moreau Envelopes论文阅读+代码解析_第2张图片
到这里就结束啦,希望大家能看懂。

你可能感兴趣的:(每日一次AI论文阅读,个性化联邦学习,联邦学习,最新论文,pytorch)