FedGen: Data-Free Knowledge Distillation for Heterogeneous Federated Learning

联邦学习（Federated Learning, FL）是一种去中心化的机器学习范式，其中全局服务器通过迭代地聚合局部用户的模型参数来工作，而无需访问这些数据。用户异构性给联邦学习带来了显著挑战，可能导致全局模型偏离并收敛缓慢。最近，知识蒸馏技术被提出以解决这一问题，它通过利用来自异构用户的聚合知识来精炼服务器模型，而不是直接聚合他们的模型参数。然而，这种方法依赖于一个代理数据集，这使得它在实际应用中除非满足这一先决条件，否则并不可行。此外，集成知识并未被充分利用来指导局部模型学习，这可能会反过来影响聚合模型的质量。

在本文中，我们提出了一种无数据知识蒸馏方法来解决异构联邦学习问题。在该方法中，服务器学习一个轻量级的生成器，以无数据的方式集成用户信息，然后将该生成器广播给用户，利用所学知识作为归纳偏置来规范局部训练。基于理论启示的实证研究表明，与最先进的方法相比，我们的方法能够以更少的通信轮次促进联邦学习，从而获得更好的泛化性能。

原文论文地址：https://arxiv.org/pdf/2105.10056.pdf

1. Introduction

联邦学习（Federated Learning, FL）是一种有效的机器学习方法，它能够实现计算和数据资源的去中心化。经典的联邦学习，以FedAvg（McMahan et al., 2017）为代表，通过迭代地平均分布式局部用户模型的参数来获得一个聚合模型，因此省去了访问这些数据的需要。作为一种通信高效且保护隐私的学习方案，联邦学习已经展示了其在促进现实世界应用中的潜力，包括医疗保健（Sheller et al., 2020）、生物识别（Aggarwal et al., 2021）和自然语言处理（Hard et al., 2018; Ammad-Ud-Din et al., 2019）等，仅举几例。

尽管联邦学习具有广阔的前景，但它也面临着来自数据异构性的实际挑战（Li et al., 2020b）。这是因为现实世界中的用户数据通常是非独立同分布（non-iid）的，这本质上会导致局部最优解的偏离（Karimireddy et al., 2020）。此外，深度神经网络的排列不变性（permutation-invariant property）进一步增加了用户模型之间的异质性（Yurochkin et al., 2019; Wang et al., 2020b）。因此，大多数现有联邦学习方法所采用的局部模型逐元素平均的方式，可能无法诱导出一个理想的全局模型（Li et al., 2020c;b）。

为了应对用户异构性，已经进行了多种努力，主要从两个互补的角度出发：一种方法侧重于稳定局部训练，通过在参数空间内调节局部模型与全局模型之间的偏差来实现（Li et al., 2020b; Dinh et al., 2020; Karimireddy et al., 2020）。然而，这种方法可能无法充分利用用户模型之间潜在的知识，这些模型的多样性表明了它们局部数据的结构差异具有信息性，因此值得进一步研究。另一种方法旨在提高模型聚合的有效性（Yurochkin et al., 2019; Chen & Chao, 2021），其中知识蒸馏作为一种有效的解决方案而出现（Lin et al., 2020; Li & Wang, 2019）。在提供一个未标记的数据集作为代理的情况下，知识蒸馏通过利用来自局部模型的集成知识来丰富全局模型，从而缓解了由异构性引起的模型漂移问题，这已被证明比简单的参数平均更为有效。

然而，代理数据的前提条件使得这种方法在许多应用中并不可行，因为在这些应用中，服务器上可能并不总是有一个精心设计的数据集可用。此外，通过仅精炼全局模型，用户模型之间的固有异构性并未得到完全解决，这可能会反过来影响知识集成的质量，特别是如果它们由于局部数据的限制而产生偏差（Khoussainov et al., 2005），这在联邦学习中是一个典型的情况。

面对用户异构性带来的挑战以及现有技术的局限性，本工作提出了一种无需数据的知识蒸馏方法用于联邦学习，我们称之为FedGen（通过生成学习进行联邦蒸馏）。具体而言，FedGen学习一个仅从用户模型的预测规则中推导出的生成模型。给定一个目标标签，该生成模型能够产生与用户预测集合一致的特征表示。之后，这个生成器被广播给所有用户，通过在潜在空间上增加样本来辅助他们的模型训练，这些样本体现了来自其他对等用户的提炼知识。由于潜在空间的维度远小于输入空间，FedGen学习的生成器可以是轻量级的，从而为当前的联邦学习框架引入最小的开销。

提出的FedGen具有多重优势：i) 它提取了用户之间的知识，这些知识在模型平均后通常会被削弱，而且不依赖于任何外部数据。ii) 与某些仅精炼全局模型的前期工作不同，我们的方法直接使用提取的知识来调节本地模型的更新。我们展示了这种知识如何为本地模型施加归纳偏置，从而在非独立同分布（non-iid）数据分布下获得更好的泛化性能。iii) 此外，所提出的方法准备好应对更具挑战性的联邦学习场景，在这些场景中，由于隐私或通信限制，共享整个模型参数是不切实际的，因为所提出的方法仅需要本地模型的预测层来进行知识提取。

通过广泛的实证研究和理论阐述，我们的研究结果表明，与最先进的技术相比，我们提出的方法在更少的通信轮次下，能够获得具有更好泛化性能的全局模型。这证明了我们的方法在提高模型性能和减少通信开销方面的有效性。

2. Notations and Preliminaries

