COSMOS：一种基于聚类服务器模型与仅伪标签通信的模型无关个性化联邦学习方法

# COSMOS：基于聚类服务器模型和仅伪标签通信的模型无关个性化联邦学习
来源: https://arxiv.org/html/2605.11165 Luise GeWilliam YeohNing ZhangYevgeniy Vorobeychik Washington University in St\. Louis

###### 摘要

在异构环境中，联邦学习（FL）仍然面临挑战，因为客户端模型通常在架构和数据分布上均存在差异。虽然最近的方法试图通过客户端聚类和知识蒸馏来解决这一挑战，但同时处理架构异构性和统计异构性仍然很困难。我们引入了 COSMOS，这是一个模型无关框架，它利用仅伪标签通信实现服务端个性化。客户端训练本地模型并在公共数据上进行预测；服务器根据预测相似性对客户端进行聚类，利用其自身计算资源为每个组训练特定的集群模型，并将生成的模型蒸馏回客户端。我们提供了首个理论分析，表明从学习到的集群模型进行蒸馏可以实现个性化风险的指数级收缩，这超越了模型无关 FL 中通常提供的收敛到平稳性的保证。在多个基准测试上的实验表明，COSMOS 始终优于所有模型无关 FL 基线，同时与最先进的个性化 FL 方法保持竞争力。更广泛地说，我们的结果突出了带有伪标签的个性化服务端学习，作为一种在高度异构环境中实现可扩展且模型无关的联邦学习的有前景范式。

关键词：联邦学习 个性化联邦学习

## 1 引言

