如何解决联邦学习中的通信开销问题？(下)

1.3 Count Sketch 压缩处理 [3]

这篇文章中提出的 FetchSGD 使用 Count Sketch 来压缩模型更新，然后利用 Sketches 的可合并性来将不同客户端的模型更新进行合并。FetchSGD 设计中的一个关键问题是，由于 Count Sketch 是线性的，动量和误差累积都可以在 Count Sketch 中进行。这使得该方法能够将动量和误差累积从客户端转移到中央服务器中，从而在克服稀疏客户端参与挑战的同时，确保高压缩率和良好的收敛性。FetchSGD 的完整方法见图 8：（1）在客户端本地计算梯度，（2）将梯度的 sketches 发送到中央服务器中，中央服务器聚合梯度（3）sketches（4）动量和误差累积（5），（6）提取近似的 top-k 值，（7）中央服务器将稀疏值更新到参与下一轮训练的客户端设备中。

图 8. FetchSGD 完整方法图示。

在 FetchSGD 的每次迭代中，第 i 个参与的客户端使用部分（或全部）本地数据计算随机梯度 (g_i)^t，然后使用称为 Count Sketch 的数据结构压缩(g_i)^t。每个客户端将 Sketch S((g_i)^t) 作为其模型更新发送到聚合器（中央服务器）。Count Sketch 是一种随机的数据结构，它可以通过将向量多次随机投影到低维空间来压缩向量，以便后续近似地恢复高幅值（High-Magnitude Elements）数据。如下式：

由于 Count Sketch 具备线性特性，中央服务器可以在给定 S((g_i)^t)的情况下，计算出每个小批量梯度 g^t 的 sketch：

此外，Count Sketch 的另一个有用特性是，对于 sketching 操作符 S()，有一个对应的解压操作符 U()返回原始向量的无偏估计，从而实现对向量高幅值元素的近似：

简而言之，U()近似地「撤消」由 S() 计算的投影，然后使用这些重建来估计原始向量。基于 Sketch S((g_i)^t)，中央服务器利用下式更新全局模型：

Top-k(g^t) 不是 g^t 的无偏估计，因此不具备经典 SGD 的聚合性。不过，研究人员也证明了如果有偏梯度压缩方法累积了有偏梯度压缩算子所产生的误差，并在以后的优化中重新引入误差，则仍然可以收敛。在 FetchSGD 中，由 Top-k 引入的 bias 可以在中央服务器端进行误差累积，从而将其重新零初始化（zero-initialized）为 sketch S_e：

相比之下，其他有偏梯度压缩方法在压缩梯度时会给客户端带来偏差，因此客户端本身必须保持单独的误差累积向量。这在联邦学习中是很难保证的，这是由于客户端仅能参与一次更新，这样就没有机会在下一轮中重新引入错误。从另一个角度看，由于 S()是线性的，并且误差累积只包含线性操作，因此在 S_e 的中央服务器上进行误差累积相当于在每个客户端上进行误差累积，并将结果 Sketch 上传到中央服务器。更进一步，我们注意到动量也只包含线性操作，因此动量可以等价地在客户端或中央服务器上执行。推广上述方程可以得到：

完整的 FetchSGD 计算流程如下：

本文实验主要在小型本地数据集和 Non-IID 数据上完成，因为作者认为这是联邦学习中一个重要且相对未解决的问题。经典的梯度稀疏化方法是将每个客户端的局部 top-k 梯度元素聚合在一起训练全局模型的，当各个客户端的本地数据集非常小或相互之间差异非常大时，这种方法在近似全局梯度的真正 top-k 梯度元素时的表现就会非常差。在这种情况下，与经典方法相比 FetchSGD 有一个关键的优势：FetchSGD 的压缩算子是线性的。在 FetchSGD 中「只使用具有 N 个数据的单个客户端执行一个步骤相当于使用 N 个客户端执行一个步骤」，因此，每个客户端只贡献一次数据，小型客户端本地数据集不会带来任何问题。独立同分布的问题也可以通过随机选择客户端得到解决，FetchSGD 将参与训练的各个客户端的数据聚合起来，因此可以得到更完整的数据分布样本。

