Federated optimization survey(2):实用的联邦学习算法设计

Federated optimization survey(1),对联邦学习的研究背景和问题定义进行介绍。本篇中将讨论实用的联邦学习算法的设计。主要包含以下两方面的内容:

  • 实用算法的设计指南:介绍联邦算法设计的通用指南,从中可以理解设计实用联邦算法时候需要考虑的关键问题
  • 几个具有代表性的算法例子:介绍一些有代表性的技术和优化算法,其中说明了如何解决联邦场景下的各种挑战,主要为了突出设计原则并阐明可以应用它们的情况

1. 实用算法的设计指南

1.1 阐明应用场景

联邦学习可以应用于许多不同的应用程序,包括前面Federated optimization survey(1)中讨论过的cross-silo和cross- device场景设置。这些应用场景具有不同的属性,因此特定算法可能不适合所有应用场景。在提出新的优化算法时,应考虑该特定设置的约束和要求,并明确为其算法指定最合适的应用设置。

比方说cross- device场景,需要设计无需跨轮次保持客户端持久状态(state)。一方面是出于隐私要求,但更多的是由于客户端数量通常非常多具体某个客户端通常只会参加一次训练。另外,即使在最好的情况下,也就是特定的客户端能够多次参与,但是在具体参与间隔之间可能会经过许多轮,因此对应的状态可能会过时,并对收敛产生不利影响。

大量客户端和间歇性可用性也意味着服务器无法计算任何需要跨所有客户端聚合的全局信息,这里的全局信息包括全局目标值 F (x)、客户端总数 M 、样本总数等。

1.2 提高通信效率

在联邦学习中(尤其是在跨设备的情况下),客户端可能会遇到严重的网络延迟和带宽限制。实用的 FL 算法通常使用减少通信机制来实现可用的计算与通信比率。接下来将介绍三种降低通信成本的常用方法: (i) 通过允许本地多次更新来降低通信频率; (ii) 通过压缩信息减少通信量; (iii) 通过限制每轮的参与客户端的数量来减少服务器的通信流量。这三种方法是正交的,可以相互结合。

关于通信效率需要注意的是,在本小节中,如果一种方法可以降低每次模型局部更新(即每次局部迭代)的通信成本,则我们认为该方法通信效率高。但是,总通信成本是局部迭代次数与每次迭代通信成本的乘积。因此,在介绍高效通信方法时,还讨论了它们对实现目标精度的迭代总数的影响。

使用多次本地更新

许多工作已经通过经验验证了本地更新在减少通信轮数以学习准确模型方面的有效性 。当客户端执行 τ 个模型更新时,每次客户端模型更新的通信成本可以有效降低 τ 倍。

虽然使用多次本地更新确实对于提高通信效率具有显著的效果,但对比 SGD 等方法来说,这会使得理论上分析采用本地更新的方法(例如 FedAvg)的收敛更为困难。许多工作表明,FedAVG算法中一个额外的误差项会随着局部steps的数量的增加而增加,但是这个误差项与其他项相比具有更高的阶数,在衰减学习率时可以忽略不计。而当客户端数据是Non IID 时,多次本地更新导致的副作用会更明显。因此,通信减少和收敛之间存在权衡当使用较大数量的本地更新时,算法会降低通信频率,节省通信;然而,它可能会在收敛时产生更高的误差。在实践中,可以选择局部步骤的最佳值来平衡这种Trade-off,或随着时间的推移对其进行调整。

采取多个局部步骤的替代性方法是使用较大的mini-batch。尽管local-update SGD方法在经验上优于大 mini-batch方法,但从理论角度来看,它的优越性并未得到很好的理解 。

应用压缩技术

减少客户端和服务器之间传输的比特是另一种节省通信成本的有效方法,并催生了大量的方法来稀疏化或量化模型更新,同时不会显着降低准确性 。压缩技术对于高维模型尤其有效,其输出可以看作是模型更新的随机估计,因此可能是无偏也可能是有偏的。尽管有偏压缩方法可以实现更高的压缩比,但直接应用它们会导致无法收敛。在全程参与且客户端具有持久状态(例如 cross-silo FL)的情况下,可以使用错误反馈策略。同时有偏和无偏压缩方法以及误差反馈方案也可以与优化方法(例如动量)相结合,以提高收敛性 。另一种方法将梯度压缩合并到模型架构本身中。例如,使用低秩近似的因式分解模型架构可以大大减少通信模型更新所需的通信量。

