Tackling the Objective Inconsistency Problem in Heterogeneous Federated Optimization(解决异构联邦优化中的目标不一致问题)

1.摘要

在联邦学习中,客户端本地数据集和计算速度的异质性导致每个客户端在每次通信轮次中执行的本地更新次数存在较大差异。对这些模型进行简单的加权聚合会导致目标不一致,即全局模型收敛到一个与真实目标函数相差任意大的不匹配目标函数的稳定点。本文提供了一个通用框架来分析异构联邦优化算法的收敛性。该框架涵盖了之前提出的方法,如FedAvg和FedProx,并提供了对解偏倚和由于目标不一致导致的收敛减慢的第一个原则性理解。利用这一分析的见解,我们提出了FedNova,这是一种归一化平均方法,它消除了目标不一致性,同时保持了快速误差收敛。

2.介绍

联邦学习是分布式优化的一个新兴子领域,在该领域中,数据收集和模型训练被推送到大量具有有限通信和计算能力的边缘客户端。与传统分布式优化不同,传统分布式优化在每次本地梯度计算后通过中央服务器或对等通信进行一致性操作,在联邦学习中,每个通信轮次中选择的客户端子集会在这些模型聚合以更新全局模型之前执行多次本地更新。
联邦学习中本地更新次数的异质性。参与联邦学习的客户端通常在本地数据集的大小和计算速度方面表现出高度异质性。联邦学习的原始论文提出,每个客户端执行 $E$ 个周期(遍历其本地数据集)的本地更新随机梯度下降(SGD),每次使用的小批量大小为 $B$。因此,如果一个客户端有 $n$ 个本地数据样本,则本地第 $i$ 次 SGD 迭代次数为 $\tau = \frac{E n}{B}$,这在不同客户端之间可能差异很大。本地 SGD 迭代次数的异质性因客户端计算速度的相对变化而加剧。在给定的时钟时间内,较快的客户端可以比较慢的客户端执行更多的本地更新。由于不可预测的滞后或由后台进程、断电、内存限制等因素引起的减速,客户端在不同的通信轮次中进行的本地更新次数也可能有所不同。最后,客户端可能使用不同的学习率和本地求解器(而不是普通的)。这些方法(如 proximal gradient 方法或自适应学习率调度)可能导致每个客户端的模型进度存在异质性。
局部更新中的异质性导致目标不一致。大多数分析联邦优化算法收敛性的最新工作 [8-37] 假设所有客户端的本地更新次数相同(即,对于所有客户端 $i$,$\tau = \tau$)。这些工作表明,周期性地在本地训练的客户端模型之间达成一致可以达到全局目标函数 $F(x) = \frac{1}{n} \sum_{i=1}^{m} F_i(x) $ 的驻点,该函数是加权局部目标的和,权重由数据集大小 $n$ 决定。然而,当前没有任何分析对本地更新或联邦优化算法在实际情况下,当本地更新次数 $\tau$ 在不同客户端之间变化时的收敛性提供见解。事实上,正如我们在第 3 节中所展示的,异质性本地更新后标准的客户端模型平均化会导致收敛到一个驻点——不是原始目标函数 $F(x)$ 的驻点,而是一个不一致的目标函数 $F’(x)$,其与 $F(x)$ 的差异可能取决于 $\tau$ 的相对值。如果客户端 1 执行更多的本地更新,则更新后的 $x^{(t+1,0)}$ 会偏离真实的全局最小值 $x^*$,向局部最小值 $x^*_1$ 偏移。
广义分析框架的需求。一种朴素的方法来克服异质性是固定每个客户端在通信轮次中必须完成的本地更新次数 $\tau$,并让快速节点空闲等待慢速客户端完成其更新。这种方法将确保目标一致性(即代理目标 $F(x)$ 等于真实目标 $F(x)$),然而,等待最慢的客户端会显著增加总的训练时间。更复杂的方法,如 FedProx、VRLSGD和 SCAFFOLD,这些方法旨在处理非 IID 本地数据集,可以在一定程度上减少(而非消除)目标不一致性,但这些方法要么导致收敛速度变慢,要么需要额外的通信和内存。到目前为止,对于联邦学习中具有异质本地更新这一挑战性的设置,尚无严格理解目标不一致性和收敛速度。同时,如何最好地结合使用不同本地进度水平训练的模型也尚不清楚。
本文贡献。据我们所知,本工作首次提供了对解中偏差(由目标不一致引起)的基本理解,以及客户端本地进度异质性如何影响收敛速率。在第4节中,我们提出了一种通用的理论框架,该框架允许不同数量的本地更新和非IID本地数据集。
此外,还包括不同的本地求解器,如GD、SGD、带近端梯度的SGD、梯度跟踪、自适应学习率、动量等。它涵盖了现有的方法,如FedAvg和FedProx,并对它们的收敛行为提供了新的见解。在第5节中,我们提出了FedNova,这是一种在平均时正确加权本地模型的方法。它确保了目标一致性,同时保持快速的误差收敛,并且在第6节中显示其性能优于现有方法。FedNova可以与任何本地求解器和服务器优化器配合使用,因此与现有的方法是互补的。

