一种基于分层抽样的个性化联邦学习方法

A personalized federated learning method based on stratified sampling in edge network

摘要

联邦学习是一种保护数据隐私的分布式学习方法，通过聚合多个客户端的本地模型更新来训练全局共享模型。为了有效应对不同客户端的数据分布异质性问题，个性化联邦学习方法在传统联邦学习框架中引入了个性化组件，使每个客户端都能够根据其特定需求训练本地模型。本文结合分层抽样策略提出了一种新的个性化联邦学习方法，提高了模型训练的效率，同时减少了通信复杂度。本文建立了算法的收敛性。仿真实验表明，合理地选择分组数量和抽样策略能够显著加快算法的收敛速度，提高其精度。

Abstract

Federated learning is a distributed learning method protecting the data privacy.It trains a global shared model by aggregating local model updates from multiple clients.To deal with the problem of heterogeneity in data distribution among clients， personalized federated learning is an effective way introducing a personalized component into the traditional federated learning framework and enabling each client to train a local model according to its specific needs.In this paper， a new personalized federated learning method combining the stratified sampling strategy is proposed to improve the efficiency of model training and reduce communication complexity.Convergence of the algorithm is established and simulation experiments are implemented to show that reasonable selection of the number of groups and sampling strategy can significantly accelerate the convergence of the algorithm and improve its accuracy.

Graphical abstract

关键词

个性化联邦学习 / 分层抽样 / 非独立同分布 / 优化算法 / 收敛性分析

Key words

personalized federated learning / stratified sampling / non-IID / optimization algorithm / convergence analysis

联邦学习（Federated Learning，FL）是一种分布式学习框架，允许多个客户端（个人电子设备和企业数据库等）联合训练机器学习模型，无需共享其本地数据。与传统的全局模型训练方法一般将大规模数据集存储在昂贵的数据中心进行处理不同，联邦学习将数据和训练过程保留在本地客户端^［1］。通过这种方式，联邦学习可以在保护本地数据隐私的同时从这些设备上的大数据中进行学习，大大提高通信效率。因其在隐私保护和通信效率方面的优势，目前联邦学习已被广泛应用于金融、太赫兹信道估计^［2］等应用领域。

联邦平均（FedAvg）算法^［1］及其变体^［3-4］是目前联邦学习中最流行的方法。在广义FedAvg框架^［5］中，每个通信轮次服务器先将当前模型广播给客户端子集，然后客户端使用本地数据更新模型并发回本地模型更新，服务器再聚合客户端的模型更新并计算新的全局模型。此类算法在实践中受到欢迎的原因有很多，其中之一便是联邦学习与其他隐私保护技术具有良好兼容性^［6］，如差分隐私^［7］和安全聚合^［8-9］等。

在实际应用中，联邦学习面临的一个主要挑战是大量的异构客户端。具体来说，客户端拥有的本地数据集通常在这些设备之间非独立同分布^［10］。研究表明，FedAvg算法在处理非独立同分布数据时效率低下^［11］。

另一方面，个性化联邦学习在本地设备上训练各自的模型而非使用全局模型，不仅能提高模型的收敛性能，同时也能够使用户获得解决自身问题的模型^［12］。以下是几种典型的个性化联邦学习方法。

1） 全局-局部模型聚合^［12］。这种方法结合了全局和局部模型聚合策略，通过平衡全局一致性和局部适应性，该方法使用定制的全局模型来适应本地数据需求，提高个性化。

2） 基于元学习的pFL^［13］。这种方法侧重于通过“学习如何学习”各种任务来使模型能够快速适应新任务。通过元学习，模型可以利用来自多个任务的经验更好地为每个客户进行个性化。

3） 基于模型拆分的pFL^［14］。该方法将模型划分为子模型，每个子模型都针对不同的用户数据段进行训练，然后将这些子模型进行集成，在增强整体模型性能的同时保留特定数据特征的个性化。

现有的个性化联邦学习算法在以下几个方面存在不足。首先，组间数据分布异质性导致的低效问题。在非独立同分布应用场景中，客户端之间的数据差异性可能会导致全局模型无法充分表征数据的分布特性，影响模型的个性化效果和整体收敛速度。其次，设备异构性导致的低效问题。客户端的计算能力和通信能力可能存在显著差异，部分低效设备可能成为训练的瓶颈。这种设备间差异在全局模型聚合和通信调度中可能进一步加剧资源浪费和训练延迟。

作为一种高效客户端选择策略，分层抽样方法有望在解决上述的个性化联邦学习问题时展现独特优势。首先，分层抽样方法能够降低组间数据分布的异质性。将某些客户端划分为同一组并进行分层抽样能够显著减小抽样数据的分布差异性，不仅能提升模型对不同数据分布的适应性，还可以提高了模型训练的稳定性和准确性。其次，分层抽样方法能提升通信效率。根据客户端的计算和通信能力，分层抽样可以将通信速率和计算速度相似的客户端分组并抽样。这种策略能够在通信调度中有效避免因异构性造成的延迟，加快模型的训练速度。最后，分层抽样的分组机制还可以显著减少跨组通信的频率，进一步降低整体通信复杂度。图1和图2分别展示了随机抽样和分层抽样的特点，其中横坐标为抽样轮次，纵坐标为客户端指标。在随机抽样中，M=\mathrm{20}个客户端每次随机抽样N=\mathrm{5}个客户端进行训练，在多次抽样后可能有的客户端被抽样多次，有的客户端则很少被抽到，甚至未被抽取。另一方面，在分层抽样中，M=\mathrm{20}个客户端被分为\overline{K}=\mathrm{4}组，每组有M/\overline{K}=\mathrm{5}个客户端，按顺序从组内抽样N=\mathrm{2}个客户端进行训练，在多次抽样后仍可以保证每个客户端均被抽取且抽取次数相近。