实验中，作者以上传和下载的总字节数表征相对于未压缩的 SGD 所实现的压缩效果。这些数据中没有考虑到的一个重要因素是，在 FedAvg 中，客户端必须在参与之前立即下载一个完整的模型，因为每个模型的权重在每一轮中都会得到更新。相比之下，局部 Top-k 和 FetchSGD 每轮只更新有限数量的参数，因此未参与的客户端可以相对更新当前模型，从而减少了必须在参与之前立即下载的新参数的数量。这使得本地 Top-k 和 FetchSGD 的上传压缩比下载压缩更重要。下载压缩对于这三种方法（FedAvg、Top-k、FetchSGD）也不那么重要，因为目前边缘设备的互联网连接的下载速度往往远远高于上传速度。作者给出了整体压缩（包括上传和下载）的结果，在图 9 中把这些图分成单独的上传和下载部分进行展示。FetchSGD 在不同的数据集和上传、下载和整体压缩的任务中，表现都较优。

图 9. CIFAR10（左）和 CIFAR100（右）的上传（顶部）、下载（中间）和整体（底部）压缩效果比对。为了提高可读性，每个图只显示该图中显示的压缩类型的运行的帕累托边界。

1.1.4 周期平均和量化的压缩处理方法 [6]

这篇文章介绍了一种周期平均和量化处理的联邦学习方法（a Communication-efficient Federated learning algorithm with Periodic Averaging and Quantization, FedPAQ）。其中的量化处理也可以看作是一种压缩方式。

在经典联邦学习框架中，为了利用客户端节点上所有可用的数据样本，参与训练的客户端在每次训练迭代中通过中央服务器同步其模型，因此，客户端和中央服务器之间要进行多次通信，从而导致网络上的通信争用造成较大通信开销。本文提出的周期平均和量化处理方法令客户端进行本地更新并定期通过中央服务器进行同步。一旦某客户端节点从中央服务器中获取到更新的模型，客户端节点每隔 τ 次本地 SGD 就向中央服务器发送更新信息以更新聚合模型。这种周期平均方案减少了中央服务器和客户端之间的通信次数，从而降低了训练模型的总体通信成本。如果在客户端中运行 T 次 SGD 迭代，则客户端需要与中央服务器进行 K=T/τ 轮通信，从而将总通信成本降低为原成本的 1/τ。这就是「周期平均」的处理思路。

从 K=T/τ 的计算公式可以直观看出，选择较大的 τ 值可以减少固定迭代次数 T 的通信次数，进而降低通信成本。但是这种降低是以牺牲模型准确度为代价的。增大 τ，会增加系统的噪声，进而客户端中的局部模型会逐渐收敛到局部最优解，而不是全局最优解。因此，作者考虑运行更多次迭代 T 来使模型达到特定的准确度。事实上，我们需要解决的一个关键问题是找到最优 τ，以使整个过程通信成本最小化。

在联邦学习网络中，通常有大量的设备（如智能电话）与中央服务器（基站）进行通信。但是，基站本身的下载带宽有限，因此只有少数设备能够同时将其消息上载到基站。由于这一限制，从客户端设备发送的消息将在中央服务器基站中进行流水线式的串行传输，这导致训练速度大大减慢。另一方面，如果让所有的客户端设备都参与到整个训练过程中，将会造成巨大的、昂贵的网络通信开销。此外，在实际应用中，并不是所有的客户端都在每一轮训练的过程中发挥作用的。有很多因素决定客户端是否参与当前的训练过程：设备需在当前状态下处于中央服务器基站通信可达的范围内、客户端设备在当前状态下为空闲可用状态、客户端设备通电且联网等等。

考虑到上述因素，FedPAQ 假设，在总共 n 个客户端设备中每轮训练中只有 r 个节点（r ≤ n）可用，且这 r 个可用设备在网络上随机且均匀地分布。在第 k 个训练周期内，中央服务器将其当前模型 x_k 发送给本轮选定的参与训练的 S_k 个客户端子集中的 r 个客户端节点，r 个客户端节点在该子集的总共 n 个客户端节点之间随机均匀分布。另一方面，联邦学习中的设备上行链路带宽有限，这使得从客户端到中央服务器的通信缓慢且昂贵，这也是前面几节中各种压缩方法所考虑的主要问题。本文所提出的方法是在传输信息中使用量化算子，通过交换量化更新来降低网络通信开销。