为了明确起见，本文讨论了一个典型的联邦学习（FL）设置，用于监督学习，即多类分类的一般问题。设$\mathcal{X}\subset {\mathbb{R}}^p$是一个实例空间，$\mathcal{Z}\subset {\mathbb{R}}^d$是一个潜在特征空间，其中d < p$，\mathrm{以}\mathrm{及}\mathcal{Y}\subset \mathbb{R}$是一个输出空间。$\mathcal{T}$表示一个域，它由实例空间$\mathcal{X}$上的数据分布$D$和一个真实标签函数$c^{\ast }:\mathcal{X}\to \mathcal{Y}$组成，即$\mathcal{T}:=⟨\mathcal{D},c^{\ast }⟩$。注意，我们将“域”和“任务”这两个术语视为等价。一个由参数$\theta :=[{\theta }^f;{\theta }^p]$参数化的模型由两个组件组成：一个由${\theta }^f$参数化的特征提取器$f:\mathcal{X}\to \mathcal{Z}$，和一个由${\theta }^p$参数化的预测器$h:\mathcal{Z}\to {\mathrm{\Delta }}^{\mathcal{Y}}$，其中${\mathrm{\Delta }}^{\mathcal{Y}}$是$\mathcal{Y}$上的单纯形（即概率分布）。给定一个非负、凸的损失函数$l:{\mathrm{\Delta }}^{\mathcal{Y}}\times \mathcal{Y}\to \mathbb{R}$，在域$\mathcal{T}$上参数化为$\theta$的模型的风险${\mathcal{L}}_{\tau }(\mathit{\theta })$定义为${\mathcal{L}}_{\tau }(\mathit{\theta }):={\mathbb{E}}_{x\sim \mathcal{D}}\left[l(h(f(x;{\mathit{\theta }}^f);{\mathit{\theta }}^p),c^{\ast }(x))\right]$，其中${\mathbb{E}}_{x\sim D}$表示对数据分布$D$的期望。监督学习的目标是找到一个模型$\theta$，该模型能够最小化风险${\mathcal{L}}_{\tau }(\mathit{\theta })$。

联邦学习（Federated Learning, FL） 旨在学习由$\theta$参数化的全局模型，使其在每个用户任务${\mathcal{T}}_k$上的风险最小化（McMahan等，2017）：

$$\underset{\theta }{min}{\mathbb{E}}_{{\mathcal{T}}_k\in \mathcal{T}}[{\mathcal{L}}_k(\theta )],(1)$$

其中，$\mathcal{T}=\{{\mathcal{T}}_k{\}}_{k=1}^K$是用户任务的集合。我们假设所有任务共享相同的标签规则$c^{\ast }$和损失函数$l$，即${\mathcal{T}}_k=⟨{\mathcal{D}}_k,c^{\ast }⟩$。在实践中，公式 (1) 被经验优化为：

$$\underset{\theta }{min}\frac{1}{K}\sum _{k=1}^K{\hat{\mathcal{L}}}_k(\theta ),$$

其中

$${\hat{\mathcal{L}}}_k(\mathit{\theta }):=\frac{1}{|{\hat{\mathcal{D}}}_k|}\sum _{x_i\in {\hat{\mathcal{D}}}_k}[l(h(f(x_i;{\mathit{\theta }}^f);{\mathit{\theta }}^p),c^{\ast }(x_i))]$$

是可观测数据集${\hat{\mathcal{D}}}_k$上的经验风险。FL 的一个隐含假设是全局数据$\hat{\mathcal{D}}$被分布到每个本地域中，且$\hat{\mathcal{D}}=\bigcup \{{\hat{\mathcal{D}}}_k{\}}_{k=1}^K$。

知识蒸馏（Knowledge Distillation, KD） 也被称为教师-学生范式，其目标是使用从一个或多个强大的教师模型中提炼的知识来学习一个轻量级的学生模型（Buciluǎ等，2006；Ba & Caruana，2014）。典型的 KD 利用代理数据集${\hat{\mathcal{D}}}_P$来最小化教师模型${\theta }_T$和学生模型${\theta }_S$之间的对数输出差异。一个典型的选择是使用 Kullback-Leibler 散度来度量这种差异（Hinton等，2015）：

$$\underset{{\theta }_S}{min}{\mathbb{E}}_{x\sim {\hat{\mathcal{D}}}_P}\left[D_{KL}(\sigma (g(f(x;{\theta }_f^T);{\theta }_p^T))\parallel \sigma (g(f(x;{\theta }_f^S);{\theta }_p^S)))\right],$$

其中$g(\cdot )$是预测器$h$的对数输出，$\sigma (\cdot )$是应用于这些对数的非线性激活函数，即$h(z;{\theta }_p)=\sigma (g(z;{\theta }_p))$。

知识蒸馏的理念已被扩展到 FL，以解决用户异质性问题（Lin等，2020；Chen & Chao，2021），通过将每个用户模型${\theta }_k$视为教师，其信息被聚合到学生（全局）模型$\theta$中，以提高其泛化性能：

$$\underset{\theta }{min}{\mathbb{E}}_{x\sim {\hat{\mathcal{D}}}_P}\left[D_{KL}(\sigma (\frac{1}{K}\sum _{k=1}^{K}g(f(x;{\theta }_f^k);{\theta }_p^k))\parallel \sigma (g(f(x;{\theta }_f);{\theta }_p)))\right].$$

上述方法的主要限制在于其对代理数据集${\hat{\mathcal{D}}}_P$的依赖，选择代理数据集需要谨慎考虑，并在蒸馏性能中起关键作用（Lin等，2020）。接下来，我们展示如何在无数据的情况下使 KD 在 FL 中可行。