3.系统模型与前期工作

联邦异构优化设置。在联邦学习中,总共m个客户端旨在共同解决以下优化问题:
$$
[
\min_{x \in \mathbb{R}^d} \left[ F(x) := \sum_{i=1}^m p_i F_i(x) \right]
]
$$
其中 $p = \frac{n_i}{n}$ 表示相对样本大小,$F(x) = \frac{1}{n_i} f(x; \xi)$ 是在第 $i$ 个客户端处的局部目标函数。这里,$f$ 是由学习模型定义的损失函数(通常假设为凸函数),$\xi$ 表示来自本地数据集的样本。在第 $t$ 次通信时,$D$ 轮中,每个客户端独立地从当前全局模型 ( x^{(t,0)} ) 开始运行 ( \tau ) 次本地求解器(例如,SGD)迭代,以优化其自身的局部目标。
在我们的理论框架中,我们将 ( \tau ) 视为一个可以随轮次变化的任意标量。实际上,如果客户端运行相同的本地周期数 ( E ),则有:
$$
[
\tau = \frac{E n}{B}
]
$$
其中 ( B ) 是小批量大小。或者,如果每个通信轮次在时钟时间上具有固定长度,则 ( \tau ) 表示客户端 ( i ) 在时间窗口内完成的本地迭代次数,并且可能在不同客户端之间(取决于它们的计算速度和可用性)以及不同的通信轮次之间发生变化。

FedAvg基准算法联邦平均(FedAvg)是第一个也是最常用的算法,在每个通信轮次结束时用于在中央服务器上聚合这些本地训练的模型。共享全局模型的更新方式如下:
$$
x^{(t+1,0)} - x^{(t,0)} = \sum_{i=1}^m p_i \Delta_i^{(t)} = -\sum_{i=1}^m p_i \cdot \eta \sum_{k=0}^{r_i - 1} g_i(x_i^{(t,k)})
\tag{2}
$$
其中,$x_i^{(t,k)}$ 表示在第 $t$ 次通信轮次中第 $k$ 次本地更新后的客户端 $i$ 的模型,$\Delta_i^{(t)} = x_i^{(t,r_i)} - x_i^{(t,0)}$ 表示客户端 $i$ 在第 $t$ 轮次中累计的本地进展。此外,$\eta$ 是客户端的学习率,$g_i$ 表示一个包含 $B$ 个样本的小批量的随机梯度。
当客户端数量 $m$ 较大时,中央服务器可能在每一轮中仅随机选择部分客户端进行计算。

FedProx通过添加一个近端项来改进 FedAvg。为了缓解由于非独立同分布数据和异构本地更新导致的一致性问题,提出在每个本地目标中添加一个近端项:
$$
[
\mu \left| x^{(t,0)} - x^{(t,i)} \right|^2
]
$$
其中 $( \mu )$ 是一个可调参数。这个近端项将每个本地模型向后拉得更接近全局模型 $( x^{(t,0)} )$。尽管研究表明 FedProx 改进了 FedAvg,但其收敛性分析受限于比之前的 FedAvg 分析更强的假设,并且仅适用于足够大的 $( \mu )$。
由于 FedProx 是我们一般框架的一个特例,我们的收敛性分析提供了对 $( \mu ) $影响的敏锐见解。我们表明,较大的 $( \mu ) $可以减轻(但不能消除)目标不一致性,尽管这会以较慢的收敛速度为代价。我们提出的 FedNova 方法可以通过保证一致性而不减慢收敛速度来改进 FedProx。
通过动量和跨客户端方差减少改进FedAvg。近期文献中通过在服务器端应用动量或使用跨客户端方差减少技术如VRLSGD和SCAFFOLD来提升FedAvg的性能。再次强调,这些工作没有考虑异构本地进度。我们提出的归一化平均方法FedNova与这些加速或方差减少技术正交,并且可以轻松结合。此外,FedNova也与梯度压缩/量化和公平聚合技术兼容并互补。

4.异构联邦优化的新理论框架

我们现在提出一个通用的理论框架,该框架涵盖了多种联邦优化算法,并有助于分析目标不一致性对其误差收敛的影响。虽然结果是针对全客户端参与设置提出的,但很容易将其扩展到每轮随机采样一部分客户端的情况。

4.1 异构联邦优化的广义更新规则