压缩方法的另一个需要考虑的重要方面是它们与分布式系统中聚合协议的兼容性。例如,许多分布式系统使用 all-reduce 风格技术来跨计算节点聚合梯度或本地模型更新。在这样的系统中,与 all-reduce 不兼容的压缩技术可能会提供较低的通信效率,尽管它们具有更高的压缩比 。在这样的设置中,重要的是要确保压缩操作与加法兼容,即压缩梯度之和必须等于(至少在预期中)梯度之和的压缩,这对于许多类似于secure aggregation的聚合协议也是必需的。

客户端采样

通过对较小的子集进行采样来减少每一轮参与客户端的数量可以隐式降低通信成本,因为每轮通信延迟随着接收者/发送者的数量单调增加。此外,当客户端具有随机或异构计算速度时,该策略可以减轻掉队效应。例如,如果我们假设每个客户端的计算延迟是一个指数随机变量,那么等待最慢的客户端的额外时间约为 O(log M),其中 M 是参与客户端的数量。参与客户端的数量不仅会影响通信时间,还会影响算法的收敛特性。每轮使用太少的客户端可能会显著增加训练过程中的随机噪声。如何设置适当数量的参与客户仍然是开放的问题,在相关工作中探索较少。

1.3 处理数据和计算异构性

由于客户端数据的不平衡性和异构性,客户端的本地优化目标通常不相同,并且可能不共享共同的minimizer。因此,当从同一个全局模型执行局部更新时,客户端将漂移到局部目标的最小值并最终得到不同的局部模型。这种现象通常被称为客户端漂移,并降低了执行多个本地更新步骤的好处

联邦学习中另一个重要但研究较少的异质性来源是计算异质性。与在专用和同类机器上执行计算的集中式设置不同,联邦设置中的客户端可能由于硬件差异而具有不同的计算速度。例如,一些客户端可能是手机,而另一些可能是配备 GPU 的笔记本电脑。同样,一些客户可能将他们的资源专门用于训练,而其他客户可能需要与后台任务共享资源。

虽然许多算法假设客户端采取相同数量的步骤,但这可能会由于计算异构性而导致掉队,并可能增加算法的运行时间。生产系统可能会强制执行超时限制,导致掉队方的更新被丢弃。为了避免此类问题,一些工作通过限制任何客户端可以在本地更新的步数,同时对客户端进行采样的概率与他们拥有的数据量成正比的策略。另一种解决方案是允许跨客户端的本地步数可变。例如,允许客户端在给定的时间窗口内执行尽可能多的本地步骤。但是在数据异质性的情况下,这种方法的加剧了客户端漂移,并导致 FedAvg 式算法的收敛点与实际经验损失函数的最小值之间存在无法弥合的不一致。但相关工作同时表明这种由多步本地更新方法所引起的不一致可以通过仔细重新加权客户端更新来完全消除。

因此设计实用联邦算法的时候需要考虑允许数据异构和可变数量的本地计算设计。如果采用这种方法,我们建议修改客户端权重策略以考虑步骤中的这种可变性(例如,对采取更多本地步骤的客户端给予较少的权重)。此外,在一些实际场景中,客户端可能只有非常有限的计算或内存容量来完成一次本地更新。为了解决这个问题,相关工作也探索了允许客户根据自己的计算或内存预算使用不同大小的模型的策略。

1.4 兼容系统架构和隐私保护协议

实际的联邦学习系统通常很复杂,并包含许多隐私保护协议(例如安全聚合和差分隐私机制)。这种架构复杂性可能会给实际的联邦学习算法带来额外的约束。虽然设计一种与所有可能的系统架构和隐私保护协议兼容的算法具有挑战性,但建议研究人员 (i) 明确指定目标系统和隐私保证,(ii) 讨论其方法的局限性和不兼容性,以及 ( iii) 考虑修改他们提出的算法,使它们能够在不同的系统中运行。

