FedeKD：在非同质化设置下用于鲁棒联邦知识蒸馏的基于能量的门控机制

# FedeKD：面向异构环境下鲁棒联邦知识蒸馏的基于能量的门控机制

**来源：** https://arxiv.org/html/2605.05553

**Quang-Huy Nguyen**
计算机科学和软件工程系
奥本大学
阿拉巴马州奥本 36849
[email protected]

&

**Jiaqi Wang**
计算机科学和软件工程系
奥本大学
阿拉巴马州奥本 36849
[email protected]

&

**Wei-Shinn Ku***
计算机科学和软件工程系
奥本大学
阿拉巴马州奥本 36849
[email protected]

###### 摘要

联邦学习（FL）运行在异构环境中，数据分布的变化和不对称的模型设计往往导致负迁移。虽然联邦知识蒸馏（FKD）避免了直接共享模型参数，但现有方法通常依赖于公共数据集或假设转移的知识具有均匀可靠性，这限制了其在实践中的鲁棒性。本文提出了 **FedeKD**，这是一种可靠性感知的 FKD 框架，它将样本级信任估计作为知识转移的显式组件，而不依赖额外的公共数据。每个客户端维护一个用于本地学习的高容量私有模型和一个用于跨客户端知识交换的轻量级共享代理模型。在训练过程中，服务器聚合代理模型以形成全局代理模型，随后用于指导私有模型的更新。FedeKD 的核心是一种基于能量的门控机制，该机制将特定任务的私有-代理分歧转化为反向蒸馏的样本级信任权重。这种机制实现了知识转移的样本级加权，其中代理模型对可靠样本贡献更大，同时降低不可靠样本的权重。在六个真实世界数据集上的大量实验表明，FedeKD 在保持强大预测性能的同时，显著减少了异构设置下的负迁移。

## 1 引言

联邦学习（FL）使多个客户端能够在不共享原始数据的情况下协作训练模型，使其非常适合对隐私敏感的应用程序。然而，现实世界的部署涉及数据分布和模型不对称性的显著异构性，使得可靠的知识共享变得具有挑战性。经典的参数聚合方法，包括 FedAvg (McMahan et al., 2017) 和 FedProx (Li et al., 2020)，在这种异构性下往往会出现性能下降和收敛不稳定。

联邦知识蒸馏（FKD）通过传输模型输出而非参数提供了一种替代方案。然而，在异构设置中，传输知识的质量在不同样本和客户端之间可能存在显著差异。这要求将知识可靠性视为 FKD 中的一等对象，而不是假设所有传输信号应对本地学习产生同等影响。由于分布偏移、有限的本地数据或架构差异，教师模型可能会产生误导性信号。盲目地传输此类知识可能导致严重的负迁移（即，与本地训练相比性能下降）。

现有的 FKD 方法通常假设传输的知识具有均匀可靠性，忽视了跨样本和客户端的知识质量变化。这一挑战反映了人类学习的过程，当教师的专长与学习者的背景不匹配，或者指导提供的可操作信息很少时，指导变得不可靠。

在这项工作中，我们试图回答一个根本性问题：**何时应该信任教师模型来指导知识转移？**

我们的关键灵感来自一个简单的现实原则：教师不会平等地对待所有知识，而是强调他们最了解且最自信的部分。基于这一见解，我们引入了一个带有**能量门控知识蒸馏**（FedeKD）的联邦学习框架，以动态确定每个代理预测应在多大程度上影响每个私有模型更新（图 1）。FedeKD 实现了跨异构客户端的鲁棒知识转移，并通过自适应的样本级加权减轻负迁移。

**图 1：FedeKD 框架。** 私有模型（Private Model）表示用于本地学习的高容量网络，而代理模型（Proxy Model）是用于客户端间通信和聚合的轻量级网络。详见附录 H。

FedeKD 分两个阶段运行。在**前向**阶段，每个客户端将其私有模型中的知识蒸馏到代理模型中，该代理模型在服务器上聚合以形成全局代理。在**后向**阶段，全局代理通过基于能量的门控机制指导私有模型的更新。此过程实现了样本级信任加权的知识转移，其中模型在传输知识看起来可靠时更依赖代理，而在看起来不可靠时较少依赖。