本文主要对全局-局部模型聚合方法进行进一步的探索。基于分层抽样，本文提出了一种个性化联邦学习模型，对模型在局部分布上的收敛性进行了分析。该收敛性取决于混合参数、局部分布之间的差异、训练数据量以及分层抽样的策略。为了学习个性化模型，本文提出了一种通信效率高的优化算法，算法通过在学习过程中利用局部和全局模型之间的相关性来自适应地学习模型。仿真实验显示，本文提出的个性化方案优于常用的联邦优化算法的全局和局部模型。

1 问题描述

1.1 目标

考虑一个包含\mathrm{ℳ}个客户端的联邦学习问题，第m个客户端收集了本地n_m个数据，即

{\mathrm{𝒟}}_m={\left\{\left({\mathbf{x}}_{m,k},{\mathbf{y}}_{m,k}\right)\right\}}_{k=\mathrm{1}}^{n_m}（1）

其中，\left({\mathbf{x}}_{m,k},{\mathbf{y}}_{m,k}\right)\in {\mathbb{R}}^S\times {\mathbb{R}}^C是{\mathrm{𝒟}}_m中的第k个样本，S表示输入数据的维度，C表示数据类别的数量，每个数据集有n_m个数据。定义

{\mathbf{x}}_m=[x_{m,\mathrm{1}},x_{m,\mathrm{2}},\dots ,x_{m,n_m}]，

{\mathbf{y}}_m=[y_{m,\mathrm{1}},y_{m,\mathrm{2}},\dots ,y_{m,n_m}]，

并令h(w;·)表示含有模型参数的机器学习模型。我们考虑通过最小化联合成本函数来训练全局模型w，即

\underset{w}{\mathrm{m}\mathrm{i}\mathrm{n}}{\displaystyle \sum _{m=\mathrm{1}}^M}\frac{\mathrm{1}}{M}{f}_m\left(w\right)（2）

其中

f_m\left(w\right)\triangleq l\left(h\left(w;{\mathbf{x}}_m\right),{\mathbf{y}}_m\right)（3）

是每个客户端m的局部成本函数，M是客户端的总数，l是针对一组数据样本的机器学习模型h的损失函数。

1.2 个性化模型

在局部模型和全局模型的帮助下，我们为每个客户端设计一个个性化模型。具体来说，假设对于任何输入\mathbf{x}，个性化模型的输出h\left({\overline{v}}_m;·\right)均可以表示为局部模型h\left(v_m;·\right)和全局模型h(w;·)输出的凸组合，即

\begin{array}{l}h\left({\overline{v}}_m;\mathbf{x}\right)=\\ {\beta }_{m}h\left(v_m;\mathbf{x}\right)+\left(\mathrm{1}-{\beta }_m\right)h(w;\mathbf{x}),\forall \mathbf{x}\end{array}（4）

其中{\beta }_m\in [\mathrm{0,1}]是预定义的混合系数。本文研究的问题是：找到一个局部模型，使（4）式成为局部成本函数的最小化器。具体来说，可以将问题写成

\underset{v_m}{\mathrm{m}\mathrm{i}\mathrm{n}}{\displaystyle \sum _{m=\mathrm{1}}^M}\frac{\mathrm{1}}{M}{F}_m\left(v_m,w^{\mathrm{*}}\right)（5）

使得

w^{\mathrm{*}}=\underset{w}{\mathrm{a}\mathrm{r}\mathrm{g}\mathrm{m}\mathrm{i}\mathrm{n}}{\displaystyle \sum _{m=\mathrm{1}}^M}\frac{\mathrm{1}}{M}{f}_m\left(w\right)（6）

其中

F_m\left(v_m,w\right)\triangleq l\left({\beta }_{m}h\left(v_m\right)+\left(\mathrm{1}-{\beta }_m\right)h(w),{\mathbf{y}}_m\right)（7）

这里的f_m(·)在式（3）中已定义。为便于表达，本文用h(w)表示h(w;·)。需要说明的是，本文将采用（5）中的全局模型和局部模型的凸组合来构建个性化模型，而非对它们的参数进行组合。

问题（5）是一个典型的双层优化问题，本文将采用两阶段优化的形式来解决这个问题。具体来说，在第一阶段，通过求解子问题（5）来优化全局模型参数w^{\mathrm{*}}，然后以w^{\mathrm{*}}作为输入；在第二阶段，每个客户端m再次求解（5），以确定其局部最优模型参数v_m。

2 算法

为了减少个性化联邦学习的通信复杂度，一般需要从\mathrm{ℳ}个客户端中随机选择N个进行通信。然而，这将导致以下问题：1） 组内数据分布的异质性可能较强；2） 客户端的计算能力和通信能力可能存在显著差异。为此，我们考虑以循环客户端参与的方式保证在选择合适的分组后能够实现通信效率提升。

2.1 分层抽样

将\mathrm{ℳ}个客户端分成\overline{K}个不重叠的组，每个组包含\mathrm{ℳ}/\overline{K}个客户端。客户端组用\sigma (i),i\in [\overline{K}]表示，组的顺序\sigma (\mathrm{1}),\dots ,\sigma (\overline{K})在整个训练过程中是固定的，以模拟客户端分层抽样结构。