Figure 1. FEDGEN 概述：服务器学习生成器$G_w(\cdot |y)$，用于聚合来自不同本地客户端的信息，而无需观察它们的数据。然后将生成器发送给本地用户，其知识被蒸馏到用户模型中，以调整其对良好特征分布的解释。

算法 1 FEDGEN 的描述如下：

1. 初始化：算法要求输入任务集$\{T_k{\}}_{k=1}^K$，全局参数$\theta$和局部参数$\{{\theta }_k{\}}_{k=1}^K$，生成器参数$w$，初始化全局标签分布的估计$\hat{p}(y)$，用于生成与标签相关的数据。初始学习率$\alpha$、$\beta$，局部迭代步数$T$，批量大小$B$，以及局部标签计数器$c_k$。
2. 训练循环：反复进行以下步骤：
   • 服务器随机选择活跃用户$A$，并将参数$w$、$\theta$和$\hat{p}(y)$广播给这些用户。
   • 对于每个用户$k\in A$，执行以下操作：
     • 将全局参数$\theta$赋值给局部参数${\theta }_k$。
     • 在$T$次迭代中，每次从任务$T_k$中采样数据$\{x_i,y_i{\}}_{i=1}^B$和生成器输出$\{{\hat{z}}_i\sim G_w(\cdot |{\hat{y}}_i),{\hat{y}}_i\sim \hat{p}(y){\}}_{i=1}^B$。
     • 更新标签计数器$c_k$。
     • 优化局部参数${\theta }_k$，使用梯度下降法更新${\theta }_k$。
   • 用户将更新后的${\theta }_k$和$c_k$发送回服务器。
3. 参数更新：服务器根据接收到的参数更新全局参数$\theta \leftarrow \frac{1}{|A|}\sum _{k\in A}{\theta }_k$，并根据$\{c_k{\}}_{k\in A}$更新$\hat{p}(y)$，最后优化生成器参数$w\leftarrow w-\alpha {\mathrm{\nabla }}_{w}J(w)$，直到训练停止。

这个算法的核心在于通过协调多用户的局部训练，达到有效的联邦学习效果。

3. FEDGEN：基于GAN的无数据联合蒸馏

本节中，我们详细阐述了我们提出的方法，概述如算法1所示，并在图1中展示了其学习过程的概述。

3.1. 知识提取

我们的核心思想是提取关于全局数据分布视角的知识，这在传统的联合学习中是不可观测的，并将这种知识蒸馏到本地模型中以指导它们的学习。我们首先考虑学习一个条件分布${\mathcal{Q}}^{\ast }:\mathcal{Y}\to \mathcal{X}$，用于描述与目标标签的真实先验分布$p(y)$和后验分布$p(x|y)$一致的知识，其中这两者都是未知的。

$$Q^{\ast }=\underset{Q:\mathcal{Y}\to \mathcal{X}}{\arg max}{\mathbb{E}}_{y\sim p(y)}{\mathbb{E}}_{x\sim Q(x|y)}[\log p(y|x)],(2)$$

为了使方程2相对于$\mathcal{Q}$可优化，我们用它们的经验近似值替换。首先，我们估计$p(y)$为：

$$\hat{p}(y)\propto \sum _k{\mathbb{E}}_{x\sim {\hat{D}}_k}[I(c^{\ast }(x)=y)]$$

其中，$I(\cdot )$是指示函数，${\hat{D}}_k$是领域${\mathcal{T}}_k$的可观察数据。在实践中，$\hat{p}(y)$可以通过在模型上传阶段要求用户的训练标签计数来获得。接下来，我们使用用户模型的集成智慧来近似$p(y|x)$：

$$\log \hat{p}(y|x)\propto \frac{1}{K}\sum _{k=1}^K\log p(y|x;{\theta }_k)$$

具备了以上的近似值，直接在输入空间$\mathcal{X}$上优化方程2仍可能具有挑战性：当$\mathcal{X}$的维度较高时，会带来计算负担，并可能泄露有关用户数据配置文件的信息。因此，更可行的方法是恢复一个诱导分布$G^{\ast }:\mathcal{Y}\to \mathcal{Z}$，覆盖一个比原始数据空间更紧凑的潜在空间，并能减轻某些与隐私相关的担忧：

$$G^{\ast }=\arg \underset{G:\mathcal{Y}\to \mathcal{Z}}{max}{\mathbb{E}}_{y\sim \hat{p}(y)}{\mathbb{E}}_{z\sim G(z|y)}[\sum _{k=1}^K\log p(y|z;{\mathit{\theta }}_k^p)].(3)$$

基于以上推理，我们的知识提取目标是通过学习一个条件生成器$G$，参数化为$w$，来优化以下目标：

$$\underset{\mathit{w}}{min}J(\mathit{w}):={\mathbb{E}}_{y\sim \hat{p}(y)}{\mathbb{E}}_{z\sim G_{\mathit{w}}(z|y)}[l(\sigma (\frac{1}{K}\sum _{k=1}^{K}g(z;{\mathit{\theta }}_k^p)),y)],(4)$$

以下是对上述段落的翻译：

其中，$g$和$\sigma$分别是定义在第2节中的logit输出和激活函数。给定任意的样本$y$，优化方程4只需要访问用户模型的预测模块${\theta }_p^k$。具体来说，为了使$G(\cdot |y)$产生多样化的输出，我们引入一个噪声向量$ϵ\sim \mathcal{N}(0,I)$给生成器，这类似于先前文献中提出的重新参数化技术（Kingma & Welling, 2014），使得$z\sim G_w(\cdot |y)\equiv G_w(y,ϵ|ϵ\sim \mathcal{N}(0,I))$。我们在补充材料中讨论更多实现细节。