联邦学习（FL）是一种分布式训练范式，客户端在其中协同训练一个或多个服务端模型，而无需公开本地数据\[22 (https://arxiv.org/html/2605.11165#bib.bib30),18 (https://arxiv.org/html/2605.11165#bib.bib15),12 (https://arxiv.org/html/2605.11165#bib.bib4)\]。受客户端数据分布异构性的驱动，大量个性化联邦学习（PFL）方法被开发出来，以针对单个客户端定制模型\[16 (https://arxiv.org/html/2605.11165#bib.bib26),3 (https://arxiv.org/html/2605.11165#bib.bib20),1 (https://arxiv.org/html/2605.11165#bib.bib19),29 (https://arxiv.org/html/2605.11165#bib.bib42)\]。然而，大多数现有的 FL 方案仍然假设客户端模型架构具有一定的结构知识或兼容性。这构成了一个重大的实际障碍，因为客户端可能希望使用最适合其需求的任何模型架构，或者使用他们不愿披露的专有架构。因此，FL 中一个重要的实际需求是成为*模型无关*的，允许使用任意模型的客户端简单地“接入”FL 方案。

模型无关 FL 的含义是它禁止任何关于模型参数或梯度的通信。尽管有明确的需求，但模型无关 FL 的问题仍然研究不足，特别是当客户端同时在数据分布和模型架构上存在差异时。因此，现有的模型无关方法依赖于输出级通信，通常利用共享的未标记数据集\[16 (https://arxiv.org/html/2605.11165#bib.bib26),3 (https://arxiv.org/html/2605.11165#bib.bib20),1 (https://arxiv.org/html/2605.11165#bib.bib19)\]。这些数据集提供了一个共同的参考集，异构模型可以在其中交换预测信号，而无需揭示参数或私有数据。在许多实际部署中，通过公开语料库（例如，网络抓取图像或文本）、合成生成或机构共享的基准池，可以很容易地获得此类数据集。因此，基于预测的通信已成为在 FL 中促进异构模型协作的最实用机制之一。此外，通过使用伪标签而不是参数或梯度，可以显著提高通信效率（见表 4 (https://arxiv.org/html/2605.11165#A3.T4) 中的比较），这已被广泛认为是 FL 中的关键问题，因为无线和其他终端用户连接通常更慢、更昂贵且可靠性较低\[26 (https://arxiv.org/html/2605.11165#bib.bib64)\]。

据我们所知，COMET\[3 (https://arxiv.org/html/2605.11165#bib.bib20)\] 是最接近针对这一交叉点的先前工作。然而，COMET 有三个主要局限性。首先，服务器仅充当被动协调器，忽略了服务端可能具有的大量计算资源，这些资源本可用于促进个性化。其次，其依赖启发式 K 均值聚类，需要预先知道集群数量 K，并且在底层客户端多样性高时缺乏正式的个性化机制。因此，COMET 的理论框架仅限于非凸目标的标准收敛到平稳性，不提供关于个性化性能或风险降低的正式保证。

我们通过 COSMOS（Clustered Output-based Server Models，基于聚类输出的服务器模型）解决了所有这些局限性。虽然高容量服务器有可能协助客户端，但服务器是否能够有效从客户端提供的嘈杂、异构伪标签中学习，这是一个非平凡的问题。COSMOS 通过仔细的算法设计明确克服了这一瓶颈。具体来说，COSMOS 允许客户端在其私有数据上训练任意本地模型，并使用它们在共享未标记数据集上生成伪标签。服务器然后执行距离控制的聚类，将具有相似数据分布的客户端分组，并为每个集群训练专用的教师模型。值得注意的是，COSMOS 不需要数据集与客户端分布完全匹配，只需要它提供输入空间的广泛覆盖即可。如我们在实验中所示（第 6 节 (https://arxiv.org/html/2605.11165#S6)），即使公共数据集仅占整体训练数据的一小部分，COSMOS 仍然有效。

现有模型无关 PFL 文献中的另一个关键缺口是缺乏风险收缩保证。为解决这一问题，我们提供了首次端到端的个性化风险收缩分析。为了实现这一严谨性，我们利用来自半监督学习（SSL）理论\[14 (https://arxiv.org/html/2605.11165#bib.bib54),31 (https://arxiv.org/html/2605.11165#bib.bib53)\] 的标准工具，包括基于扩展的连接性和有界伪标签误差。它们允许我们推导指数风险收缩的一般充分条件，而不限制模型类。重要的是，这些假设仅用于分析，并不对 COSMOS 的实际实施施加约束。

我们的主要贡献如下：

1. **算法框架**。我们提出了 COSMOS，这是第一个模型无关的 PFL 框架，其中服务器使用客户端的伪标签主动训练特定于集群的模型。
2. **理论保证**。我们在充分条件下建立了 COSMOS 个性化风险边界的端到端指数收缩，提供了模型无关 PFL 中首个一般性的风险收缩保证。
3. **实证评估**。我们证明了 COSMOS 不仅始终优于现有的模型无关 FL 方法，而且在同构设置中保持竞争力，同时将通信量从参数共享减少了 1-2 个数量级。

## 2 相关工作

**经典联邦学习**。联邦学习通过 FedAvg\[22 (https://arxiv.org/html/2605.11165#bib.bib30)\] 引入，该方法通过聚合客户端权重更新来训练单个全局模型。虽然简单且通信效率高，但 FedAvg 在非独立同分布（non-IID）数据表现不佳，这促成了诸如 FedProx\[18 (https://arxiv.org/html/2605.11165#bib.bib15)\] 和 SCAFFOLD\[12 (https://arxiv.org/html/2605.11165#bib.bib4)\] 等方法，这些方法通过近端或控制方差校正来稳定优化。尽管如此，这些方法收敛于单一全局模型，不提供个性化。

**模型无关、模型异构和基于知识蒸馏的联邦学习**。由于经典联邦学习传输参数或其更新，它要求客户端和服务端上的所有模型共享相同的架构。为了提供更大的灵活性，提出了一些模型异构方法，例如 FedHeNN\[21 (https://arxiv.org/html/2605.11165#bib.bib39)\] 中传输实例级表示，或 FedProto\[30 (https://arxiv.org/html/2605.11165#bib.bib40)\] 中抽象类原型，以放宽架构同质性假设，许多个性化 FL 方法也是如此（见下文）。然而，虽然模型无关方法必然是模型异构的，但大多数模型异构方法*不是*模型无关的，因为它们仍然施加某些架构约束。此外，模型无关 FL 必然是基于知识蒸馏的 FL（KD-FL）\[23 (https://arxiv.org/html/2605.11165#bib.bib38)\]。尽管如此，许多 KD-FL 方法在某些阶段仍依赖参数聚合\[19 (https://arxiv.org/html/2605.11165#bib.bib9),25 (https://arxiv.org/html/2605.11165#bib.bib46),34 (https://arxiv.org/html/2605.11165#bib.bib44),2 (https://arxiv.org/html/2605.11165#bib.bib2)\]。据我们所知，只有 FedMD\[16 (https://arxiv.org/html/2605.11165#bib.bib26)\] 和 COMET\[3 (https://arxiv.org/html/2605.11165#bib.bib20)\] 探索使用软标签的纯模型无关 FL，而 FedCT\[1 (https://arxiv.org/html/2605.11165#bib.bib19)\] 依赖硬标签。最近，伪标签的通信效率也被用于联邦多视图聚类（例如，CeFMC\[20 (https://arxiv.org/html/2605.11165#bib.bib63)\]），尽管其目标与我们的正交。

**个性化联邦学习（PFL）**。几乎所有实际的联邦学习设置都表现出统计异构性，不同客户端的本地分布可能存在显著差异。个性化联邦学习（PFL）通过学习适应客户端的模型来解决这个问题\[29 (https://arxiv.org/html/2605.11165#bib.bib42)\]。现有的 PFL 方法大致可以根据它们是维护单个共享服务器模型还是少量服务端模型来分组。在第一组中，单个全局模型通过元学习（Per-FedAvg\[8 (https://arxiv.org/html/2605.11165#bib.bib3)\]）、正则化（pFedMe\[6 (https://arxiv.org/html/2605.11165#bib.bib61)\], Ditto\[17 (https://arxiv.org/html/2605.11165#bib.bib59)\]）、本地和全局模型的自适应混合（APFL\[5 (https://arxiv.org/html/2605.11165#bib.bib28)\]）、表示细化（FedBABU\[24 (https://arxiv.org/html/2605.11165#bib.bib62)\]）或基于超网络的参数生成（pFedHN\[27 (https://arxiv.org/html/2605.11165#bib.bib60)\], FedSelect\[28 (https://arxiv.org/html/2605.11165#bib.bib65)\]）适应到每个客户端。在第二组中，聚类 PFL 方法明确维护多个服务端模型并将客户端分配给它们，如 IFCA\[10 (https://arxiv.org/html/2605.11165#bib.bib10)\], FedGroup\[7 (https://arxiv.org/html/2605.11165#bib.bib8)\], AutoCFL\[11 (https://arxiv.org/html/2605.11165#bib.bib5)\], 和 pFedCK\[33 (https://arxiv.org/html/2605.11165#bib.bib27)\]。据我们所知，只有 COMET 也是模型无关的 PFL 方法\[3 (https://arxiv.org/html/2605.11165#bib.bib20)\]。然而，它要求预先指定集群数量 K，并应用标准 K 均值聚类，这是启发式的，不易于进行理论个性化保证。

**模型无关 FL 的理论**。虽然传统 FL 方案的理论收敛结果层出不穷，但模型无关设置以及个性化使得这些结果具有更大的挑战性。最早的模型无关 FL 方法 FedMD\[16 (https://arxiv.org/html/2605.11165#bib.bib26)\] 不提供任何理论保证。最近的 FedCT 方法\[1 (https://arxiv.org/html/2605.11165#bib.bib19)\] 需要一个过度简化的假设，即训练算法*总是*产生单调增加的精度以实现收敛。另一方面，COMET\[3 (https://arxiv.org/html/2605.11165#bib.bib20)\] 的理论分析需要线性模型和高斯数据分布以获得泛化结果。因此，没有针对模型无关个性化 FL 的风险边界收缩的一般充分条件。我们通过采用半监督学习文献\[31 (https://arxiv.org/html/2605.11165#bib.bib53)\] 中的分析工具来弥合这一差距。

## 3 模型

我们考虑一个具有 $N$ 个客户端和服务器的 $M$ 类分类的联邦学习场景。每个客户端 $i \in [N]$ 拥有一个私有标记数据集 $D_i = \{(x_{ij}, y(x_{ij}))\}_j$，其中 $x_{ij}$ 从其输入空间 $\mathcal{X}$ 上的本地分布 $\mathcal{D}_i$ 中独立同分布（i.i.d.）抽取，$y(x_{ij})$ 是 $x_{ij}$ 的真实标签。我们假设每个 $\mathcal{D}_i$ 允许一个密度函数 $p_i(x)$。客户端 $i$ 从假设类 $\mathcal{H}_i$ 中训练一个模型 $f_i: \mathcal{X} \to [0,1]^M$，其中 $\|f_i(x)\|_1 = 1$，对于每个 $i$，$\mathcal{H}_i$ 可以是不同的，例如，代表不同客户端的不同神经网络架构。服务器拥有一个假设类 $\mathcal{H}_S$，并可以训练一个模型*集合* $H = \{h_1, ..., h_K\} \subset \mathcal{H}_S$，其中对于每个 $k$，$h_k: \mathcal{X} \to [0,1]^M$ 且 $\|h_k(x)\|_1 = 1$。在我们的设置中，$K$ 的值是*作为训练过程的一部分内生获得的*。此外，令 $\pi: [N] \to [K]$ 为映射（也是在训练期间获得的），它将每个客户端 $i$ 分配给相应的服务器模型 $h_{\pi(i)}$。最后，遵循先前关于 PFL 的工作，我们假设存在一个大小为 $n = \|U\|$ 的*未标记公共数据集* $U = \{x_j\}$，其中每个 $x_j$ 从 $\mathcal{X}$ 上的全局分布 $\mathcal{Q}$ 中 i.i.d. 抽取。我们假设 $\mathcal{Q}$ 允许一个密度函数 $q(x)$。获取此类未标记数据集并不难，例如，从互联网上抓取（开放许可证）图像或文本，或合成生成。

为了形式化学习目标，我们引入一些额外的符号。令 $g$ 为一个分类模型，其输出为 $M$ 个类别上的分布（即，伪标签）。我们用 $A \circ g(x) = \text{argmax}_{m \in [M]} [g(x)]_m$ 来表示预测的类别（即，在 $g$ 下概率最高的类别）。进一步，令 $R_{\mathcal{D}}(g) = \mathbb{E}_{x \sim \mathcal{D}}[\ell_{0-1}\{A \circ g(x), y(x)\}]$ 表示在分布 $\mathcal{D}$ 下 $g$ 的风险（错误概率），其中 $\ell_{0-1}$ 是 0-1 损失，且 $\text{Err}_{x \in D}(g) = \frac{1}{\|D\|} \sum_{x} \ell_{0-1}\{A \circ g(x), y(x)\}$ 表示数据集 $D$ 上的硬标签预测误差。

**学习目标**：我们的目标是训练 $K$ 个个性化服务器模型 $\{h_k\}$ 的集合，以及客户端到模型的映射 $\pi$，以最小化*个性化风险*，定义为