综上所述，本工作的贡献有两方面：

- 我们引入了一种可靠性感知的后向蒸馏目标，将焦点从如何交换知识转移到应信任每个传输信号的程度。为了支持此目标且无需公共数据，我们设计了具有非对称私有-代理架构的 FedeKD。该目标通过批归一化能量门实现，将特定任务的私有-代理分歧映射为连续的样本级信任权重。对于分类任务，能量使用熵校准的分布分歧；对于回归任务，能量使用连续预测分歧。该机制降低了不可靠知识的权重，同时保留了信息性信号，从而减轻了异构 FL 环境中的负迁移。
- 在六个真实世界数据集上的大量实验表明，FedeKD 在保持强大预测性能的同时，显著改善了平均情况和最坏情况下的负迁移。额外的消融研究表明，所提出的门控机制在不同的异构程度和超参数设置下仍然有效。

## 2 相关工作

##### 参数聚合方法。

参数聚合仍然是 FL 中的主导范式，其中全局模型是通过聚合本地训练的客户端模型获得的。FedAvg (McMahan et al., 2017) 和 FedProx (Li et al., 2020) 在同质设置下展示了强大的实证性能。FedDyn (Durmus et al., 2021) 通过引入动态正则化项，在训练期间对齐本地和全局目标，进一步提高了异构数据下的鲁棒性。这些方法假设所有客户端共享常见的模型架构并在相对同质的数据分布下运行，允许直接聚合模型更新以提高全局性能。

然而，这些假设在现实世界的 FL 环境中经常被违反。在异构设置中，由于客户端之间的分布不匹配，直接参数聚合通常会失败，导致收敛退化和性能次优。此外，共享模型参数可能会暴露有关本地数据分布的敏感信息，在 FL 设置中引发潜在的隐私问题。相比之下，FedeKD 避免直接聚合私有模型参数，而是在代理空间中进行聚合，使用知识蒸馏作为知识交换的主要机制。这种设计减少了私有模型参数的暴露，从而增强了联邦设置中的隐私保护。此外，FedeKD 并不假设客户端或样本之间的均匀可靠性，而是通过特定任务的模型分歧来近似传输知识的质量，对于分类使用熵校准，对于回归使用连续预测分歧。

##### 联邦知识蒸馏和异构 FL。

为了解决异构设置下参数聚合的局限性，最近的研究探索了基于知识蒸馏的替代通信范式。这些方法不直接聚合模型参数，而是交换辅助信息，如 logits (Huang et al., 2022)、类分数 (Li and Wang, 2019) 或标签级表示 (Yi et al., 2023; Tan et al., 2022)，以促进异构模型之间的协作。最近的方法通过集成学习 (Li et al., 2020)、相互学习 (Yue et al., 2022; Shen et al., 2023) 或模型重组 (Wang et al., 2023) 等技术进一步扩展了这一想法。

虽然这些方法提高了异构 FL 中的灵活性，但它们通常依赖额外的公共数据集或共享数据表示来稳定训练。然而，这种依赖性在实践中引入了两个关键限制。（1）合适公共数据的可用性和选择在现实世界应用中仍然具有挑战性。（2）交换中间表示或模型相关信息可能会暴露有关本地数据分布的敏感信息，引发隐私问题。此外，这些方法通常假设传输的知识具有均匀可靠性，忽视了跨样本和客户端的知识质量变化。

相比之下，FedeKD 不需要额外的公共数据，并将通信限制在轻量级代理模型上。FedType (Wang et al., 2024) 是最接近的先前工作，因为它也考虑了小代理模型和大客户端模型之间的不对称互惠性。然而，FedeKD 在几个关键方面与 FedType 不同。**首先**，关键的方法论差异在于如何表示、优化可靠性并将其与私有模型更新耦合。FedType 依赖于两个一致性模型为客户端和代理模型构建离散的、样本相关的置信集，而 FedeKD 用连续的训练目标替换了集合值可靠性过滤：特定任务的私有-代理分歧直接通过样本级信任权重调节后向蒸馏损失（公式 3）。这种公式提供了在样本层面对知识转移的更细粒度控制，允许模型平滑地降低不可靠信号的权重，而不是依赖基于集合的过滤。**其次**，FedType 的特定一致性集合互惠公式面向分类，因为它是在离散标签空间上定义的。相比之下，FedeKD 直接在模型输出上定义可靠性，使得相同的门控原理可以通过分布分歧应用于分类，通过连续预测分歧应用于回归。