根据任意目标标签$y$，所提出的生成器可以产生特征表示$z\sim G_w(\cdot |y)$​，这些表示能够从用户模型的集成中诱导出理想的预测。换句话说，生成器近似了一个共识分布的诱导形象，与全局视角下的用户数据保持一致。

3.2. 知识蒸馏

学习到的生成器$G_w$会被广播到本地用户，使得每个用户模型可以从$G_w$中采样，获得特征空间上的增强表示$z\sim G_w(\cdot |y)$。因此，本地模型${\theta }_k$的目标变为最大化其对增强样本产生理想预测的概率：

$$\underset{{\theta }_k}{min}J({\mathit{\theta }}_k):={\hat{\mathcal{L}}}_k({\mathit{\theta }}_k)+{\hat{\mathbb{E}}}_{y\sim \hat{p}(y),z\sim G_{\mathit{w}}(z|y)}[l(h(z;{\mathit{\theta }}_k^p);y)],(5)$$

其中

$${\hat{\mathcal{L}}}_k({\mathit{\theta }}_k):=\frac{1}{|{\hat{\mathcal{D}}}_k|}\sum _{x_i\in {\hat{\mathcal{D}}}_k}[l(h(f(x_i;{\mathit{\theta }}_k^f);{\mathit{\theta }}_k^p),c^{\ast }(x_i))]$$

是给定本地数据${\hat{D}}_k$​ 的经验风险。稍后我们将展示，增强样本可以为本地用户引入归纳偏差，增强他们的模型学习，进而提升泛化性能。至此，我们提出的方法已经实现了无数据的知识蒸馏，通过交互地学习一个主要依赖于本地模型预测规则的轻量级生成器，并利用该生成器将共识知识传递给本地用户。我们在第6.2节中证明了我们的方法可以有效应对联邦学习中的用户异质性，并且在第4节中分析了其理论优势。

3.3. 灵活参数共享的扩展

除了应对数据异质性外，FEDGEN还可以处理一种具有挑战性的联邦学习（FL）场景，即共享整个模型违反了通信或隐私的前提条件。一方面，具有深度特征提取层的高级网络通常包含数百万个参数（He et al., 2016; Brown et al., 2020），这会给通信带来显著负担。另一方面，已有研究表明，常规FL方法存在后门攻击的可能性（Wang et al., 2020a）。对于医疗或金融等实际的FL应用，共享整个模型参数可能会带来相当大的隐私风险，如先前的研究所讨论的（He et al., 2020）。

FEDGEN可以缓解这些问题，只共享本地模型的预测层${\theta }_p^k$，这是优化方程4所需的主要信息，同时将特征提取器${\theta }_f^k$保持本地化。与共享整个模型的策略相比，这种部分共享范式更高效，同时数据泄露的风险也更低。第6.4节的实证研究表明，即使不共享特征提取模块，FEDGEN也显著有利于本地用户。我们将在补充材料中提供这种变体方法的算法摘要。

4. FEDGEN分析

在本节中，我们从多个角度来理解我们提出的方法。首先，我们可视化FEDGEN所学习和蒸馏的知识，然后分别从分布匹配和领域适应的角度分析为什么蒸馏的知识是有利的。我们主要集中在解释FEDGEN背后的原理，详细的讨论和推导留待补充材料中。

Figure 2 . 在应用知识蒸馏（KD）后，一个用户的准确率从81.2%提高到98.4%（图2a - 图2b），而通过参数平均（不使用KD）获得的全局模型准确率为93.2%（图2c），相比之下，一个理想模型（oracle model）的准确率为98.6%（图2d）。

Figure 3. 生成器样本逐渐接近真实数据分布，其中每个用户模型（教师）仅看到有限且不相交的本地数据。背景颜色表示基于全局数据学习到的理想决策边界

4.1 用于归纳偏差的知识蒸馏

我们在一个包含三个用户的FL原型上展示了FEDGEN的知识蒸馏过程，每个用户分配到一个不相交的数据集$\widehat{{\mathcal{D}}_k}$，$k\in \{1,2,3\}$。当仅使用本地数据进行训练时，用户模型容易学习到有偏的决策边界（见图2a）。

接下来，根据用户模型的预测规则，学习生成器$G_w(·|y)$。为了清晰可视化，我们在原始特征空间$\mathcal{Y}\to \mathcal{X}\subset {\mathbb{R}}^2$上学习$G_w(·|y)$而不是在潜在空间上。如图3所示，$r(x|y)$表示从$G_w(x|y)$导出的分布，即使个体本地模型有偏差，它也逐渐与真实分布$p(x|y)$一致（图2d）。换句话说，$G_w(x|y)$可以融合用户模型的聚合信息来近似全局数据分布。

然后，让用户从$G_w(x|y)$中采样，这作为具有有限数据的用户的归纳偏差。结果是，每个用户可以超越其自身的训练数据，调整其决策边界以接近集成智慧（图2b）。

4.2 用于分布匹配的知识蒸馏

FEDGEN与先前工作的显著区别在于，知识是蒸馏到用户模型而不是全局模型。因此，蒸馏的知识传递归纳偏差给用户，可以通过在潜在空间$\mathcal{Z}$上执行分布匹配来直接调节他们的学习：