在进行 τ 轮本地 SGD 更新后，每个客户端 i∈S_k 中拥有本地模型 (x_k,τ)^i，其中 x_k 为最新从中央服务器中获取的全局模型。每个客户端对 x_k 和(x_k,τ)^i 的差值应用量化算子 Q()，并将量化结果 Q((x_k,τ)^i-x_k) 发送至中央服务器。中央服务器接收到量化结果后进行反量化解码处理，并基于处理结果生成新的全局模型 x_k+1。本文使用的量化算子为：

其中，ξ_i(x,s) 为以概率 |x_i|/||x||)s-l 确定值为 (l+1)/s 的随机变量。控制参数 s 表征量化水平，l∈[0,s) 为满足下式的整数：

完整的 FedPAQ 方法如下：

总的来说，FedPAQ 通过使用三个模块来降低通信负载：周期平均、部分客户端参与和量化处理。然而，这种通信减少带来了收敛准确度降低的问题，因此需要更多次的训练迭代。作者在原文中还进行了 FedPAQ 的收敛性分析，并给出了 FedPAQ 强凸和非凸损失函数的近似最优理论证明。我们在这里不再详述。

作者在实验中对比了通信开销和收敛性的 tradeoff 的结果分析。实验以总的训练时间为代价目标，包括通讯时间和计算时间。首先，定义网络带宽（Bandwidth，BW），每轮的通讯时间为上传的总比特数除以 BW。每轮的总比特数计算为 r ・ | Q(p,s)|，其中 | Q(p,s)| 表示根据具有 s 级（s levels）的特定量化算子对 p 维向量进行量化编码所需的比特数。在本文实验中使用的量化算子中，假定它需要 pF 位来表示长度为 p 的未量化向量，其中 F 为 32 位。然后，利用梯度计算时间的位移指数模型（Shifted-Exponential Model）确定计算时间。假设对于任何客户端，计算一个周期内的 τ 次迭代和批量大小为 B 的梯度需要确定位移 τ ・ B ・ scale^(-1) ，其中 shift 和 scale 分别是位移指数模型的位移和尺度参数，实验中 B 确定为 10。每轮的总计算时间就是 r 个贡献客户端节点中最大的本地计算时间。最后，计算通信 - 计算比为：

在图 10 中，前四个图展示了在 MNIST 数据集（'0'和'8'位）上，T=100 次迭代的正则化逻辑回归问题的训练时间。联邦学习网络中共有 n=50 个客户端节点，每个节点加载 200 个样本。设置 C_comm/C_comp=100/1 来捕获通信瓶颈。第一列图中曲线显示了在中央服务器上每轮的训练时间与所获得的训练损失之间的关系。第二列图中曲线显示了参与更新的客户端数量 n 的影响。第三列论证了周期长度 τ 在通信 - 计算 tradeoff 中的作用。最后一列图比较了 FedPAQ 与其他两个基线方法 FedAvg 和 QSGD 的训练时间。后四个图为 CIFAR-10 数据库中神经网络的实验结果。具体图例与 MNIST 中结果相同。

图 10. 训练损失与训练时间：MNIST 的 Logistic 回归分析（上），CIFAR-10 的神经网络结果（下）。

1.2 其它处理方法

1.2.1 集成方法 [4]

针对联邦学习的通信开销问题，这篇文章提出了利用集成方法（Ensemble method）的思路。集成方法是机器学习中的一种通用技术，用于组合多个基本预测因子（Base predictors）或专家（Experts）来创建一个更精确的模型。作者认为，联邦学习中的通信开销问题是由每轮从中央服务器发送到客户端（下行）和从客户端发送到中央服务器（上行）的参数数量引起的。在每轮训练过程中，中央服务器将当前模型的迭代状态发送给全部参与的客户端，直接将集成方法应用于这种联邦学习框架中会由于每轮都需要传递预测值而导致通信爆炸。本文提出的 FedBoost 能够在降低通信成本的同时实现计算加速、收敛保证和隐私保护。这种方法可以通过联邦学习使用客户端设备上的数据来训练一个原本可能超过客户端的通信带宽和存储容量的模型。此外，FedBoost 能够同时降低中央服务器到客户端（下行）和客户端到中央服务器（上行）的通信成本。

