引言：个性化联邦学习是解决联邦学习中异构问题（如，包含非独立同分布数据的问题）的新范式。然现有的个性化联邦学习方法缺乏关于如何设计模型的个性化和共享部分的标准，例如，手动设计的模型可能会导致比非个性化更差的性能。最近，为自动个性化每个客户端的模型架构和权重，有研究者在人工智能顶级会议AAAI上提出一种新的个性化联邦神经架构搜索算法，即 PerFedRLNAS (Personalized Federated Reinforced Learning Neural Architecture Search)。该算法可以解决以往联邦神经架构搜索工作效率低、不能适应新搜索空间的问题。该研究从理论上证明了，自动分配不同的客户端架构可以解决联邦学习中数据分布的异构性、效率和内存问题。通过与当前先进的相关对比实验分析，也验证了该算法在个性化准确性和总体时间方面可以表现出更好的性能。此外，PerFedRLNAS 对新客户端具有良好的泛化能力，易于在实践中部署。让我们一起来读一读吧！
原文链接：https://ojs.aaai.org/index.php /AAAI/article/view/29576
一、研究背景
联邦学习（Federated learning , FL）已经成为分布式机器学习的一个重要范式，其中大量客户端参与协作训练，同时保持其各自数据的私密性。在联邦学习中，客户端可能具有不同的计算速度和不同的数据分布，这种异质性可能会显著影响联邦学习的性能，特别是使用传统的联邦平均算法 FedAvg [1]。因此，个性化联邦学习作为解决这种异质性并实现更好性能的新范式而出现。个性化联邦学习的目标是为每个客户端训练一个性化的模型。
尽管已经为个性化联邦学习提出了许多解决方案，但每个解决方案都有自己对模型中要个性化的部分的定义。当用户选择解决方案时，缺乏传统的智慧和标准可能会让他们感到困惑。在不同的联邦学习设置下，用户需要选择不同的人工定义个性化策略来达到最优效果。例如，图 1 包含了个性化联邦学习的几种典型方法，并总结了它们如何定义模型的共享部分和个性化部分。不同的方法对我们应该共享或个性化哪些部分有不同的手动设置定义。这种手工设计导致这些方法的性能变化，即使有大量的现有工作，也未能充分回答两个问题：（1）为什么个性化模型可以解决联邦学习中的这种异构问题？（2）既然有这么多的个性化策略，是否有一种标准和自动化的方式来决定客户端模型的个性化和共享部分，以达到最佳性能？
除了传统的个性化联邦学习方法外，还提出了利用自动机器学习（Automated Machine Learning，AutoML）的新方法来寻求实现更好性能的自动方法。联邦神经架构搜索（Federated Neural Architecture Search, FedNAS）[2] 被提出，用于在联邦学习设置下直接进行神经架构搜索（Neural Architecture Search, NAS），每个客户端可以获得不同架构和不同权值的模型。然而，以前的工作要么只搜索所有客户端共享的全局神经架构，要么手动定义客户端集群，而不是为每个客户端搜索个性化架构或逐步对每个客户端进行神经架构搜索。在传统的个性化联合学习中，使用 NAS 的超级网络对固定模型进行这种简单的替换，仍然无法在不进行手动设置的情况下展示个性化模型的优势。
在本文中，为了建立一个以自动化方式决定模型的个性化和共享部分的新标准，我们提出了一个新的框架，称为通过强化学习的个性化联邦神经架构搜索(PerFedRLNAS)[3]。PerFedRLNAS 作为一个通用框架工作，并且可以部分地自动个性化模型，而无需手动设置个性化和共享部分。在我们的框架中，每个客户端都有一个具有自定义权重的自定义架构模型。他们的模型的哪些部分需要共享，哪些部分需要个性化，哪些部分应该在客户端之间共享，这些客户端将根据最优策略自动决定。换句话说，模型的个性化和共享部分是由目标直接决定的。回到第一个问题，我们展示了在每个客户机上拥有部分个性化的架构和模型可以帮助解决数据和系统异构问题。
除了无法回答这两个核心问题外，目前个性化联邦学习和联邦神经架构搜索解决方案存在的另一个限制是缺乏适应各种深度神经网络模型的灵活性，这使得它们无法获得更好的性能。然而，神经网络的结构发展非常迅速，我们提出的新框架可以用任意的 NAS 超级网络模型来实现。如果以后提出了一个新的有前途的 NAS 超级网络模型，PerFedRLNAS 可以将其应用到联邦设置中。
此外，还有另一个限制，使以前的联邦 NAS 方法无法有效地应用于个性化联邦学习：低效率。我们将架构搜索阶段完全保留在服务器上，并且在服务器和客户端之间只传输一个模型，而不是整个超级网络。使用我们的方法，客户只需要在正常比例的神经网络模型上执行局部训练，并且不会产生额外的通信和局部计算开销。总体效率与 FedAvg 相当。
最后，我们的经验表明，PerFedRLNAS 优于 FedAvg、本地训练和现有的最先进的个性化和异构联邦学习方法。PerFedRLNAS 在不同的目标、不同的异构设置、不同的骨干模型和搜索空间中，在准确率和效率方面都能取得更好的性能。我们还验证了 PerFedRLNAS 可以很好地泛化到新客户机。
二、本文方法
2.1 问题描述
在个性化联邦学习中，我们总共有 个客户端，每个客户端都有自己的本地目标。我们表示它们为 。在个性化联邦学习中，自定义权重 不需要与全局权重相同，通常也不相同。它是特定于每个客户机的。所以，问题是描述为
其中，我们试图使每个本地数据集的性能都很好。 如果 ，则只使用全局模型。如果 ，其中 是一个共享的全局模型和 是损失函数，这与 FedAvg 的标准公式是一致的。我们使用联邦神经架构搜索算法来自动确定共享权值，并自定义客户端的局部权值。公式描述为
其中， 是超级网络， 表示各局部模型的架构参数， 是一个采样函数：例如，通过 softmax 函数传递 ，在超级网络的每条边（Edge）上选择一个操作（Operation）。
2.2 建立个性化的联邦 NAS 框架
我们用 表示放置在服务器上的超级网络，或者我们称之为搜索空间。超级网络可以是任何架构，每个架构都有许多不同的搜索维度和不同的候选操作。每个搜索维度不需要有相同数量的候选操作。我们使用表 1 中所示的示例搜索空间来更好地说明我们的方法。
在一个超级网络中，有 个搜索维度。在这个例子中，卷积块的搜索维度为宽度、块数和核大小；transformer 块的搜索维度为宽度和块数；输入分辨率的搜索维度为宽度。因此，有 6 个搜索维度。对于每个搜索维度，我们分配每个选择的概率： ，其中 表示每个维度的索引， 表示选择的个数。在本例中，当 时，表示卷积的内核大小维度。我们有三个选项：3, 5, 7。所以，我们有 。我们知道 是选 3 的概率， 是选 5 的概率，和 是选 7 的概率。我们有 。 接下来，我们为每个概率 分配一个架构参数 ，其中 是客户端索引。我们可以通过 softmax 得到概率。
对于每个客户端，我们可以分配这样的架构参数。我们用 来表示每个客户端的模型架构。对于每个 ，我们可以采样并生成一个独热向量（one-hot vector） 。基于采样策略 ，我们可以为客户端 采样 和所有 。因此，我们可以用这样的公式来表示任意的超级网络模型。
为了在每一轮的聚合过程中将单个模型聚合到超级网络中，我们遵循了以前联邦神经架构搜索方法中的常规聚合方法。我们直接将单个模型扩展到与超级网络相同的维度，填充零并进行聚合。当我们进行平均操作时，我们将模型权重除以每个参数的非零出现次数。
2.3 通过强化学习解决 PerFedNAS
在下一步中，我们将我们的问题视为马尔可夫决策过程(Markov Decision Process，MDP)，并通过强化学习（Reinforced Learning，RL）说明架构搜索的解决方案。我们在这里使用政策在线 RL（On-policy Online RL）。将 RL 方法应用于 FL 的优点是，每个客户端都可以通过评估动作的单个智能体（Agent）来解决。每个客户单独观察奖励。
马尔科夫决策过程：我们将每个通信回合视为 MDP 的一个步骤，并遵循强化学习中的常规步骤来定义状态和动作:
1. 状态（State） : 状态包括两个部分:超级网络 和 客户端参数 。虽然这些是客户端的状态，但它们只在服务器上存储和处理。 的初始值均设为 0，表示所有操作被选择的概率相同。
2. 动作（Action）：每个客户端将在服务器上放置一个虚拟智能体，负责搜索每个客户机的架构。在每一轮通信中，该代理将根据 从超级网络中采样一个个性化的模型服务器发送给客户端。然后，每个客户端将对其本地数据集进行本地训练，并将训练好的模型连同测试精度一起返回给服务器。我们用 表示每个智能体的动作。
3. 策略（Policy）：我们可以使用我们设计的架构参数来表示我们的策略。 表示每个虚拟智能体在第 轮通信时的策略。状态为 ，根据奖励 ，采取动作 使状态变为 。
在最终状态下，每个客户端的 将表示搜索到的每个客户端的最佳架构，客户端的模型权重可以直接部署，无需进一步训练。
通过策略梯度进行搜索：因此，关键是对策略参数进行优化，为每个客户找到最佳的个性化模型。在每一轮通信中，对于每个参与的客户端，它将有一个奖励（Reward），我们利用这个奖励 来更新策略网络，该策略网络由 参数化。我们通过梯度递增来最大化从超级网络中采样的子模型的期望奖励:
的含义是在搜索维度 上，选择的选项是 。我们可以进一步计算：
其中假设 的第 项是 1，对于 :
有了这样的公式，我们就可以优化策略网络（这是一层 softmax）。这使得服务器除了聚合之外几乎没有额外的开销。
奖励函数的设计：为了适应联邦学习中任意程度的异质性，我们可以根据我们的目标设计奖励函数。在这里，我们通过设计一个考虑准确性和效率的奖励函数来解决统计和计算速度异质性的目标： ，以添加每个客户端的一轮时间（Round Time， RT）和本地推理精度之间的权衡。由于每个客户在每轮通信中都有一个奖励，以反映他们各自的表现，因此每个奖励不会有极端的急剧变化。一轮时间 是从发送模型到接收客户端更新的模型参数所花费的时间。然而，直接使用精度和一轮时间的值会导致训练不稳定。因此，在每一轮中，我们只需要在历史信息中更新这一轮参与的客户的准确率，并计算客户之间的平均准确率。并且我们仍然可以像平常的 FL 一样，每轮选择一部分客户端参与。客户端除了发送模型参数外，只需要发送精度，这大致不增加通信开销。一轮时间 是实时的。min 表示每轮通信中参与的客户端所花费的最短时间。为了解决稳定性问题，可能是由于不同客户的训练时间差异太大造成的，我们使用 时间来调节精度和一轮时间 到相同的尺度.
2.4 PerFedNAS 算法
利用 PerFedRLNAS 框架和强化学习解决方案，我们可启动个性化的联邦学习，如算法 1 所示。整个过程如图 2 所示，服务器将对客户端模型进行采样并将其分发给客户端。站在客户端的角度来看，它们的职责与 FedAvg 中的完全相同，因此，我们的算法不会在客户机上增加额外的开销。在服务器上，它执行架构参数的详细更新是通过虚拟 RL 的智能体们对每个客户机的聚合。由于策略网是一个带有一层 softmax 的浅网，因此服务器上的开销可以忽略不计。在每一轮中，服务器接收 并将 聚合到超级网络。
三、实验结果
我们将 PerFedRLNAS 与最先进的个性化联邦学习算法和以前的联邦神经架构搜索算法进行了实证比较，以了解我们的方法如何很好地解决异构问题。
3.1 数据集、任务和模型
我们使用 CIFAR10 和 CIFAR100 研究图像分类任务。为了衡量不同方法的灵活性并进行综合评估，我们采用 ViT（NASViT [4] 作为搜索空间，LeViT [5] 作为固定模型）和 CNN （MobileNetV3 [6]和 DARTS [7]作为搜索空间，MobileNetV3- large 作为固定模型）。NASViT 是一次性和块搜索的超级网络。MobileNetV3 是块搜索和进化搜索的超级网络，而 DARTS 是路径搜索和可微搜索的超级网络。
3.2 非独立同分布数据集和实现
我们使用参数 的狄利克雷分布（Dirichlet Distribution）来生成非独立同分布数据集。 越小，说明假设的非独立同分布程度越强。所有实验都用的相同的非独立同分布数据集。训练样本和测试样本在所有客户端上被平均划分。我们主要考虑跨设备场景的情况，其中一些客户端参与联邦学习过程。
3.3 个性化的性能
为了公平地比较结果并模拟真实的联邦学习设置，在每一轮中，从 100 个客户端中选择 5 个。在每一轮沟通中，每个客户端进行 5 次局部训练。我们根据准确性和效率来评估性能。关于准确性，我们报告了这些 100 个客户的本地数据集的平均测试精度和标准偏差。为了提高效率，我们直接测量整个系统达到报告准确性的延迟，包括花费在客户端、传输和服务器上的时间。我们将上传和下载数据的传输速率设置为 100Mbps。在所有实验中，随机种子都是固定的。不同运行的实验之间的方差远小于不同客户端的局部精度之间的方差（例如，在 NASViT 搜索空间上使用 PerFedRLNAS 时，使用不同的种子运行 5 次，结果是 0.5 比 4.11）。所有的方法都被训练成收敛的。
3.4 与传统方法比较
对于传统的个性化 FL，我们选择了 FedRep [8]、FedBABU [1] 和 FedTP [9]，因为它们是个性化客户端部分模型的代表性工作。为了建立一个基本基准，我们还比较了 PerFedRLNAS 和 FedAvg（在训练后对局部数据集进行 10 次微调），并且只进行局部训练。在 FedTP 方法中，我们直接使用了本文自行设计的结构。由于我们的方法对不同的搜索空间具有灵活性，因此我们将 ViT 和 CNN 结构化模型进行比较，并将数值结果显示在表 2 中。
共享是个性化的必要条件：从结果来看，只有本地训练表现最差，并且在所有客户端之间存在很大差异。即使是 FedAvg 也可以提供比当地训练更大的改进。另一方面，如果我们采取合理的个性化策略，可以大大提高性能，经验表明我们需要个性化来解决异质性。我们的方法在共享部件和个性化部件之间进行了很好的权衡，取得了很好的性能。
PerFedRLNAS 在准确性和效率方面表现出更好的性能：当使用最先进的手动调优模型时，传统的最先进的方法带来的改进很少。当数据是非独立同分布时，PerFedRLNAS 在准确性和效率方面都有更大的改进。FedTP 虽然有很大的改进，但存在较大的方差。除此之外，它还缺乏将自己应用于任意 transformer 模型的灵活性，例如 LeViT，其注意图（Attention Maps）在不同层上具有不同的大小。虽然 FedTP 有更小的定制模型，但我们的方法在 CIFAR10 上只需要 3.29 小时，在 CIFAR100 上只需要 1.13 小时就能达到相同的精度。尽管我们搜索的模型与手动调整的模型具有相似的架构，但我们需要更少的沟通开销和训练时间。
3.5 与联邦 NAS 方法比较
在表 3 中，将我们的方法与以前的联邦 NAS 工作进行比较。对于 FedNAS [2]、FedRLNAS [10]、FedorAS [11] 等多阶段方法，结果是在应用了所有阶段之后生成的，而不仅仅是超级网络训练阶段。对于 FedorAS，它将客户机划分为四个集群，并在最后阶段为每个集群训练一个模型。因此，我们报告了每个集群的个性化模型的准确性。为了公平比较，所有方法都是在 DARTS 搜索空间上实现的。由于 DARTS 是一个 CNN 结构化搜索空间，我们可以看到，使用 CNN 模型 MobileNetV3，联邦 NAS 方法获得的结果优于传统个性化联邦学习方法获得的结果。原因是联邦 NAS 可以在更大的空间中寻找需要个性化的部分。此外，与 FedNAS、SPIDER [12] 等方法相比，PerFedRLNAS 的传输开销要小得多，因为前者需要跨客户机和服务器传输整个搜索空间。
PerFedRLNAS 是一种单阶段算法，自然地结合了预热、后期微调和再训练。 在 PerFedRLNAS 的开始阶段，所有的架构都具有接近的采样概率，这与预热阶段相同。在接近趋同的最后阶段，由于架构大多是固定的，模型的个性化部分仅由特定的客户集群更新。此过程具有与后微调阶段相似的功能。因此，PerFedRLNAS 不需要预热或后期微调，但与 FedNAS、FedRLNAS 和 FedorAS 等多阶段方法相比具有竞争力。
3.6 为什么个性化模型可以解决联邦学习中的异构问题?
回到第一个问题，很直观地，不同的架构可以解决系统异构问题，因为更强大的设备可以得到更大的模型，反之亦然。关于统计异质性，如果所有客户端共享相同的分布，我们只需要一个全局模型。但是，当不同的客户端具有不同的分布时，每个客户端上优化的架构应该是不同的。另一方面，每一种架构彼此之间都有相似之处。具有更紧密数据分布的客户机将具有更相似的架构。因此，这些客户端将有更多的模型部分可以共享。通过这种方式，客户端可以自动形成任何规模的集群，而无需手动定义客户端集群或向服务器泄露数据集隐私。
3.7 资源预算意识
除了解决数据异构和效率异构之外，我们还可以考虑不同客户端可能具有不同内存预算的情况。因此，我们设计了另一个奖励函数来实现资源(内存)预算意识，根据客户的内存限制为客户分配不同的架构。我们在设计这样一个奖励函数时考虑了内存利用率。
其中， 是每个客户端设备的内存利用率。如果设备不能将当前分配的模型装入物理内存(GPU)，则其 将为 1，精度将为 0。在这种情况下，我们将 设置为当前一轮参与的所有客户的最大的一轮时间。
本文在 CIFAR10 0.3 数据集和 CIFAR100 0.1 数据集上重复采用 NASViT 超级网络进行训练的实验。NASViT 涵盖了从 159.22 M Flops 到 1293.98 M Flops 的广泛模型。我们用从 2GB 到 16GB 的统一样本模拟每个客户机的内存预算。在图 3 中，我们展示了与奖励函数中没有内存预算的用例相比，推理精度的训练曲线。CIFAR10 的最终测试准确率为 82.87±4.87%，训练时间为 27.84 小时。CIFAR100 的最终测试准确率为 63.24±4.07%，训练时间为 26.56 小时。我们还在图 4 中显示了每个客户端的内存预算和每个客户端上搜索到的架构所消耗的内存。我们可以看到，每个设备可以根据其内存预算接收不同架构的模型。通过这个例子，我们展示了 PerFedRLNAS 可以灵活地通过设置不同的奖励函数来实现不同类型的异质性。
3.8 泛化能力
在超级网络训练之后，不参与联邦训练的新客户端只测试接收到的模型。我们选择了 100 个拥有非独立同分布数据集的客户端，在联邦训练中未使用。图 5 显示经过 轮架构搜索后这 100 个新客户机的准确性。超级网络的权重保持不变；在此过程中，客户端既不训练模型，也不上传权重。
3.9 不同联邦设置的消融研究
在这里，我们研究了不同联邦设置下的鲁棒性。在控制其他变量的同时，我们将每轮的参与者人数改为 50 人。接下来，我们将总客户 K 更改为 200，但使用相同的样本率（10%）来控制变量。从图 6 中所示的沿时钟时间的收敛来看，我们可以看到 PerFedRLNAS 的性能仍然优于基线，因为它花费更少的时间来达到相同的精度，并最终收敛到更高的精度。我们还研究了移动时间项目对奖励函数的影响。在使用使用 MobileNetV3 时，设置 CIFAR10 0.3 时，我们将 设置为 0。精度可提高到 83.33%±4.97，延迟为 10.65 小时。当我们不把效率作为目标时，精度得到了提高，但效率却降低了。
四、总结
在本文中，我们提出了一个新的联邦神经结构搜索框架 PerFedRLNAS。它可以通过对各种异构目标的优化，包括非独立同分布，效率和内存预算，直接在每个客户端上自动搜索定制架构和权重。该框架为个性化联邦学习中模型个性化部分的自动设计提供了解决方案。最后，我们展示了 PerFedRLNAS，在非独立同分布数据集的局部精度和效率指标，优于其他最先进的个性化联邦学习方法。此外，它对新的客户端具有很好的泛化性，这对我们在推理中轻松部署客户端是一个很好的特性。PerFedRLNAS 在不同的联邦设置中也是鲁棒的。
参考文献
[1] OH, JAE HOON, Sangmook Kim, and Seyoung Yun. FedBABU: Toward Enhanced Representation for Federated Image Classification, International Conference on Learning Representations (ICLR), 2022.
[2] He C, Mushtaq E, Ding J, et al. FedNAS: Federated Deep Learning via Neural Architecture Search. CVPR 2020 Workshop on Neural Architecture Search and Beyond for Representation Learning, 2021.
[3] Yao D, Li B. PerFedRLNAS: One-for-All Personalized Federated Neural Architecture Search, Proceedings of the AAAI Conference on Artificial Intelligence, 38(15), 16398-16406. 2024.
[4] Gong C, Wang D. NASViT: Neural Architecture Search for Efficient Vision Transformers with Gradient Conflict-aware Supernet Training. ICLR, 2022.
[5] Graham B, El-Nouby A, Touvron H, et al. LeVit: A Vision Transformer in Convnet's Clothing for Faster Inference. In Proc. the IEEE/CVF International Conference on Computer Vision (ICCV), 12259–12269.
[6] Howard A, Sandler M, Chu G, et al. Searching for Mobilenetv3. In Proc. the IEEE/CVF International Conference on Computer Vision (ICCV), 1314–1324.
[7] Liu H, Simonyan K, Yang Y. DARTS: Differentiable Architecture Search. In Proc. International Conferenceon Learning Representations (ICLR). 2018.
[8] Collins L, Hassani H, Mokhtari A, et al. Exploiting Shared Representations for Personalized Federated Learning. International Conference on Machine Learning. PMLR, 2021.
[9] Li H, Cai Z, Wang J, et al. FedTP: Federated Learning by Transformer Personalization. IEEE Transactions on Neural Networks and Learning Systems, 2023.
[10] Yao D, Wang L, Xu J, et al. Federated Model Search via Reinforcement Learning. 2021 IEEE 41st International Conference on Distributed Computing Systems (ICDCS). 2021.
[11] Laskaridis S, Fernandez-Marques J, Dudziak Ł. Cross-device Federated Architecture Search. Workshop on Federated Learning: Recent Advances and New Challenges (in Conjunction with NeurIPS 2022). 2022.
[12] Mushtaq E, He C, Ding J, et al. SPIDER: Searching Personalized Neural Architecture for Federated Learning. arXiv preprint arXiv:2112.13939, 2021.
初稿 ：廖鹏
复审 ：颜学明
终审 ：金耀初