Remark 1. 令$p(y)$为标签的先验分布，$r(z|y):\mathcal{Y}\to \mathcal{Z}$为从生成器$G_w$导出的条件分布。然后使用从$r(z|y)$采样来调节用户模型${\theta }_k$，可以最小化从生成器和用户分别导出的两个分布之间的条件KL散度：

$$\begin{array}{rl}& \underset{{\theta }_k}{max}{\mathbb{E}}_{y\sim p(y),z\sim r(z|y)}[\log p(y|z;{\mathit{\theta }}_k)]\\ & \equiv \underset{{\theta }_k}{min}{D}_{\mathrm{KL}}[r(z|y)\parallel p(z|y;{\mathit{\theta }}_k)],& \text{(6)}\end{array}$$

其中我们定义$p(z|y;{\theta }_k)$为在输入特征给预测器${\theta }_k$为$z$且给出标签$y$的情况下的概率。

实际上，通过使用生成器的经验样本优化方程6：$\{(z,y)|y\sim \hat{p}(y),z\sim G_w(z|y)\}$，这与本地模型目标（方程5）的第二项一致，对于所有$y\in \mathcal{Y}$：

$$\underset{{\mathit{\theta }}_k}{max}{\mathbb{E}}_{z\sim r(z|y)}[\log p(y|z;{\mathit{\theta }}_k)]\approx \underset{{\mathit{\theta }}_k}{min}{\mathbb{E}}_{z\sim G_{\mathit{w}}(z|y)}[l(h(z;{\mathit{\theta }}_k^p);y)].$$

与先前的应用权重正则化到本地模型的方法不同（Li et al., 2020b; Dinh et al., 2020），FEDGEN可以作为一种替代和兼容的解决方案来应对用户异质性，本质上弥合了用户模型在对理想特征分布解释上的差距。

4.3. 用于改进泛化的知识蒸馏

我们还可以从理论上将 FEDGEN 学习到的知识与改进的泛化界限联系起来。为此，我们首先提出一个基于领域适应（Ben-David 等，2007；2010）的 FL 中聚合模型的性能界限：

定理 1：（FL 的泛化界限）考虑一个具有$K$个用户的 FL 系统。设${\mathcal{T}}_k=⟨{\mathcal{D}}_k,c^{\ast }⟩$和$\mathcal{T}=⟨\mathcal{D},c^{\ast }⟩$分别是第$k$个本地领域和全局领域。设$R:\mathcal{X}\to \mathcal{Z}$是一个在用户之间同时共享的特征提取函数。记$\{h_k\}$为在领域${\mathcal{T}}_k$上学习的假设，$h=\frac{1}{K}\sum _{k=1}^K{h}_k$为用户假设的全局集成。那么，以至少$1-\delta$的概率：

$$\begin{array}{rl}{\mathcal{L}}_{\mathcal{T}}(h)& \equiv {\mathcal{L}}_{\mathcal{T}}\left(\frac{1}{K}\sum _k{h}_k\right)\\ & \le \frac{1}{K}\sum _k{\hat{\mathcal{L}}}_{{\mathcal{T}}_k}(h_k)+\frac{1}{K}\sum _k(d_{\mathcal{H}\mathrm{△}\mathcal{H}}({\tilde{\mathcal{D}}}_k,\tilde{\mathcal{D}})+{\lambda }_k)\\ & +\sqrt{\frac{4}{m}(d\log \frac{2em}{d}+\log \frac{4K}{\delta })},\end{array}$$

其中${\hat{\mathcal{L}}}_{{\mathcal{T}}_k}(h_k)$是${\mathcal{T}}_k$上的经验风险，${\lambda }_k:=\underset{h}{min}({\mathcal{L}}_{{\mathcal{T}}_k}(h)+{\mathcal{L}}_{\mathcal{T}}(h))$表示一个 oracle 性能。$d_{\mathcal{H}\mathrm{\Delta }\mathcal{H}}({\tilde{\mathcal{D}}}_k,\tilde{\mathcal{D}})$表示在对称差假设空间上测量的发散性。${\tilde{\mathcal{D}}}_k$和$\tilde{\mathcal{D}}$分别是${\mathcal{D}}_k$和$\mathcal{D}$在$R$上的诱导图像， 使得对于一个概率事件$B$，有${\mathbb{E}}_{z\sim \tilde{\mathcal{D}}k}[B(z)]=\mathbb{E}x\sim {\mathcal{D}}_k[B(R(x))]$,对于$\tilde{\mathcal{D}}$也如此。

具体而言，${\mathcal{L}}_{\mathcal{T}}(h)$通常被认为是 FL 中全局模型的理想上界（Peng 等，2019；Lin 等，2020）。从定理 1 中可以得出两个关键含义：i）较大的用户异质性导致较高的分布发散性（$d_{\mathcal{H}\mathrm{\Delta }\mathcal{H}}({\tilde{\mathcal{D}}}_k,\tilde{\mathcal{D}})$），这会削弱全局模型的质量；ii）更多的经验样本（$m$）有利于泛化性能，这会缓解数值约束。

换句话说，通过用与全局分布对齐的增强数据丰富本地用户可以提高泛化性能：