## 3 FedeKD

### 3.1 问题设置

我们考虑一个包含 $K$ 个客户端的 FL 设置。每个客户端 $k$ 访问本地数据集 $\mathcal{D}_k=\{(x_i, y_i)\}$，该数据集不与其他客户端或服务器共享。目标是在异构数据分布下协作改进本地模型，同时保护数据隐私。

每个客户端维护两个模型：一个用于主要学习任务的私有模型 $f_k$ 和一个用于跨客户端知识交换的轻量级代理模型 $g_k$。代理模型在客户端之间共享常见架构并在服务器上聚合，而私有模型保持本地且从不共享。

### 3.2 FedeKD 概述

在每个通信轮次中，FedeKD 实施三个顺序阶段：前向代理蒸馏、代理聚合和基于能量的门控私有模型更新。在前向阶段，每个客户端保持其私有模型固定，仅训练轻量级代理模型以在本地数据上模仿当前私有模型。然后将代理模型上传到服务器并聚合以形成全局代理模型，该模型广播回所有客户端。在后向阶段，每个客户端使用由监督损失和来自全局代理的基于能量的门控后向蒸馏损失组成的组合目标更新其私有模型一次。算法 1 总结了一个通信轮次。

### 3.3 前向代理蒸馏

在每个通信轮次开始时，每个客户端保持其私有模型 $f_k$ 固定，仅在本地数据上训练代理模型 $g_k$。代理模型通过前向知识蒸馏使用特定任务的蒸馏损失来训练以模仿私有模型。对于分类任务，这对应于匹配预测分布，而对于回归任务，则简化为匹配连续预测。这一步骤使代理模型能够从私有模型中捕获与任务相关的知识，同时在客户端之间保持共享表示空间。

### 3.4 代理聚合

本地更新后，每个客户端将其代理模型上传到服务器。服务器聚合这些代理模型以形成全局代理，然后广播回所有客户端。重要的是，聚合仅在代理空间中进行，避免直接共享私有模型参数，从而减少敏感信息的暴露。

### 3.5 基于能量的门控后向蒸馏

为了将知识从全局代理转移回私有模型，FedeKD 引入了一种可靠性感知的后向蒸馏目标，通过基于能量的门控实现样本级信任。给定输入 $x$，计算能量分数 $E(x)$ 以测量私有模型和代理模型之间的分歧，即，$E(x) = \mathcal{E}(g(x), f_k(x))$，其中 $\mathcal{E}(\cdot)$ 可以使用各种差异度量实例化。

对于分类任务，我们采用熵归一化的对称 KL 散度，定义为：

$$
E(x) = \frac{\frac{1}{2}\big(\mathrm{KL}(p \| q) + \mathrm{KL}(q \| p)\big)}{H(p) + H(q) + \epsilon_H}, \quad (1)
$$

其中 $p=\sigma(f_k(x))$，$q=\sigma(g(x))$，$H(\cdot)$ 表示熵，$\epsilon_H=10^{-8}$ 是一个小的数值常数。此公式测量私有-代理分歧，同时通过两个模型的总预测熵校准其幅度。

对于回归任务，我们使用平方误差作为预测分歧的代理，能量简化为：

$$
E(x) = \frac{1}{2}\|f_k(x) - g(x)\|_2^2, \quad (2)
$$

这提供了点预测分歧的连续度量。对于我们实验中使用的标量输出回归任务，这简化为私有预测和代理预测之间的平方差。与分类能量不同，这种回归能量不建模预测不确定性；相反，它测量私有预测和代理预测之间的功能分歧。

##### 设计原理。

所提出的分类能量函数基于对称 KL 散度，并由总熵归一化，如公式 (1) 中所定义。此公式具有三个关键特性。首先，对称 KL 散...