集成方法通过联邦学习的框架在中央服务器端只需要学习混合权重，所需要经由客户端发送给中央服务器的数据量非常小。因此，作者认为在集成方法中客户端到中央服务器（上行）的通信成本可以忽略不计。本文提出了标准（Standard）和任务不可知（Agnostic）的联邦学习集成方法，以解决中央服务器到客户端（下行）的通信瓶颈问题。

首先介绍标准联邦学习集成方法。给定一组预先训练的基本预测因子或假设：H≜{h_1, ..., h_q}。在标准集成方法中，将全部的假设都发送给每个参与的客户端。然而，在实践中，由于中央服务器和客户端之间的通信带宽以及客户端的内存和计算能力的限制，这种处理方式是不可行的。作者提出了一种抽样方法，只将其中一部分假设发送给客户端。这虽然可以降低通信复杂度，但同时会带来整体梯度偏差的问题以及集成收敛性的不确定性问题。经典的集成方法为：

其中，D_k 表示第 k 个域分布，h 为预先确定的假设，B_F 表示域分布 D_λ与假设估计 h 之间的 Bregman 散度。

给定确定的基本预测因子，引入一个集成权重 α。令 L_k(α)表示经验损失函数，L_D(k)(h_α) 表示域 k 中 m_k 个样本的集成：

其中的集成权重 h_α，定义为：

设 C 为每轮发送给客户端的最大基本预测因子数，即 C 能够表征通信效率。中央服务器端的目标函数是学习一组针对预先训练的基本估计量 h_k 的系数 α：

FedBoost 的完整算法流程如下：

在每轮训练中，FedBoost 在中央服务器上抽取两个子集：一个预训练假设子集，其中每一个子集以概率 γ_k,t 抽取得到（用 H_t 表示）；N 个客户端的随机子集（用 S_t 表示）。定义以下 Bernoulli 指标：

在随机采样的条件下，t 时刻的集成为：

针对该集成的偏置估计为：

通过除以概率γ_k,t 的方式来解决这一问题，从而得出对集成的无偏估计：

作者提出了两种计算概率 γ_k,t 的方法。一是均匀抽样，γ_k,t=C/q。二是加权随机抽样，

然后，我们介绍任务不可知的联邦学习集成方法。作者将上述标准联邦学习环境中的集成方法推广到任务不可知的联邦学习中。在不可知损失函数中，优化问题涉及两组参数：集成权重 α 和混合权重 λ。不可知损失函数如下：

其中，L_k(α) 为标准联邦学习中域 k 的经验损失。针对 L(α,λ)的优化问题为一个两人博弈问题，找到最小化目标函数和对手的 α，同时使用λ最大化目标函数。最终目标是找到给定 α_opt 的极小极大博弈的均衡，它使得对于混合权重λ_opt 的损失最小化。l 为凸函数，可以使用一般镜像下降（generic mirror descent）或其他基于梯度的算法来优化解决这个问题。

作者提出了任务无关的 AFLBoost 方法优化上述目标函数。AFLBoost 可以看作是 FedBoost 和针对任务无关损失函数的随机镜像下降算法 [7] 的结合。AFLBoost 的详细算法流程如下：

作者在实验中证明了 FedBoost 在不同通信成本下对密度估计（Density estimation）任务的有效性。具体包括三种方法：（1）无通信效率处理（无采样）：γ_k,t=1。（2）均匀抽样：γ_k,t=C/q。（3）加权随机抽样：γ_k,t ∝ α_k,t·C。

作者首先创建了一个 p=100 的合成数据集，其中每个 h_k 是单个元素上的点质分布(Point-mass distribution)，初始化每个 α_k 为 1/p，混合权重λ遵循幂律分布(Power law distribution)。实验结果见图 11。由图 11 左，加权抽样的性能优于均匀抽样，两种方法的损失都在稳步下降。由图 11 中，在通信预算为 64 的情况下，不考虑通信效率均匀抽样和加权抽样的性能与 FedBoost 相同。此外，作者还在经典 Shakespeare 数据集中进行实验。如图 11 右，加权随机抽样比均匀抽样的表现更好，在这个库中加权抽样的表现甚至优于 FedBoost，具有更好的收敛性能。