推论 1：设$\mathcal{T},{\mathcal{T}}_k,R$如定理 1 中定义。${\mathcal{D}}_A$表示一个增强分布，${\mathcal{D}}_k^{\prime }=\frac{1}{2}({\mathcal{D}}_k+{\mathcal{D}}_A)$是分布的混合体。相应地，${\tilde{D}}_A,{\tilde{D}}_k^{\prime }$分别表示${\mathcal{D}}_A$和${\mathcal{D}}_k^{\prime }$在$\mathcal{R}$上的诱导图像。设${\hat{\mathcal{D}}}_k^{\prime }={\hat{\mathcal{D}}}_k\cup {\hat{\mathcal{D}}}_A$是${\mathcal{D}}_k^{\prime }$的经验数据集，其中$|{\hat{\mathcal{D}}}_k|=m,|{\hat{\mathcal{D}}}_k^{\prime }|=m^{\prime }>m$。如果$d_{\mathcal{H}\mathrm{\Delta }\mathcal{H}}({\tilde{\mathcal{D}}}_A,\tilde{\mathcal{D}})$是有界的，使得存在$ϵ>0$，满足$d_{\mathcal{H}\mathrm{\Delta }\mathcal{H}}({\tilde{\mathcal{D}}}_A,\tilde{\mathcal{D}})\le ϵ$，则以概率 1 -$\delta$：

$$\begin{array}{rl}{\mathcal{L}}_{\mathcal{T}}(h)& \le \frac{1}{K}\sum _k{\mathcal{L}}_{{\mathcal{T}}_k^{\prime }}(h_k)+\frac{1}{K}\sum _k(d_{\mathcal{H}\mathrm{△}\mathcal{H}}({\tilde{\mathcal{D}}}_k^{\prime },\tilde{\mathcal{D}})+{\lambda }_k^{\prime })\\ & +\sqrt{\frac{4}{m^{\prime }}(d\log \frac{2e{m}^{\prime }}{d}+\log \frac{4K}{\delta })},& \text{(7)}\end{array}$$

其中${\mathcal{T}}_k^{\prime }=\{{\mathcal{D}}_k^{\prime },c^{\ast }\}$是更新后的本地领域， 当$ϵ$较小时，有$d_{\mathcal{H}\mathrm{\Delta }\mathcal{H}}({\tilde{\mathcal{D}}}_k^{\prime },\tilde{\mathcal{D}})\le d_{\mathcal{H}\mathrm{\Delta }\mathcal{H}}({\tilde{\mathcal{D}}}_k,\tilde{\mathcal{D}})$​，并且

$$\sqrt{\frac{4}{m^{\prime }}(d\log \frac{2e{m}^{\prime }}{d}+\log \frac{4K}{\delta })}<\sqrt{\frac{4}{m}(d\log \frac{2em}{d}+\log \frac{4}{\delta K})}。$$

这样的增强分布${\mathcal{D}}_A$可以从多个方面促进 FL：不仅通过更多的经验样本（$m^{\prime }>m$）放宽数值约束，而且还减少了本地和全局特征分布之间的差异（$d_{\mathcal{H}\mathrm{\Delta }\mathcal{H}}({\tilde{\mathcal{D}}}_k^{\prime },\tilde{\mathcal{D}})$）。这一发现与 FEDGEN 的优点一致：由于生成器$G_w(z|y)$是为了恢复特征空间上的聚合分布而学习的，可以将从生成器获得的样本$\{z|y\sim \hat{p}(y),z\sim G_w(z|y)\}$视为来自${\tilde{\mathcal{D}}}_A$的增强数据，这自然与全局诱导分布$\tilde{\mathcal{D}}$有小的偏差。沿着这条线的更严格的分析留待我们未来的工作中进行。我们在附录中详细说明了这种增强分布${\mathcal{D}}_A$的角色。

5. Related Work

联邦学习（Federated Learning, FL） 最初由McMahan等人（2017）提出，作为一种去中心化的机器学习范式。随后，沿着这一方向的研究工作解决了联邦学习面临的不同挑战，包括异质性（Karimireddy等人，2020；Li等人，2020b；Mansour等人，2020）、隐私（Duchi等人，2014；Agarwal等人，2018）、通信效率（Guha等人，2019；Konečný等人，2016）以及收敛性分析（Kairouz等人，2019；Qu等人，2020；Yuan & Li, 2019）。特别地，大量工作已经被提出以处理用户异质性，包括通过正则化模型权重更新（Li等人，2020b）、允许个性化的用户模型（Fallah等人，2020；Dinh等人，2020）或引入新的模型聚合方案（Yurochkin等人，2019；Mansour等人，2020）。我们推荐读者参考Li等人（2020a）的工作，以获取对联邦学习领域最新进展的有组织讨论。

知识蒸馏（Knowledge Distillation, KD） 是一种技术，用于将一个或多个教师模型的知识压缩到一个空白的学生模型中（Hinton et al., 2015; Buciluǎ et al., 2006; Ba & Caruana, 2014; Jacobs et al., 1991）。传统的知识蒸馏依赖于一个代理数据集（Hinton et al., 2015）。最近的工作使得知识蒸馏在更少数据参与的情况下成为可能，例如数据集蒸馏（Wang et al., 2018）或核心数据选择（Tsang et al., 2005; Sener & Savarese, 2018）。随后，出现了无数据知识蒸馏方法，其目标是重建用于训练教师的样本（Yoo et al., 2019; Micaelli & Storkey, 2019）。特别是，Lopes et al.（2017）从教师的激活层中提取元数据。Yoo et al.（2019）学习了一个条件生成器，该生成器能够产生样本，这些样本能够最大化教师模型对于目标标签的预测概率。沿着相同的思路，Micaelli & Storkey（2019）通过对抗训练来学习一个生成器。与先前的工作不同，我们学习了一个为联邦学习（FL）量身定制的生成模型，该模型通过在潜在空间中对多个用户模型的知识进行集成，使得学习和通信过程更加轻量级。