下面讨论由于系统架构和隐私保护协议而导致的额外约束的两个例子。

不活动客户端采样策略

虽然cross-device设置在训练过程中可能有大量客户端,但由于各种系统限制,在给定时间,服务器可能只能与一小部分客户端通信。在许多实验模拟中,用于在给定轮次执行训练的客户队列是通过从总客户端群中随机均匀抽样获得的。理论上通常基于所有客户始终可以进行抽样的假设,但是这种假设在实际的跨设备联邦学习系统中通常是不正确的,在这种系统中,客户通常根据自己的状态决定是否可用,例如,移动客户端只有在连接到 WiFi 并正在充电时才可用。这种对客户端采样的约束被称为不活动客户端

隐私保护和加权策略

隐私保护方法通常应用于 FL 算法的聚合过程(客户端到服务器通信),例如安全聚合 和差分隐私 。在FedAvg的许多变体中,聚合过程可以表示为客户端更新的加权求和聚合权重的自然选择是客户端上的样本数量,因此当每个客户端仅执行一次梯度更新迭代(即 FedSGD)时,加权和的期望等于 ERM 目标的批量梯度。这种基于样本的加权方案被用于原始的 FedAvg 算法和许多其他联邦学习论文。然而,不均匀的权重会增加单个客户端的更新对聚合更新的影响,这与隐私目标相悖(边界敏感性是差分隐私的关键要求,更具挑战性)。

在实践中,DP-FedAvg在联邦学习中应用差分隐私时应用统一的加权方案(即 $p_i = 1/S$),但是统一加权本身是否有利于隐私保护、对异常值的鲁棒性或对样本数量少的客户的公平性,仍然是一个未解决和验证的问题。

2. 提高技术的代表性方法

2.1 与动量和自适应方法的配合

动量和自适应优化方法(如 Adam、Adagrad、Yogi )已成为训练深度神经网络的关键组成部分,并且可以经验上提高 SGD 的优化和泛化性能。最近的许多论文工作研究了如何在 FL 中结合动量和自适应,并验证它们在加速收敛方面的有效性。具体包括以下几种结合方式

Server和client的优化框架

工作[1]中证明了 vanilla FedAvg 隐式使用了双级优化结构。据此,他们提出了一个名为 FedOpt 的通用算法框架(Algorithm 1)。当客户端使用 ClientOpt 来最小化本地训练损失时,服务器将聚合的本地模型更新作为伪梯度,并将其用作 ServerOpt 的输入以更新全局模型。FedOpt 框架可以将 ClientOpt 或 ServerOpt 设置为 SgdM、Adam、Yogi 等,允许使用任何带动量和自适应的优化器

截屏2023-02-21 19.18.19

在server上应用动量和自适应方法

直觉上,将聚合的客户端本地更新视为伪梯度,服务器端的动量与每个轮次增加所选客户端数量具有类似的效果由于伪梯度的指数加权平均,前几轮选择的客户端的模型变化也有助于本轮的全局模型更新。

本着同样的思想,工作[1]进一步证明了在跨设备联邦学习中使用服务器端自适应性的好处。特别是,通过将 Adam 和 Yogi 分别设置为 ServerOpt 和 SGD 作为 ClientOpt,FedAdam 和 FedYogi 将自适应学习率纳入联邦优化,并在多个基准联邦学习数据集上实现比 vanilla FedAvg 更快的收敛和更高的验证性能。此外,相关结果表明自适应方法更容易调节(即对超参数变化更稳健)。

在将 ClientOpt 保持为 SGD 的同时使用服务器端动量或自适应方法的另一个好处是,它不会增加客户端的计算复杂性或每轮的通信成本。它还兼容极低的客户端采样率(小于 1%)。上述所有好处都突出了服务器端动量或自适应方法在cross-device FL 设置中的实用性。

在Clients上应用动量和自适应方法