图 11. 实验对比图。左：合成数据集的损失曲线比较；中：合成数据集中取样方法的比较；右：Shakespeare 联邦学习数据集的损失曲线比较。

1.2.2 分布式不动点优化方法 [5]

针对联邦学习的通信开销问题，一些研究人员的解决思路是利用客户端的本地计算。也就是说，在通信和模型聚合处理之前，在每个客户端设备中进行更多的本地计算，从而减少获得全局有意义的解决方案所需的通信总轮数。在这种思路下，研究人员集中考虑了一些局部梯度下降算法以改进本地计算的效果。本文就是这种思路的工作之一。但是本文的工作并不局限于通过梯度下降来最小化目标函数，作者引入了计算 M 个操作算子平均值的不动点的方法。实际上，大多数迭代方法都属于不动点方法（Fix-Point method），其目的是寻找某个算子的不动点。

为了从不动点方法角度进行分析，首先介绍经典的分布式不动点优化模型。令分布式系统中共包括 M 个并行计算节点（客户端）。每个节点中处理的变量可看作是欧氏空间中的向量。令Τ_i 表示欧氏空间的操作算子，平均算子为：

本文方法的核心是找到Τ的不动点，即找到向量 x* 满足 Τ(x*)=x*。可以通过在每个节点重复应用Τ_i，同时进行平均化处理以达到共识来最终获得目标解 x*。本文作者考虑，经过多次迭代后，每个客户端节点将其变量同步传递到远程中央服务器中。然后中央服务器计算所接收到的向量的平均值并将其广播到所有节点。

本文考虑两种不动点优化策略：第一种策略是对于每个客户端计算节点，迭代执行若干次某个操作序列（称之为局部步骤，local steps）。第二种策略是减少通信步骤的数量，即仅以一定的较低概率到中央服务器中共享信息，并且只在中间过程进行局部计算。

将 T 定义为欧氏空间中的一个操作算子，令 Fix(T) 表示 T 的不动点集合，对于欧氏空间中的每个 x 和 y，如果 T 满足下式，则称 T 是χ-Lipschitz 连续的：

此外，如果 T 是 1-Lipschitz 连续和χ- 收缩的，则称 T 是非扩张的。如果 T 是收缩的，则它具有存在且唯一的不动点。对于任意 α∈(0,1]，如果对于一些非扩张算子 T’ ，如果存在 T=αT’+Id，则称 T 是 α 平均的。如果 T 是 1/2 平均的，则称 T 是坚决不扩张的。

令(t_n)_n∈N 表示通信过程中的整数序列。在每次迭代过程中，操作算子Τ_i 应用于节点 i，并利用参数λ进行松弛。对于一定数量的迭代，M 个计算节点全部将其本地向量传输给中央服务器的主节点，主节点计算平均值并将平均值广播给全部节点。全部节点在新的一轮迭代开始时拥有相同的变量 (x^)^k。该算法是局部梯度下降法（Federated Averaging）的推广。

我们把一系列本地迭代称为一个 epoch，然后求取平均值。也就是说，第 n 个 epoch 是指数 k+1=t_(n-1)+1,...,t_n 的迭代序列。假设在两个聚合计算平均值步骤之间的每个 epoch 的迭代次数由某个整数 H 限定。则对于每个 n≥1，有 1≤t_n - t_(n-1) ≤H。具体，第一种策略的详细算法流程如下：

接下来，作者提出了第二种策略。第一种策略中的本地处理步骤可以看作是两个通信步骤之间的内环（Inner Loop），将内环用概率聚合（Probabilistic Aggregation）来代替，即可得到第二种策略。在以下意义上，它是通信 - 有效的：在第一种策略中，通信轮数除以 H（或是非均匀情况下 t_n - t_(n-1)的平均值），而在第二种策略中该值乘以概率 p≤1。第二种策略的详细算法流程如下：

在原文中，作者对两种策略算法的收敛性能进行了充分的论证，我们在这里不再详述。

作者选择经典分类问题的逻辑回归进行实验。相应的目标函数如下：

其中，a_i∈R^d，b_i∈{-1,+1} 为数据样本。使用 LIBSVM 库中的「a9a」和「a4a」数据集，并将 k 设为 L/n，其中 n 是数据集的大小，L 是ᐁf 第一部分的 Lipschitz 常数，且没有经过正则化处理。