在每个通信轮次中，假如选择第\sigma (i)个客户端组，服务器就从第\sigma (i)组中随机选择N个客户端。当这些客户端被选择后，至少在接下来的\overline{K}-\mathrm{1}轮中将不会参与训练。本文将“循环周期”定义为服务器按顺序处理所有客户端组\{\sigma (i),i\in [\overline{K}]\}一次的间隔。换句话说，每个循环周期由\overline{K}个通信轮次组成。本文用k表示循环周期的索引，i\in [\overline{K}]表示循环周期内当前可用的客户端组的索引。然后，服务器将全局模型w^{(k,i-\mathrm{1})}发送给从客户端组\sigma (i)中选择出的N个客户端集合{\mathrm{𝒮}}^{(k,i)}，进行本地训练。

2.2 算法设计

本文算法主要分为两步。

第1步 对于第(k,i)轮，服务器将全局模型参数w^{(k,i-\mathrm{1})}广播到S^{(k,i)}的客户端，每个选定的客户端m\in {\mathrm{𝒮}}^{(k,i)}依次获得以下子问题的近似解：

w_m^{(k,i)}=\underset{w_m}{\mathrm{a}\mathrm{r}\mathrm{g}\mathrm{m}\mathrm{i}\mathrm{n}}{f}_m\left(w_m\right)（8）

v_m^{(k,i)}=\underset{v_m}{\mathrm{a}\mathrm{r}\mathrm{g}\mathrm{m}\mathrm{i}\mathrm{n}}{F}_m\left(v_m,w_m^{(k,i)}\right)（9）

第2步 服务器收集并取某个平均值，用于下一轮更新。实际上，使用局部 SGD 来近似问题（8）和（9）的解就足够了。具体来说，每个客户端m\in {\mathrm{𝒮}}^{(k,i)}对w_m^{(k,i)}和v_m^{(k,i)}执行\tau \ge \mathrm{1}个连续SGD步骤，即对于t=\mathrm{0},\cdots ,\tau -\mathrm{1}，

w_m^{(k,i,t+\mathrm{1})}=w_m^{(k,i,t)}-\eta g_m\left(w_m^{(k,i,t)}\right)（10）

v_m^{(k,i,t+\mathrm{1})}=v_m^{(k,i,t)}-{\eta }_c{G}_m\left(v_m^{(k,i,t)},w_m^{(k,i,t)}\right)（11）

其中，\eta >\mathrm{0}和{\eta }_c>\mathrm{0}是两个步长， g_m(w_m^{(k,i,t)})表示小批量\mathrm{𝒮}上f_m(·)的随机梯度下降（Stochastic Gradient Descent，SGD），G_m\left(·,·\right)是F_m在\mathrm{𝒮}上的 SGD，初始值定义为

w_m^{(k,i,\mathrm{0})}=w^{(k,i)},v_m^{(k,i,\mathrm{0})}=v_m^{(k-\mathrm{1},i,\tau )}（12）

我们定义

w_m^{(k,i)}=w_m^{(k,i,\tau )},v_m^{(k+\mathrm{1},i)}=v_m^{(k,i,\tau )},\forall m\in {\mathrm{𝒮}}^{(k,i)}。

算法1 分层抽样学习算法

输入： 混合权重{\beta }_{\mathrm{1}},\cdots ,{\beta }_m，客户端训练步数\tau，客户端组\sigma (i),i\in [\overline{K}]，全局模型w^{(\mathrm{0,0})}，本地模型v_m^{(\mathrm{0,0})},m\in [M]；

输出： 全局模型w^{(K+\mathrm{1,0})}和本地模型v_m^{(K,i)}。

1） 从 k=\mathrm{1} 到 K：

2）  从 i=\mathrm{1} 到 \overline{K}：

3）  选出客户端组\sigma (i)中N个客户端的集合S^{(k,i)}；

4） 若m\in S^{(k,i)}则

5） 用\mathrm{L}\mathrm{U}函数更新w_m^{(k,i)},v_m^{(k,i)}，即

\mathrm{L}\mathrm{U}\left(m,w^{(k,i-\mathrm{1})},v_m^{(k-\mathrm{1},i)}\right)；

6） 否则v_m^{(k,i)}=v_m^{(k-\mathrm{1},i)}；

7） 结束判断；

8）   w^{(k,i)}=\frac{\mathrm{1}}{N}\sum _{m\in S^{(k,i)}}{w}_m^{(k,i)}；

9） 结束循环；

10）   w^{(k,\overline{K})}=w^{(k+\mathrm{1,0})}

11） 结束循环。

在算法1中，\mathrm{L}\mathrm{U}\left(m,w,v\right)函数是如下定义的：

从j=\mathrm{1}到\tau

取出小批量样本{\xi }_m；

w\leftarrow w-\eta \nabla f\left(w,{\xi }_m\right)；

v\leftarrow v-{\eta }_{\mathrm{c}}{\nabla }_v{F}_m\left(w,v,{\xi }_m\right)；

结束循环。

3 收敛性分析

本节给出算法的收敛性分析。为此我们需要以下的假设，这些假设在 FL相关文献中是常见的^{［9， 12］}。

假设1（L-光滑性） 客户端的局部目标函数f_{\mathrm{1}}(w),\cdots ,f_M(w)都是L-光滑的，即对所有的m，w_{\mathrm{1}}和w_{\mathrm{2}}，均有