鉴于服务器端动量和自适应方法,自然会出现一个问题,即我们如何将动量或自适应直接应用于客户端?直觉上,在每一步更新使用一阶和二阶矩可能比在服务器端自适应方法中每 τ 步骤使用它们更好。然而 ,当客户端在本地执行动量或自适应优化方法时,它们的优化器状态是单独更新的,因此,由于Non IID 数据分布,它们可能会彼此偏离。为了解决这个问题,一些工作提出了同步状态策略,其中本地模型和客户端优化器状态在每一轮都由服务器同步。

最近一些工作表明,即使在每一轮开始时将客户端优化器状态重置为默认值,客户端自适应方法也是有益的,并且可以与服务器自适应方法结合以进一步提高性能。此外,这种重置状态策略不会在服务器和客户端之间产生任何额外的通信。

动量和自适应的惰性更新框架

动量和自适应的惰性更新框架的主要出发点是避免客户端优化器状态之间的偏差(即自适应优化器中的一阶和二阶矩),具体做法是是在本地训练期间修复优化器状态,并在每轮结束时在服务器上延迟更新它们。与上述服务器端自适应方法相比,有两个关键区别

(i)惰性更新算法允许客户端在本地使用动量或自适应方法,但它们的优化器状态是固定的或保持同步的;

(ii) 在这些惰性更新算法中,优化器状态通过同步所有客户端的本地优化器状态或使用在当前全局模型上评估的批梯度在服务器上更新。

2.2 降低本地模型更新的偏差

如前所述,FedAvg 和其他联邦优化算法通常通过允许客户端在与服务器同步之前每轮采取多个 SGD 步骤来降低通信成本。然而,在具有异构客户端数据的设置中,采取更多本地步骤实际上会阻碍收敛,因为生成的客户端更新变得偏向于本地最小值。下面介绍了两种有助于减少每个客户端本地模型更新偏差的方法。

控制偏差

控制变化是在标准凸优化相关工作中开发的一种技术,用于减少有限和最小化问题中随机梯度的方差,从而加快收敛速度。可以使该技术适应联邦学习设置以减少客户端之间的差异。在 FL 中采用这种技术的代表性算法是 SCAFFOLD [2],如下所述。

SCAFFOLD 算法适用于cross-silo设置,并利用每个客户端存储的持久状态。该状态是客户端 i 的控制变化的值 $c_i$, $c_i ≈ ∇F_i(x)$ 用于估计关于客户端本地数据的损失梯度。服务器维护所有客户端状态的平均值作为它自己的可控制变量 $c$,它在每一轮中被传递给所有选定的客户端。客户端在每一轮中执行多个步骤的 SGD,就像 FedAvg 一样,但为每个随机梯度添加一个校正项 $c-c_i$, $∇F_i(x) + c − ci ≈ ∇F (x)$,校正的效果是消除每个客户端更新步骤的偏差,确保它们更接近全局更新,这使得SCAFFOLD 能够比普通的 FedAvg 收敛得更快,而无需对数据异质性做任何假设

在实施 SCAFFOLD 时,有许多可能的可控制变量选择。例如,可以选择使用上一轮的平均局部梯度作为 $c_i$。具体来说,局部更新后,局部控制变量更新如下:

\(\boldsymbol{c}_{i}^{(t+1)}=\left\{\begin{array}{ll} \boldsymbol{c}_{i}^{(t)}-\boldsymbol{c}^{(t)}+\frac{1}{\eta \tau}\left(\boldsymbol{x}^{(t)}-\boldsymbol{x}_{i}^{(t, \tau)}\right), & \text { if } i \in \mathcal{S}^{(t)} \\ \boldsymbol{c}_{i}^{(t)}, & \text { otherwise } \end{array}\right.\qquad\) 其中上标$(t)$ 表示通信轮次的索引。

值得注意的是,SCAFFOLD要求客户端具有跨轮的持久状态。为了克服这个限制,Mime 算法 [3] 探索了另一种选择,即使用在全局模型上评估的随机梯度\(\nabla F_i(x;\xi _i)\)作为局部变量\(c_i\),同步的full-batch梯度\(\frac{1}{ \mathcal{S} } \sum _{i\in\mathcal{S}}\nabla F_i(x)\)作为全局控制变量 c。通过这样做,Mime 还可以减少局部 mini-batch 梯度的方差,并且适用于跨设备的情况

本地后验采样

另一种消除客户端更新偏差的方法是计算客户端deltas,如下所示: \(\hat{\Delta}_{i}^{(t)}=\hat{\Sigma}_{i}^{-1}\left(\hat{\mu}_{i}-\boldsymbol{x}^{(t)}\right) \qquad(5)\) 其中 $\hat{\mu}{i}$ 和 $\hat{\Sigma}{i}$ 是第 $i$ 个客户端上局部后验分布的均值和协方差的估计值。随着用于估计 $\hat{\mu}{i}$ 和 $\hat{\Sigma}{i}$ 的后验样本的增加,$\hat{\Delta}_{i}^{(t)}$ 中的偏差消失。换句话说,当客户每轮可以进行一些额外的计算时,可以运行局部马尔可夫链蒙特卡洛 (MCMC) 来生成近似的局部马尔可夫链蒙特卡洛 (MCMC)后验样本并使用它们来减少偏差,而不是运行许多步骤或epoch的 SGD(这会导致异构设置中的偏差增量)。在统一先验假设下,全局后验可以通过局部后验来估计。

2.3 正则化本地优化目标

在 vanilla FedAvg 和 FedOpt 中,客户端在每一轮执行多个本地更新。虽然更多的局部步骤将减少每次迭代的平均通信延迟,但由于本地目标的异构性,它可能会在训练结束时产生额外的误差。如果让客户端执行过多的局部更新步骤,那么局部模型将直接收敛到局部目标的最小值,这可能与全局目标相距甚远。为了避免局部模型向局部最小值漂移,一个自然的想法是通过正则化局部目标来惩罚远离全局模型的局部模型。形式上,可以在第 t 轮增加局部目标函数,如下所示: \(\widetilde{F}_{i}^{(t)}(\boldsymbol{x})=F_{i}(\boldsymbol{x})+\psi_{i}\left(\boldsymbol{x}, \boldsymbol{x}^{(t)}\right)\qquad(6)\) 其中 $x^{(t)}$ 是全局模型,也是第 t 轮局部训练的起点,函数 $ψ_i$ 表示局部模型 x 与当前全局模型 $x^{(t)}$ 之间的一种距离度量。通过仔细选择正则化器 $ψ_i$ ,可以确保局部模型接近全局模型,或者替代局部目标 (6) 具有与原始全局目标 (1) 相同的固定点。一个关于 $ψ_i$的例子是FedProx算法[4]

FedProx算法使用局部模型与全局模型之间的欧氏距离作为正则化函数。也就是说,$\psi_{i}\left(\boldsymbol{x}, \boldsymbol{x}^{(t)}\right)=\frac{\mu}{2}\left|\boldsymbol{x}-\boldsymbol{x}^{(t)}\right|^{2}$ 其中 $\mu \geq 0$ 是可调超参数。这种简单的正则化形式特别适用于 FL 的cross- device场景,与普通 FedAvg 相比不需要任何额外的计算或通信。但理论上,它无法提供比普通 FedAvg 更好的收敛速度。

2.4 使用替代的聚合函数

如前所述,全局目标 $F (x)$ 定义了客户端目标 $F_{i} (x)$ 的权重。通常,客户端的本地更新可以按样本数量加权或在每一轮统一加权。尽管加权方案可能不会影响 FL 算法的收敛速度,但它们确实控制了全局模型最终收敛到的位置。通过调整权重方案,最终的全局模型可能会偏向某些客户端的最小值。下面说明了加权策略对于学习的几点影响:

目标不一致问题

当所有客户端在每一轮都参与训练并且具有同质的计算速度时,服务器只需要以与全局目标相同的方式聚合所有局部变化。然而,当存在跨设备 FL 设置中的客户端采样机制或客户端每轮采取不同的本地步骤时,加权方案与许多实现选择密切相关。如果加权方案选择不当,则全局模型将收敛到替代目标的固定点,这与原始目标不一致。

为了保证收敛,应根据客户端抽样方案选择加权方案。具体来说,如果根据客户的本地样本数量大小对客户进行放回抽样,则应该对本地模型变化进行统一平均;如果客户是统一抽样而不放回,那么每个本地模型的权重应该根据他们的样本大小重新加权。关键是我们需要确保以与全局目标相同的方式对每个本地客户的聚合本地变化的期望进行加权。为了简化理论分析,研究人员通常假设加权抽样具有替换和均匀平均。然而,在现实世界的跨设备 FL 设置中,服务器对总体的了解有限,对客户端采样的控制也有限,因此聚合中的权重最终实际上是一个time-varying函数,具体取决于活动客户端集 $S^{(t)}$ 。

最近的工作 [5] 进一步表明,加权方案也受到客户端计算异质性的影响。如果我们仍然使用原来的加权方案,那么当客户端在 vanilla FedAvg 和许多其他 FL 算法中执行不同数量的本地步骤时,该算法将隐式地为具有更多本地步骤的客户端分配更高的权重。这项工作还给出了原始 FedAvg 实际优化的代理目标函数的解析表达式,作者提出了一种简单的方法 FedNova,通过将局部更新与局部步数“归一化”(或除以)来消除代理损失与原始损失之间的不一致。这种方法也可以被认为是一种新的加权方案,因为它等效地为具有更多本地步骤的客户端分配较低的权重

虽然客观不一致问题会极大地损害 FL 算法的性能,但影响可能会因数据集和神经网络模型的属性而异,具体取决于实际训练任务。由于神经网络的非凸性,替代目标的最小值可能也是原始目标之一。

加权平均之外的策略——神经元匹配算法

在 FedAvg 中,局部模型的参数按坐标平均,权重与客户端数据集的大小成比例 。 FedAvg 的一个潜在缺点是权重的坐标平均可能会对平均模型的性能产生严重的不利影响,并显着增加通信负担。这个问题的出现是由于神经网络中隐藏层的置换不变性,即对于大多数神经网络,只需排列每一层内的神经元,就可以形成许多功能等效的网络(相同输入映射到相同的输出)。但是将坐标平均应用于此类功能等效的网络可能会导致网络产生截然不同的输出值。

最近的几种技术通过在对客户端神经网络进行平均之前匹配客户端神经网络的神经元来解决这个问题。例如,概率联合神经匹配 (PFNM)[6] 利用贝叶斯非参数方法根据数据的异质性调整全局模型大小。在许多情况下,PFNM 比 FedAvg 具有更好的经验性能和通信效率。但是该方法的结构限制将其限制在相对简单的神经体系结构中,例如深度有限的全连接神经网络

Reference

[1]Sashank J. Reddi, Zachary Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Koneˇcn ́ y, Sanjiv Kumar, and Hugh Brendan McMahan. Adaptive federated optimization. In International Conference on Learning Representations, 2021. URL https://openreview.net/ forum?id=LkFG3lB13U5.

[2]Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank J Reddi, Sebastian U Stich, and Ananda Theertha Suresh. SCAFFOLD: Stochastic controlled averaging for on-device federated learning. International Conference on Machine Learning (ICML), 2020.

[3]Sai Praneeth Karimireddy, Martin Jaggi, Satyen Kale, Mehryar Mohri, Sashank J Reddi, Sebastian U Stich, and Ananda Theertha Suresh. Mime: Mimicking centralized stochastic algorithms in federated learning. arXiv preprint arXiv:2008.03606, 2020.

[4]Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization in heterogeneous networks. Proceedings of Machine Learning and Systems, 2020.

[5]Jianyu Wang, Qinghua Liu, Hao Liang, Gauri Joshi, and H Vincent Poor. Tackling the objective inconsistency problem in heterogeneous federated optimization. In Advances in Neural Information Processing Systems (NeurIPS), 2020.

[6]Mikhail Yurochkin, Mayank Agarwal, Soumya Ghosh, Kristjan Greenewald, Nghia Hoang, and Yasaman Khazaeni. Bayesian nonparametric federated learning of neural networks. In International Conference on Machine Learning, pages 7252–7261. PMLR, 2019.