本文实验中考虑梯度下降（Gradient Descent，GD）作为操作算子。也就是说，我们考虑最小化有限和问题：

其中，函数 f_i 为凸且 L 平滑的。令：

使用 1/L 作为步长，则每个Τ_i 都是坚决不扩张的。

使用第一种策略的算法 1 和使用第二种策略的算法 2 的实验结果见图 12 和图 13。参数 H 和λ的值越大，初始收敛速度越快，但邻域半径越大。就计算时间而言，该算法没有太大的优势，因为实验是在一台机器上进行的，通信时间可以忽略不计。但是在通信速度较慢的分布式环境中，本文的算法就有明显的优势。我们也可以观察到图中实验结果曲线没有出现振荡。因此，当只需要达到有限的准确度时，本文提出的本地方法具有明显的优势。

图 12. 具有梯度下降步长的算法 1 在均匀通信时间 t_n=nH 下的收敛性，（a）不同 H 值的通讯轮数，λ=0.5，（b）不同 H 值的计算时间，λ=0.5，（c）不同λ值的计算时间，H=4。

图 13. 梯度下降步长算法 2 的收敛性，λ=0.5，（a）梯度步长相同情况下，不同 p 值的通讯轮次数量，（b）梯度步长相同情况下，不同 p 值的计算时间，（c）梯度步长与 p 是成比例的，不同 p 值的通讯轮次数量。

2. 总结

我们在这篇文章重点关注了联邦学习框架中的通信开销研究进展。目前，大多数文章都从压缩的角度出发解决通信开销问题，这种方法的思路很直观：压缩后需要上行、下行传递的数据量就会减小，从而减轻通信开销。当然，压缩的方法有很多，例如有损压缩、提取 sketch、量化等等。此外，我们也分析了两篇非压缩思路的文章，作者分别使用了集成方法和加强本地计算的方法。在不同的文章中，作者对比和分析的实验指标各不相同，这说明目前还没有标准化、统一化、权威性的衡量联邦学习通信开销的指标，毕竟通信开销和计算效率是一对 tradeoff 的指标。单纯用通信时间或通信数据量去衡量方法的优劣并不客观。

目前，随着 5G 技术的发展，5G 网络中通信速率问题变得不再是问题。依托 5G，使用边缘设备的应用场景也越来越多，例如校园安全监控、明厨亮灶监控、移动执法等等。在这种情况下，是否能够缓解联邦学习中的通信开销问题，进而推动联邦学习更快的发展和应用？让我们拭目以待吧！

参考文献

[1] Jakub Koneˇcný, H Brendan McMahan, Felix X Yu, Peter Richtárik, Ananda Theertha Suresh, and Dave Bacon.

Federated learning: Strategies for improving communication efficiency. arXiv preprint arXiv:1610.05492, 2016b. https://arxiv.org/pdf/1610.05492.pdf.

[2] Caldas, S., Koneˇcny, J., McMahan, H. B., and Talwalkar, A. Expanding the reach of federated learning by reducing

client resource requirements. arXiv preprint arXiv:1812.07210, 2018.

[3] Rothchild, D. , Panda, A. , Ullah, E. , Ivkin, N. , Stoica, I. , & Braverman, V. , et al. (2020). FetchSGD: communication-efficient federated learning with sketching, ICML 2020, http://arxiv.org/abs/2007.07682v1.

[4] Jenny Hamer, Mehryar Mohri , Ananda Theertha Suresh, FedBoost: A Communication-Efficient Algorithm for Federated Learning,ICML 2020.

[5] Malinovsky, G. , Kovalev, D. , Gasanov, E. , Condat, L. , & Richtarik, P. . (2020). From local sgd to local fixed point methods for federated learning, ICML 2020, https://arxiv.org/abs/2004.01442.

[6] Reisizadeh A , Mokhtari A , Hassani H , et al. FedPAQ: A Communication-Efficient Federated Learning Method with Periodic Averaging and Quantization[J]. arXiv, 2019.https://arxiv.org/pdf/1909.13014.pdf.

[7] Mohri, M., Sivek, G., and Suresh, A. T. Agnostic federated learning. In International Conference on Machine Learning, pp. 4615–4625, 2019.