‖\nabla f_m(w_{\mathrm{1}})-\nabla f_m\left(w_{\mathrm{2}}\right)‖\le L‖w_{\mathrm{1}}-w_{\mathrm{2}}‖（13）

同时，局部成本函数F_m(·,·)L_F-光滑，函数l(·)\tilde{L}-光滑，函数h(·)L_h-光滑。

假设2（\mu-Polyak-Łojasiewicz假设） 对于某个\mu \ge \mathrm{0}，个性化目标满足

\frac{\mathrm{1}}{\mathrm{2}}\Vert \nabla F(v,w){\Vert }^{\mathrm{2}}\ge \mu \left(F(v,w)-\underset{v^{\text{'}},w^{\text{'}}}{\mathrm{m}\mathrm{i}\mathrm{n}}F\left(v^{\text{'}},w^{\text{'}}\right)\right)（14）

假设3（组内和组间数据异质性） 存在常数\gamma ,\alpha \ge \mathrm{0}，使得对于所有w，i\in [\overline{K}]以及所有m\in \sigma (i)，有

‖\nabla f_m(w)-\frac{\mathrm{1}}{|\sigma (i)|}\sum _{m\in \sigma (i)}\nabla f_m(w)‖\le \gamma（15）

和

‖\frac{\mathrm{1}}{|\sigma (i)|}\sum _{m\in \sigma (i)}\nabla f_m(w)-\nabla f(w)‖\le \alpha（16）

若假设3成立，则存在\nu =\gamma +\alpha \ge0，使得对所有客户端m\in [M]以及所有w，均有

‖\nabla f_m(w)-\frac{\mathrm{1}}{M}{\displaystyle \sum _{m=\mathrm{1}}^M}\nabla f_m(w)‖\le \nu。

假设4（有界方差） 对于随机均匀采样的数据样本，问题（5）所得的随机梯度是无偏的，且具有有界方差，即

\mathbb{E}\left[{\nabla }_w{f}_m\left(w_m;{\xi }_m\right)\right]={\nabla }_w{f}_m\left(w_m\right)（17）

\mathbb{E}\left[{‖{\nabla }_w{f}_m\left(w_m;{\xi }_m\right)-{\nabla }_w{f}_m\left(w_m\right)‖}^{\mathrm{2}}\right]\le {\sigma }^{\mathrm{2}}（18）

\mathbb{E}\left[{\nabla }_v{F}_m\left(v_m^i,w_m;{\xi }_m\right)\right]={\nabla }_v{F}_m\left(v_m^i,w_m\right)（19）

以及

\mathbb{E}\left[{‖{\nabla }_v{F}_m\left(v_m^i,w_m;{\xi }_m\right)-{\nabla }_v{F}_m\left(v_m^i,w_m\right)‖}^{\mathrm{2}}\right]\le {\sigma }^{\mathrm{2}}（20）

由假设4，我们能够得到f_m\left(w\right)是无偏的，方差受{\sigma }^{\mathrm{2}}/S约束，其中S是本地数据集的小批量。类似地，F_m\left(v,w\right)也是无偏的，方差受{\sigma }^{\mathrm{2}}/S约束。

假设5（模型差异） 对于任意两个数据输入{\mathbf{x}}_{\mathrm{1}}、{\mathbf{x}}_{\mathrm{2}}，它们的模型输出之间的距离是有界的，即

‖h\left({\mathbf{x}}_{\mathrm{1}}\right)-h\left({\mathbf{x}}_{\mathrm{2}}\right)‖\le \mathrm{\Gamma }（21）

在实际应用中，假设5比较温和，因为h(\mathbf{x})的输出通常位于有界空间中。例如，对于分类任务，有h(\mathbf{x})\in {[\mathrm{0,1}]}^n，从而\mathrm{\Gamma }\le \mathrm{2}。

下面我们给出满足\mu-Polyak-Łojasiewicz假设的光滑函数的局部梯度下降算法的收敛性分析。

定理3.1 设步长\eta和{\eta }_c满足

\eta \le \mathrm{m}\mathrm{i}\mathrm{n}\left\{\frac{\mathrm{1}}{\mathrm{10}\tau LN\overline{K}},\frac{\mathrm{2}}{\mathrm{3}L\tau },\frac{\mathrm{1}}{L}\right\},

\frac{M\mathrm{l}\mathrm{o}\mathrm{g}(M{K}^{\mathrm{2}})}{\tau \mu NK\overline{K}}\le {\eta }_c\le \mathrm{m}\mathrm{i}\mathrm{n}\left\{\frac{\mathrm{1}}{\mathrm{2}\tau L_F},\frac{M\eta L^{\mathrm{2}}{N}^{\mathrm{2}}{\overline{K}}^{\mathrm{2}}}{L_F^{\mathrm{2}}}\right\},

且

D_{\mathrm{0}}=\frac{\left(\mathrm{1}-{\beta }_m\right)\left({\beta }_m\tilde{L}{L}_h\mathrm{\Gamma }+L\right)}{\mathrm{2}L}\le \mathrm{4}。

在假设1~假设5成立的条件下，设

\eta =O\left(\mathrm{1}/NT\right),S=O(M),

其中T=K\overline{K}， 则以下收敛结果成立：