根据(2)中的回顾,联邦优化算法的更新规则可以表示为 $ x(t+1,0) - x(t,0) = \sum {i=1}^{m} p_i \Delta(it)$,其中$\Delta(it) := x(t,\tau i) - x(t,0)$表示第 $(i)$ 个客户端在第 $(t) $轮的局部参数变化,(p_i = \frac{n_i}{n}) 是客户端 (i) 的数据比例。我们将其更新规则重写为更一般的形式如下:
$$
[
x^{(t+1, 0)} - x^{(t, 0)} = -\tau
{\text{eff}} \sum
{i=1}^{m} w_i \cdot \eta d_i^{(t)} ]
$$
以下三个更新规则的关键元素在不同算法中表现出不同的形式:
局部平均梯度 $d(t)$:不失一般性,我们可以将累积的局部变化重写为
$$
\Delta(t) = \eta G(t)a
$$
其中
$$
G(t) = [ \text{sgn}(c x h(a, s, t)), g g(a, d x i(e, n, t)), \text{sgn}(t h(e, r, \tau u, n)) ]
$$
是一个非负向量,并定义了随机梯度如何在局部累积。然后,通过归一化梯度权重 $a$,局部平均梯度定义为
$$
d(t) = \frac{G(t)a}{| a |_1}
$$
归一化因子 $| a |1$ 是向量 $a$ 的 $(l_1)$ 范数。通过设置不同的 $a$,(4) 式适用于大多数常见的客户端优化器,如带有近端更新的 SGD、局部动量和可变学习率,更广泛地说,适用于任何累积变化
$$
\Delta(t) = \eta G(t)a
$$
即局部梯度的线性组合的求解器。
具体来说,如果客户端优化器是普通的 SGD(即 FedAvg 的情况),那么
$$
a = [1, 1, …, 1] \in \mathbb{R}^{\tau_i}
$$
新框架与 FedAvg 在规范化梯度方面的比较。因此,规范化梯度只是当前模型参数空间内所有随机梯度的简单平均值。实心黑色圆圈表示
$$
d(t) = \frac{G(t)a_i}{\tau_i} = \frac{\sum
{k=-01}^{g_i(x(t,k))}}{\tau_i}
$$
稍后,在本节中,我们将展示更多关于如何通过新的通用更新规则和 FedAvg 进行更新的具体示例。

聚合权重 $(w_i)$:每个客户端的局部平均梯度 (d_i) 在计算聚合梯度
$$
\sum_{i=1}^{m} w_i d_i
$$
时乘以权重 (w_i)。根据定义,这些权重满足
$$
\sum_{i=1}^{m} w_i = 1
$$
请注意,这些权重决定了全局目标
$$
F(x) = \sum_{i=1}^{m} w_i F_i(x)
$$
的逼近方向,FedAvg 隐式地为本地步数更多的客户端分配更高的权重,导致全局方向偏移——我们将在定理 1 中正式证明这一点。

有效步数 $(\tau)$:由于客户端 $(i)$ 执行了 $(\tau_i) $次本地更新,每次通信轮次的平均本地 SGD 步数为
$$
\bar{\tau} = \frac{\sum_{i=1}^{m} \tau_i}{m}
$$
然而,服务器可以通过设置参数 (\tau) 大于或小于 (\bar{\tau}) 来调整聚合更新的影响(类似于选择全局学习率 [25, 40])。我们称比值
$$
\frac{\bar{\tau}}{\tau}
$$
为减速比,并且它在第 4.2 节的收敛性分析中起着重要作用。

规则(4)使我们能够为给定的局部求解器 $(a)$ 自由选择 $(\tau)$ 和$ (w)$,这有助于设计如 FedNova 等快速且一致的算法,FedNova 是第 5 节中提出的一种归一化平均方法。在图 3 中,我们进一步说明了上述关键元素如何影响算法,并在模型参数空间中比较了新的广义更新规则和 FedAvg。
此外,在实现方面,服务器不需要知道局部累积向量 (a) 的具体形式。每个客户端可以将归一化的更新$\eta d(t)$发送到中央服务器,这只是累积局部变化 $\Delta(t)$的一个重新缩放版本。这种重新缩放确保了服务器能够有效地执行聚合,而不必深入了解每个客户端的具体更新机制。这种设计不仅提升了算法的灵活性,还增强了其适应不同类型客户端优化器的能力,使得在分布式学习环境中能够更好地处理异构数据和计算能力差异。
通过这种方式,FedNova 可以在一定程度上减轻因客户端间的非独立同分布(non-IID)数据引起的训练不均匀性,从而提高全局模型的收敛速度和准确性。我们在后续部分将详细探讨 FedNova 的具体实现与性能评估,及其在实际应用中的优势。
$$
x^{(t+1,0)} - x^{(t,0)} = \sum_{i=1}^{m} p_i \Delta_i^{(t)} = - \sum_{i=1}^{m} p_i |a_i|_1 \cdot \frac{\eta G_i^{(t)} a_i}{|a_i|1}
$$
$$
= -\left(\sum
{i=1}^{m} p_i |a_i|1 \right) \eta \sum{i=1}^{m} \left( \frac{p_i |a_i|1}{\sum{i=1}^{m} p_i |a_i|1} \cdot \frac{G_i^{(t)} a_i}{|a_i|1} \right).
$$
$$
\tau
{\text{eff}} : \text{effective local steps} \hspace{20pt} w_i : \text{weight} \hspace{20pt} d_i : \text{normalized gradient}
$$
不同于更一般的公式(4),公式(6)概括了以下先前的方法。在公式(6)中,(\tau) 和 (w
{\text{eff}, i}) 由局部求解器的选择(即 (a) 的选择)隐式固定。这意味着在特定的情况下,我们可以通过简单调整局部更新的步数和相应的权重来实现对全局模型更新的有效控制。
由于空间限制,以下示例的推导被移到了附录。