联邦学习中的知识蒸馏 最近已成为解决用户异质性的有效方法。大多数现有工作是数据依赖的（Lin et al., 2020; Sun & Lyu, 2020; Guha et al., 2019; Chen & Chao, 2021）。特别是，Lin et al.（2020）提出了FEDDFUSION，该方法通过知识蒸馏来精炼全局模型，假设存在一个来自相同或相似领域的未标记数据集。为应对数据异质性，已进行了补充性的知识蒸馏工作（Li & Wang, 2019; Sattler, 2021）。具体而言，Li & Wang（2019）传输代理数据集而不是模型参数。FEDAUX（Sattler, 2021）通过利用辅助数据集来初始化服务器模型并对用户模型进行加权集成，从而进行数据依赖的蒸馏，而FEDGEN则以无数据的方式进行知识蒸馏。FEDMIX（Yoon, 2021）是一个数据增强的联邦学习框架，其中用户将他们的批量平均数据共享给其他用户以辅助本地训练。相比之下，FEDGEN从现有的用户模型参数中提取知识，面临的隐私风险较小。FEDDISTILL（Federated Distillation）由Seo et al.（2020）提出，该方法从用户模型中提取逻辑向量输出的统计数据，并将这些元数据共享给用户进行知识蒸馏。我们在第6节中将与这一领域的工作进行详细比较。

6. Experiments

在本节中，我们将比较我们提出的方法与其他关键相关工作的性能。具体的实现细节和扩展的实验结果将放在补充材料中。

6.1 实验设置

基线方法：除了FEDAVG（McMahan et al., 2017）之外，FEDPROX通过在模型目标中加入一个近端项来正则化本地模型训练（Li et al., 2020b）。FEDENSEMBLE将FEDAVG扩展到对用户模型的预测输出进行集成。FEDDFUSION是一种基于数据的知识蒸馏方法（Lin et al., 2020），我们为其提供了未标记的训练样本作为代理数据集。FEDDISTILL（Jeong et al., 2018）是一种无数据的知识蒸馏方法，它在用户之间共享按标签平均的逻辑向量。由于它不共享网络参数，因此与其他基线方法相比，性能下降不可忽视。为了公平比较，我们从FEDDISTILL中派生出一个基线，该基线同时共享模型参数和按标签的逻辑向量。我们将这个更强的基线称为FEDDISTILL+。

数据集：我们在三个图像数据集上进行实验：MNIST（LeCun & Cortes, 2010）、EMNIST（Cohen et al., 2017）和CELEBA（Liu et al., 2015），这些数据集是根据LEAF联邦学习基准（Caldas et al., 2018）的建议选择的。其中，MNIST和EMNIST数据集用于数字和字符图像分类，而CELEBA是一个名人脸部数据集，用于学习一个二分类任务，即预测图片中的名人是否在笑。

配置：除非另有说明，我们运行200轮全局通信，总共有20个用户模型，并且活跃用户比例r=50%。我们采用本地更新步数T=20，每步使用大小为B=32的小批量。我们使用最多50%的总训练数据集并将其分配给用户模型，并使用所有测试数据集进行性能评估。对于分类器，我们遵循（McMahan et al., 2017）的网络架构，并将最后一个多层感知机（MLP）层视为预测器${\theta }_k^p$，将所有前面的层视为特征提取器${\theta }_k^f$。生成器$G_w$是基于MLP的。它接受一个噪声向量ε和一个独热标签向量$y$作为输入，经过一个维度为$d_h$的隐藏层后，输出一个维度为$d$的特征表示。为了进一步增加生成器输出的多样性，我们还利用了先前工作中多样性损失（Mao et al., 2019）的思想来训练生成器模型。

用户异质性：对于MNIST和EMNIST数据集，我们遵循先前的工作（Lin et al., 2020; Hsu et al., 2019）使用Dirichlet分布$Dir(\alpha )$来模拟非独立同分布（non-iid）的数据分布，其中较小的α表示更高的数据异质性，因为它使得用户$k$的$p_k(y)$​分布更加偏斜。我们在图4中可视化了在MNIST数据集上采用不同α值对统计异质性的影响。对于CELEBA，原始数据本身就是非独立同分布的。我们进一步通过将属于不同名人的图片聚合到不相交的组中，并将每个组分配给一个用户，来增加数据的异质性。

Figure 4. MNIST数据集上用户间统计异质性的可视化，其中x轴表示用户ID，y轴表示类别标签，散点的大小表示该用户可用于训练的该标签样本数量。

6.2. 性能概述

从表1可以看出，FEDGEN在性能上显著优于其他基线方法。

Table 1. 不同数据设置下的性能概述。对于MNIST和EMNIST数据集，较小的α表示更高的异质性。对于CELEBA数据集，r表示活跃用户与总用户的比率。T表示本地训练步骤（通信延迟）。

数据异质性的影响FEDGEN是唯一一个能够在不同程度用户异质性下保持稳定且表现优异的算法。如图5所示，当数据分布高度异质（$\alpha$较小）时，FEDGEN的优势更为显著。这一结果验证了我们的动机，即FEDGEN的优势来自于传递给本地用户的知识，有助于减轻用户间潜在分布的差异。这种知识对于FEDAVG或FEDPROX等基线方法通常是不可访问的。

作为竞争力较强的基线方法之一，FEDDFUSION的优势随着数据异质性的减少逐渐消失，逐渐与FEDAVG相当，如图5a和图5c所示。与FEDDFUSION不同，我们方法的性能提升在大多数情况下都显著，超过了FEDDFUSION。这种差异表明，我们提出的直接将知识传递给用户模型的方法，在引导本地模型学习方面可能比使用代理数据对全局模型进行微调更为有效，尤其是当传递的知识包含归纳偏差以指导本地模型学习时。