\mathbb{E}\left[F_m\left(v_m^{(K,i)},w^{(K,i)}\right)\right]-F_m^{\mathrm{*}}+\frac{{\overline{K}}^{\mathrm{2}}\left(E\left[F_m\left(v^{(\mathrm{0,0})},w^{(\mathrm{0,0})}\right)\right]-f_m^{\mathrm{*}}\right)}{M{T}^{\mathrm{2}}}+O\left(\frac{{\sigma }^{\mathrm{2}}{L}^{\mathrm{2}}{\tau }^{\mathrm{2}}{\overline{K}}^{\mathrm{3}}}{M^{\mathrm{2}}{N}^{\mathrm{3}}{T}^{\mathrm{3}}}\right)+O\left(\frac{M^{\mathrm{2}}{N}^{\mathrm{2}}{L}^{\mathrm{4}}{\overline{K}}^{\mathrm{4}}{\sigma }^{\mathrm{2}}\tau }{L_F^{\mathrm{2}}\mu S{T}^{\mathrm{2}}}\right)+

O\left(\frac{\mathrm{5}{\sigma }^{\mathrm{2}}{\tau }^{\mathrm{2}}{L}_F^{\mathrm{2}}}{\mathrm{4}\mu M{N}^{\mathrm{2}}{T}^{\mathrm{2}}}\right)+O\left(\frac{L\tau {\sigma }^{\mathrm{2}}\overline{K}}{NT}\right)+O\left(\frac{L^{\mathrm{2}}{(\tau -\mathrm{1})}^{\mathrm{2}}\tau {\nu }^{\mathrm{2}}}{N^{\mathrm{2}}{T}^{\mathrm{2}}}\right)+O\left(\frac{L^{\mathrm{2}}{(\tau -\mathrm{1})}^{\mathrm{2}}\tau {\alpha }^{\mathrm{2}}}{N^{\mathrm{2}}{T}^{\mathrm{2}}}\right)+O\left(\frac{L^{\mathrm{2}}{\overline{K}}^{\mathrm{2}}{\tau }^{\mathrm{3}}{\alpha }^{\mathrm{2}}}{T^{\mathrm{2}}}\right)+O\left(\frac{L{\tau }^{\mathrm{2}}{\gamma }^{\mathrm{2}}}{NT}\left(\frac{M/\overline{K}-N}{M/\overline{K}-\mathrm{1}}\right)\right)+\underset{(\mathrm{a})}{\underbrace{\frac{NK{\mathrm{\Gamma }}^{\mathrm{2}}\tilde{L}}{\mathrm{l}\mathrm{o}\mathrm{g}\left(M{K}^{\mathrm{2}}\right)}{\beta }_m\left(\mathrm{1}-{\beta }_m\right)}}（22）

证明 篇幅所限，我们仅给出证明思路。首先，由假设1易得

\begin{array}{l}\mathbb{E}\left[F_m\left(v_m^{(k+\mathrm{1},i)},w^{(k+\mathrm{1},i)}\right)\right]-\mathbb{E}\left[F_m\left(v_m^{(k,i)},w^{(k,i)}\right)\right]\le \frac{N\overline{K}}{M}\mathbb{E}\left[〈{\nabla }_v{F}_m\left(v_m^{(k,i)},w^{(k,i)}\right),{\tilde{v}}_m^{(k+\mathrm{1},i)}-v_m^{(k,i)}〉\right]+\\ \mathbb{E}\left[〈{\nabla }_w{F}_m\left(v_m^{(k,i)},w^{(k,i)}\right),\left(w^{(k+\mathrm{1},i)}-w^{(k,i)}\right)〉\right]+\\ \frac{N\overline{K}{L}_F}{\mathrm{2}M}\mathbb{E}\left[{‖{\tilde{v}}_m^{(k+\mathrm{1},i)}-v_m^{(k,i)}‖}^{\mathrm{2}}\right]+\\ \frac{L_F}{\mathrm{2}}\mathbb{E}\left[{‖w^{(k+\mathrm{1},i)}-w^{(k,i)}‖}^{\mathrm{2}}\right]\end{array}（23）

接下来，可以证明以下几个不等式成立：

\mathbb{E}\left[{‖{\tilde{v}}_m^{(k+\mathrm{1},i)}-v_m^{(k,i)}‖}^{\mathrm{2}}\right]\le {\eta }_c^{\mathrm{2}}{\tau }^{\mathrm{2}}\mathbb{E}\left[{‖\frac{\mathrm{1}}{\tau }{\displaystyle \sum _{t=\mathrm{0}}^{\tau }}{\nabla }_v{F}_m\left({\tilde{v}}_m^{(k,i,t)},{\tilde{w}}_m^{(k,i,t)}\right)‖}^{\mathrm{2}}\right]+\frac{{\eta }_c^{\mathrm{2}}{\tau }^{\mathrm{2}}{\sigma }^{\mathrm{2}}}{S}（24）