香草SGD作为本地求解器(FedAvg):在FedAvg中,本地求解器是SGD,使得$a_i =[1, 1, \ldots, 1] \in \mathbb{R}^{\tau_i} \text{ 并且 } |a_i|1 = \tau_i$。因此,局部平均梯度$d_i$是$\tau_i$次迭代的简单平均值,$\tau{\text{eff}} = \frac{\sum_{i=1}^n p_i \tau_i}{\sum_{i=1}^n p_i \tau_i}$。
并且$w_i = p_i \tau_i / \sum_{i=1}^n p_i \tau_i$。也就是说,具有更多本地步骤的标准化梯度将隐式地分配更高的权重。
近端SGD作为本地求解器(FedProx)在FedProx中,本地SGD步骤通过近端项校正。可以证明
$$
a_i = [(1 - \alpha) \tau_i^{-1}, (1 - \alpha) \tau_i^{-2}, \ldots, (1 - \alpha), 1] \in \mathbb{R}^{\tau_i},
$$
其中$\alpha = \eta \mu$,并且$\mu$是一个可调参数。在这种情况下,我们有$|a_i|1 = [1 - (1 - \alpha)^{\tau_i} ] / \alpha $
因此,$\tau
{\text{eff}} = \alpha^{-1} \sum_{i=1}^n p_i [1 - (1 - \alpha)^{\tau_i}]$,并且$w_i = p_i [1 - (1 - \alpha)^{\tau_i}] / \sum_{i=1}^n p_i [1 - (1 - \alpha)^{\tau_i}]$。
当$\alpha = 0$时,FedProx等价于FedAvg。随着$\alpha = \eta \mu$增加,FedProx中的$w_i$。更类似于$p_i$,从而使代理目标$\tilde{F}(x)$更加一致。然而,更大的$\alpha$对应于较小的$\tau_{\text{eff}}$,从而减慢收敛速度,正如我们在下一小节中讨论的那样。
带衰减学习率的SGD作为本地求解器假设客户的本地学习率呈指数衰减,则我们有$a_i = [1, \gamma_i, \ldots, \gamma_i^{\tau_i^{-1}}]$。其中$\gamma_i \ge 0$可以在客户间变化。因此我们有$|a_i|_1 = (1 - \gamma_i^{\tau_i}) / (1 - \gamma) \text{ 和 } w_i \approx p_i (1 - \gamma_i^{\tau_i}) / (1 - \gamma_i)$。与FedProx (7)的情况相比,改变$\gamma_i$的值有类似于改变$(1 - \alpha)$。
动量SGD作为本地求解器如果我们使用动量SGD,其中本地动量缓冲区在每轮开始时被重置为零 [25],由于跨装置FL [2]的无状态特性,则我们有$a_i = [1 - \rho^{\tau_i}, 1 - \rho^{\tau_i - 1}, \ldots, 1 - \rho] / (1 - \rho)$,其中$\rho$是动量因子,并且$|a_i|_1 = [\tau_i - \rho(1 - \rho^{\tau_i})] / (1 - \rho)] / (1 - \rho)$。更一般地,新公式(6)表明${w_i \ne p_i }$当客户具有不同的$|\alpha_i|_1$
时,这可能是由于不平衡本地更新(i.e., 客户有不同的维度),或各种本地学习率/动量时间表(i.e., 客户具有不同的时间尺度)。

4.2 光滑非凸函数的收敛性分析

在下面的定理1和定理2中,我们对一般更新规则(4)进行了收敛性分析,并量化了由于目标不一致导致的解偏差。该分析依赖于SGD标准分析中使用的假设1和假设2,以及联邦优化文献中常用的前提假设3,用以捕捉局部目标的差异。