作为一种无数据知识蒸馏基线，FEDDISTILL的性能下降明显，这表明在联邦学习中参数共享的重要性。另一方面，FEDDISTILL+对数据异质性较为敏感。如表1所示，当数据分布接近iid（例如$\alpha \ge 1$）时，由于共享的logit统计作为传递的知识，它可以胜过FEDAVG，但当$\alpha$变小时，其表现不如FEDAVG，这表明仅共享这些元数据可能不足以有效应对用户异质性。

FEDENSEMBLE从所有用户模型的集成预测中获益，尽管其收益相对于FEDGEN较小。我们将我们方法的领先性能归因于本地模型的更好泛化性能。在传递的知识指导下，FEDGEN中的用户模型可以快速摆脱局部最优，其聚合效果可能优于FEDENSEMBLE中潜在有偏模型的集成。

Figure 5. 数据异质性相关的性能可视化。

学习效率如图6所示，FEDGEN具有最快的学习曲线，可以快速达到高性能并优于其他基线方法。尽管FEDDFUSION在某些数据设置下享有较高的学习效率，由于从主数据中获得的优势，我们的方法可以直接为每个本地用户提供积极学习的知识，其效果更为明显和一致（更多细节见附录）。

Figure 6. 选择的学习曲线，平均了3个随机种子

关于共享生成模型的评论在紧凑的潜在空间中，生成模型可以轻量化学习或下载。在实践中，我们使用一个具有2个紧凑MLP层的生成器网络，其参数尺寸与用户分类模型相比较小。上述实证结果还表明，通过更快的收敛速率获得的领先性能增益可以抵消通过共享生成模型带来的通信负载。

6.3. 敏感性分析

Straggler用户的影响我们在CELEBA数据集上探索不同总用户数与活跃用户数的影响，活跃比例$r$从0.2到0.9变化。图5b显示，在存在大量Straggler用户（$r=0.2$，每轮学习中有25个活跃用户中的5个）时，我们的方法仅次于FEDENSEMBLE，并且在具有适度活跃用户数量时始终优于所有基线方法的渐近性能。结合图6a和图6b，可以观察到我们的方法在达到高性能时需要更少的通信轮次，无论Straggler用户的设置如何。

不同网络架构的影响我们在MNIST数据集上进行了CNN和MLP网络架构的分析。如图5d和图5c所示，尽管使用CNN网络的整体性能明显优于MLP网络，FEDGEN的卓越性能在两种不同的网络设置下都是一致的。

通信频率的影响我们在EMNIST上探索了不同的本地更新步骤$T$，较高的$T$意味着全局通信之前的较长通信延迟。表1中的结果表明，我们的方法对不同水平的通信延迟具有稳健性（更多结果见附录）。

生成器网络架构和采样大小的影响扩展分析验证了FEDGEN在不同生成器网络架构下的稳健性（表2）。此外，从生成器中抽取合成数据仅为本地用户增加了较小的训练工作量（表2）。在不同合成样本大小的情况下，FEDGEN相对于FEDAVG的增益始终显著，而足够数量的合成样本带来了更好的性能（图7）。特别是在表2中，我们探讨了输入噪声（$d_{ϵ}$）和生成器的隐藏层（$dh$）的不同维度，同时保持其输出层维度不变（即特征空间$Z$的维度）。表3显示了一个本地更新的训练时间，对用户和通信轮次进行了平均化。$BG$表示用于每个小批量优化的合成样本数。默认情况下，我们设置$BG=B$，其中$B$​是从本地数据集中抽取的实际样本数（参见算法1）。

Figure 7. 合成样本的影响。

Table 3. 使用EMNIST数据集 (α = 0.1) 时合成样本数量的影响。

6.4. 扩展到灵活参数共享

为了减轻隐私和通信问题，FEDGEN已经准备好在不共享整个模型参数的情况下受益于分布式学习。为了探索这一潜力，我们对FEDAVG、FEDDISTILL+和FEDGEN进行了案例研究，其中用户模型仅共享最后的预测层，并保留其特征提取层本地化。需要注意的是，FEDDFUSION并非设计用于部分参数共享的FL，这需要整个用户模型进行KD。为了公平比较，我们修改了FEDDFUSION，使其在模型聚合阶段上传整个用户模型，但禁止下载特征提取器，以便服务器模型仍然可以使用代理数据进行微调。

表4中的结果显示，我们的方法始终以显著的优势超越其他基线方法，尤其在高数据异质性的情况下（图8）。与FEDDFUSION的显著性能差异验证了在这种具有挑战性的场景下无数据蒸馏的有效性。这些有前途的结果表明，FEDGEN有潜力进一步减少通信负载，不仅通过快速收敛，还通过灵活的参数共享策略。

表4. 仅共享最后预测层时在MNIST上的性能概览

Figure 8. 在MNIST数据集上有限参数共享下的学习曲线。

7. Conclusions

本文提出了一种联邦学习范式，通过有效的知识蒸馏来解决用户异质性问题，而无需使用任何外部数据。基于理论推论指导的大量实证实验表明，我们提出的方法能够在与最新技术相比的通信轮次更少的情况下，显著改善联邦学习的泛化性能。

Acknowledgments

本研究得到了国家科学基金会（NSF）IIS-1749940号、海军研究办公室（Office of Naval Research）N00014-20-1-2382号以及国家老龄化研究所（National Institute on Aging）RF1AG072449号的联合支持。