\begin{array}{l}\mathbb{E}\left[〈{\nabla }_v{F}_m\left(v_m^{(k,i)},w^{(k,i)}\right),{\tilde{v}}_m^{(k+\mathrm{1},i)}-v_m^{(k,i)}〉\right]\le -\frac{{\eta }_c\tau (\mathrm{1}-{\eta }_c^{\mathrm{2}}{\tau }^{\mathrm{2}}{L}_F^{\mathrm{2}})}{\mathrm{2}}\mathbb{E}\left[{‖\frac{\mathrm{1}}{\tau }{\displaystyle \sum _{t=\mathrm{0}}^{\tau }}{\nabla }_v{F}_m\left({\tilde{v}}_m^{(k,i,t)},{\tilde{w}}_m^{(k,i,t)}\right)‖}^{\mathrm{2}}\right]+\\ \frac{{\eta }_c^{\mathrm{3}}{L}_F^{\mathrm{2}}{\sigma }^{\mathrm{2}}{\tau }^{\mathrm{2}}}{\mathrm{2}S}+\frac{{\eta }_c{L}_F^{\mathrm{2}}\tau }{\mathrm{2}}\mathbb{E}\left[{\displaystyle \sum _{t=\mathrm{0}}^{\tau }}\frac{\mathrm{1}}{\tau }{‖{\tilde{w}}_m^{(k,i,t)}-w^{(k,i)}‖}^{\mathrm{2}}\right]-\frac{{\eta }_c\tau }{\mathrm{2}}\mathbb{E}\left[{‖{\nabla }_v{F}_m\left(v_m^{(k,i)},w^{(k,i)}\right)‖}^{\mathrm{2}}\right]\end{array}（25）

（26）

根据以上几个不等式可证，当{\eta }_c\tau L_F\le \frac{\mathrm{1}}{\mathrm{2}}及L_F\le (\mathrm{1}-{\beta }_m)L时，有

（27）

（27）式右端项可由文献［8］得到，重新排列其两边可得

\begin{array}{l}\mathbb{E}\left[F_m\left(v_m^{(k,i)},w^{(k,i)}\right)\right]-\mathbb{E}\left[F_m\left(v_m^{(k+\mathrm{1},i)},w^{(k+\mathrm{1},i)}\right)\right]-\frac{{\eta }_c\tau N\overline{K}}{\mathrm{2}M}\mathbb{E}\left[{‖{\nabla }_v{F}_m\left(v_m^{(k,i)},w^{(k,i)}\right)‖}^{\mathrm{2}}\right]+\\ \frac{{\eta }_c{L}_F({\eta }_c{L}_F+\mathrm{1}){\sigma }^{\mathrm{2}}\tau }{S}+\frac{\mathrm{5}{\eta }^{\mathrm{2}}{\sigma }^{\mathrm{2}}{\tau }^{\mathrm{2}}{L}_F^{\mathrm{2}}}{\mathrm{2}S}+\frac{\mathrm{328}ML{\eta }^{\mathrm{2}}\overline{K}{\sigma }^{\mathrm{2}}}{\mathrm{25}{\eta }_c}+\frac{\mathrm{206}M{\eta }^{\mathrm{3}}{L}^{\mathrm{2}}{(\tau -\mathrm{1})}^{\mathrm{2}}{\nu }^{\mathrm{2}}}{{\eta }_c}+\\ \frac{\mathrm{2}M{\eta }^{\mathrm{3}}{L}^{\mathrm{2}}{(\overline{K}-\mathrm{1})}^{\mathrm{2}}{(\tau -\mathrm{1})}^{\mathrm{2}}{\alpha }^{\mathrm{2}}}{\mathrm{2}{\eta }_c{\overline{K}}^{\mathrm{2}}}+\frac{\mathrm{262}M{\eta }^{\mathrm{3}}{L}^{\mathrm{2}}{N}^{\mathrm{2}}\overline{K}(\overline{K}-\mathrm{1}){\tau }^{\mathrm{2}}{\alpha }^{\mathrm{2}}}{{\eta }_c}+\frac{\mathrm{16}ML{\eta }^{\mathrm{2}}\tau {\gamma }^{\mathrm{2}}}{\mathrm{5}{\eta }_c}\left(\frac{M/\overline{K}-N}{M/\overline{K}-\mathrm{1}}\right)+\\ \frac{\mathrm{2}M{\mathrm{\Gamma }}^{\mathrm{2}}\tilde{L}}{{\eta }_c\tau N\overline{K}}{\beta }_m\left(\mathrm{1}-{\beta }_m\right)\end{array}（28）

由假设2可得

\begin{array}{l}\mathbb{E}\left[F_m\left(v_m^{(k+\mathrm{1},i)},w^{(k+\mathrm{1},i)}\right)\right]-f_m^{\mathrm{*}}\le \\ \left(\mathrm{1}-\frac{{\eta }_c\tau N\overline{K}\mu }{M}\right)\left(\mathbb{E}\left[F_m\left(v_m^{(k,i)},w^{(k,i)}\right)\right]-f_m^{\mathrm{*}}\right)-p\end{array}（29）

其中的p代表式（28）右端第二项及之后的项。注意到式（29）已经构成一个迭代格式，我们考虑将其迭代K次后作为最优解的界，最终能够证明以下的不等式：

\mathbb{E}\left[F_m\left(v_m^{(K,i)},w^{(K,i)}\right)\right]-F_m^{\mathrm{*}}\le \frac{E\left[F_m\left(v^{(\mathrm{0,0})},w^{(\mathrm{0,0})}\right)\right]-F_m^{\mathrm{*}}}{M{K}^{\mathrm{2}}}+\frac{\mathrm{5}{\eta }^{\mathrm{3}}{\sigma }^{\mathrm{2}}{L}^{\mathrm{2}}{\tau }^{\mathrm{2}}}{\mathrm{4}SM{K}^{\mathrm{2}}}+\frac{M^{\mathrm{2}}{\eta }^{\mathrm{2}}{L}^{\mathrm{4}}{N}^{\mathrm{4}}{\overline{K}}^{\mathrm{4}}{\sigma }^{\mathrm{2}}\tau }{\mathrm{2}{L}_F^{\mathrm{2}}\mu S}+\frac{\mathrm{5}{\eta }^{\mathrm{2}}{\sigma }^{\mathrm{2}}{\tau }^{\mathrm{2}}{L}_F^{\mathrm{2}}}{\mathrm{4}\mu S}+\frac{\mathrm{162}L\tau NK{\eta }^{\mathrm{2}}{\overline{K}}^{\mathrm{2}}{\sigma }^{\mathrm{2}}}{\mathrm{25}\mathrm{l}\mathrm{o}\mathrm{g}\left(M{K}^{\mathrm{2}}\right)}+\frac{\mathrm{103}{\eta }^{\mathrm{3}}{L}^{\mathrm{2}}{(\tau -\mathrm{1})}^{\mathrm{2}}\tau NK\overline{K}{\nu }^{\mathrm{2}}}{\mathrm{l}\mathrm{o}\mathrm{g}\left(M{K}^{\mathrm{2}}\right)}+\frac{{\eta }^{\mathrm{3}}{L}^{\mathrm{2}}{(\overline{K}-\mathrm{1})}^{\mathrm{2}}K{(\tau -\mathrm{1})}^{\mathrm{2}}\tau N{\alpha }^{\mathrm{2}}}{\mathrm{2}\mathrm{l}\mathrm{o}\mathrm{g}\left(M{K}^{\mathrm{2}}\right)\overline{K}}+\frac{\mathrm{131}{\eta }^{\mathrm{3}}{L}^{\mathrm{2}}{N}^{\mathrm{3}}K{\overline{K}}^{\mathrm{2}}(\overline{K}-\mathrm{1}){\tau }^{\mathrm{3}}{\alpha }^{\mathrm{2}}}{\mathrm{l}\mathrm{o}\mathrm{g}\left(M{K}^{\mathrm{2}}\right)}+\frac{\mathrm{8}L{\eta }^{\mathrm{2}}{\tau }^{\mathrm{2}}NK\overline{K}{\gamma }^{\mathrm{2}}}{\mathrm{5}\mathrm{l}\mathrm{o}\mathrm{g}\left(M{K}^{\mathrm{2}}\right)}\left(\frac{M/\overline{K}-N}{M/\overline{K}-\mathrm{1}}\right)+\frac{K{\mathrm{\Gamma }}^{\mathrm{2}}\tilde{L}}{\mathrm{l}\mathrm{o}\mathrm{g}\left(M{K}^{\mathrm{2}}\right)}{\beta }_m\left(\mathrm{1}-{\beta }_m\right)（30）

证毕。

推论3.2 考虑定理3.1中影响收敛的主项

C_{\overline{K}}=O\left(\frac{L\tau {\sigma }^{\mathrm{2}}\overline{K}}{NT}\right)+O\left(\frac{L{\tau }^{\mathrm{2}}{\gamma }^{\mathrm{2}}}{NT}\left(\frac{M/\overline{K}-N}{M/\overline{K}-\mathrm{1}}\right)\right),

并假设组内数据异质性\gamma满足

{\gamma }^{\mathrm{2}}\ge \frac{{(M-\mathrm{1})}^{\mathrm{2}}{\sigma }^{\mathrm{2}}}{\tau M(N-\mathrm{1})},

则有C_{\overline{K}}\ge C_{\overline{K}+\mathrm{1}}，，即\overline{K}越大收敛速率越快。

注1 定理3.1的主导项\widetilde{\mathrm{𝒪}}(\mathrm{1}/T)受两个因素影响：组内数据异质性{\gamma }^{\mathrm{2}}和随机梯度方差 {\sigma }^{\mathrm{2}}。这种依赖性意味着，对于\sigma >\mathrm{0}，即使\overline{K}=M/N或{\gamma }^{\mathrm{2}}\approx \mathrm{0}，也无法实现\widetilde{\mathrm{𝒪}}(\mathrm{1}/T^{\mathrm{2}})收敛速度。但是，如果考虑通信的复杂性就会发现当\gamma =\mathrm{0}或\overline{K}=M/N时，项 O\left(\frac{L{\tau }^{\mathrm{2}}{\gamma }^{\mathrm{2}}}{NT}\cdot \frac{M/\overline{K}-N}{M/\overline{K}-\mathrm{1}}\right)将变为零。因此，组内异质性确实会影响通信成本。当\overline{K}>\mathrm{1}时，本文的方法的通信成本相比本地SGD较低。

注2 （22）式右端项中的（a）是由于同时训练w和v而产生的。如果在优化v期间w保持不变，这些项就会消失。然而，在算法中联合优化v和w可以显著提高效率，因为它允许我们同时获得理想的参数。因此，收敛性能和算法效率之间存在权衡。幸运的是，一旦w发生最小变化，就可以通过运行几轮额外的轮次来更新v来解决这个问题。这样，本文提出的算法可以同时实现高效率和模型准确率。

注3 本文算法对于非凸损失函数实现了O(\mathrm{1}/NT)的次线性收敛速度。该收敛速率表明，随着客户端数N和通信轮数T的增加，算法的收敛性能会提升，从而确保大规模部署中的可扩展性和效率。

4 仿真分析

4.1 实验设置

1） 数据集。我们使用经典的EMNIST数据集进行实验。EMNIST数据集有62个标签，总共有500个客户端，每轮通信分别选择10个客户端。

2） 模型。我们使用卷积神经网络，用两个卷积层和两个全连接层作为输出，其中卷积核大小为2*2。

3） 参数设置。我们使用狄利克雷分布 \mathrm{D}\mathrm{i}\mathrm{r}\mathrm{K}(\alpha )^［15］在客户端之间划分数据，其中\alpha决定了客户端之间数据异质性的程度。较小的\alpha表示数据异质性较大。在具体实验中，我们设置固定混合参数\beta ={\beta }_m，批处理大小为32，\eta =\mathrm{0.05}。

4.2 结果与分析

1） \overline{K}的影响。图3和图4展示了不同\overline{K}值和客户端本地程序对EMNIST数据集的训练损失和测试准确率，这些程序适用于高数据异质性（\alpha=0.05）。\overline{K}=\mathrm{1}的结果类似于APFL算法的结果，且（对于EMNIST，M/N=\mathrm{50}）， 意味着客户端进行分层抽样。较高的\overline{K}表示服务器在一个周期内访问了更多的客户端组。结果表明，对于高数据异质性，对于所有不同的客户端本地程序，较高\overline{K}值可实现更好的测试准确率，与理论分析一致。

2） \beta的影响。我们知道，当\beta =\mathrm{0}时，算法等同于有分层抽样的FedAvg算法。从表2中可以看到，只要\beta \ne \mathrm{0}，就能获得更好的收敛速度和测试准确率。同时，不同的\alpha值不会导致准确率相差太大。

5 结论

本文设计了一种高效的联邦学习通信优化算法，以便同时训练全局模型和局部模型。该算法不仅能加快联邦学习框架的收敛速度，还有效降低了通信成本，尤其是在客户端参与度受限的情况下。收敛性分析进一步强调了算法通过调整混合参数来平衡全局一致性和局部适应性的能力，确保在不同数据分布下实现稳健的个性化模型。此外，算法在准确性和效率之间实现了良好的权衡，不仅能保证实用性，还显著提升了模型性能。本文在EMNIST数据集上进行仿真实验，成功验证了所提算法的有效性。与现有的全局和个性化联邦学习算法（如FedAvg算法和APFL算法）相比，基于分层抽样的个性化联邦学习算法在数据异质性显著的场景中展现出更快速的收敛速度和更高的测试准确率。仿真实验结果还揭示了关键参数（例如分组数\overline{K}和混合系数\beta）对算法性能的重要影响。

参考文献

[1] McMahan H B， Moore E， Ramage D， et al.Communication-efficient learning of deep networks from decentralized data ［C］// Proceedings of the 20th International Conference on Artificial Intelligence. New York： PMLR， 2017： 1273-1282.

[2] Elbir A M， Shi W， Mishra K V， et al.Federated multi-task learning for THz wideband channel and DoA estimation ［C］// 2023 IEEE International Conference on Acoustics， Speech， and Signal Processing Workshops （ICASSPW）.Piscataway： IEEE， 2023： 1-5.

[3] Sahu A K， Li T， Zaheer M， et al.Federated optimization in heterogeneous networks ［J］.Proc Mach Learn Syst， 2020， 2： 429-450.

[4] Reddi S， Charles Z， Zaheer M， et al.Adaptive federated optimization ［EB/OL］.［2024-09-24］.

[5] Wang J， Charles Z， Xu Z， et al.A field guide to federated optimization ［EB/OL］.［2024-09-24］.

[6] Bonawitz K， Eichner H， Grieskamp W， et al.Towards federated learning at scale： System design ［J］.Proceedings of Machine Learning and Systems， 2019， 1： 374-388.

[7] McMahan H B， Ramage D， Talwar K， et al.Learning differentially private recurrent language models ［EB/ OL］.［2024-09-24］.

[8] Bonawitz K， Ivanov V， Kreuter B， et al.Practical secure aggregation for federated learning on user-held data ［EB/OL］.［2024-09-24］.

[9] Cho Y J， Sharma P， Joshi G， et al.On the convergence of federated averaging with cyclic client participation ［C］// Proceedings of the 40th International Conference on Machine Learning. New York： PMLR， 2023： 5677-5721.

[10] Zhao Z， Feng C， Hong W， et al.Federated learning with non-IID data in wireless networks ［J］.IEEE Trans Wireless Commun， 2021， 21（3）： 1927-1942.

[11] Zhang Z， Ma S， Yang Z， et al.Robust semi-supervised federated learning for images automatic recognition in Internet of Drones ［J］.IEEE Internet of Things Journal， 2022， 10（7）： 5733-5746.

[12] Deng Y， Kamani M M， Mahdavi M.Adaptive personalized federated learning ［EB/OL］.［2024-09-24］.

[13] Fallah A， Mokhtari A， Ozdaglar A.Personalized federated learning with theoretical guarantees： A model-agnostic meta-learning approach ［J］.Advances in Neural Information Processing Systems， 2020， 33： 3557-3568.

[14] Zhang J， Hua Y， Wang H， et al.GPFL： Simultaneously learning global and personalized feature information for personalized federated learning ［C］// Proceedings of the IEEE/CVF International Conference on Computer Vision.Piscataway： IEEE， 2023： 5041-5051.

[15] Hsu T-M H， Qi H， Brown M.Measuring the effects of non-identical data distribution for federated visual classification ［EB/OL］.［2024-09-24］.

基金资助

科技部重点研发计划(2020YFA0714003)

四川大学青年教师科技创新能力提升项目(2024SCUQJTX002)