FLMix: 联邦学习新范式——局部和全局的结合

文章链接：Federated Learning of a Mixture of Global and Local Models

发表期刊（会议）: ICLR 2021 Conference（机器学习顶会）

本博客从优化函数角度出发，学习传统联邦学习$◊$ 和 新型联邦学习 $♣$ 的差异

1. 背景介绍

菲利普和彼得两位学者在阿卜杜拉国王科技大学发表的一篇文章中，对于联邦学习（Federated Learning）和混合专家（MoE）的结合进行了早期的数理讨论。

有意思的是这两位学者的研究动机是为了保护自己的移动设备数据不外露的同时，还可以用这些数据进行机器学习。他们给了两个很简单的理由。

• First, many device users are increasingly sensitive to privacy concerns and prefer their data to never leave their devices.
• Second,moving data from their place of origin to a centralized location is very inefficient in terms of energy and time.
  
  一个理由是不安全，还有一个理由是不方便。

2. 传统联邦学习

目前为止，FL 已经成为一个跨学科领域，专注于通过直接在边缘设备上训练机器学习模型来解决问题。传统的FL框架，每个客户参与FL训练。

参数定义：训练客户数量 N；全局模型结构 $M_G$；全局模型参数 $\theta (d_1)维$
其中 $\theta \in {\mathbb{R}}^{d_1}$ and ${\mathbb{R}}^{d_1}\in \mathbb{R}$
FL的学习目标为：
$$◊\underset{\theta \in {\mathbb{R}}^{d_1}}{\min }F(\theta )=\frac{1}{N}\sum _{i=1}^N{f}_i(\theta )$$
对于每一个$f_i$，由于数据分布不同，假设第 $i$ 个客户的数据分布定义为 ${\mathcal{D}}_i$ 则：
$$f_i(\theta )={\mathbb{E}}_{(x,y)\sim {\mathcal{D}}_i}[f(x,\xi )]$$
其中$L_i(\cdot )$是客户 $i$ 的损失函数

求解 $F(\theta )$ 最流行的方法是FedAvg算法，在FedAvg最简单的形式中，即当不使用部分参与、模型压缩或随机近似时，FedAvg缩减为局部梯度下降(LGD)。这是GD在聚合之前对每个设备执行多个梯度步长的扩展。

FedAvg已被证明在经验上是有效的，特别是对于非凸问题（存在多个局部极小值的问题）。但在数据异质时，与非本地对应的算法相比，FedAvg的收敛保证较差。

FL 虽然已经有了诸多理论证明其可行性，但是它的最终结果是全局性的，我们需要思考，对于那些数据异构的个体而言，使用全局方案解决个体问题效用一定好吗？

答案是否定的，数据异构性不仅对设计新的训练方法来解决 $◊$ 提出了挑战，而且不可避免地对这种全局解决方案对个人用户的效用提出了质疑。事实上，在所有设备的所有数据中训练的全局模型可能会从个人用户体验的典型数据和使用模式中删除，以至于使其几乎无用。

---

3. FL新范式

本文提出了一种新的训练联邦学习模型的优化公式。标准FL旨在从存储在所有参与设备上的私人数据中找到一个单一的全局模型。相比之下，新方法寻求全局模型和局部模型之间的权衡，每个设备可以从自己的私有数据中学习而无需通信。

本文开发了有效的随机梯度下降(SGD)变体来求解新公式，并证明了通信复杂性的保证。该工作的主要贡献包括结合全局和局部模型的FL新范式、新范式的理论性质、无环路局部梯度下降(L2GD)、L2GD的收敛理论以及对局部步骤在联邦学习中的作用的见解。该文件还强调了本地SGD在通信复杂性和个性化联邦学习的好处方面优于传统SGD的潜力。

本文提出的训练监督联邦学习新范式如下：

♣ min ⁡ x 1 , . . . , x n ∈ R d { F ( x ) : = f ( x ) + λ ψ ( x ) } f ( x ) : = 1 n ∑ i = 1 n f i ( x i ) ψ ( x ) : = 1 2 n ∑ i = 1 n ∥ x i − x ‾ ∥ 2 \clubsuit \quad \min_{x_1,...,x_n \in \mathbb{R}^d } \{ F(x): = f(x)+ \lambda \psi (x)\} \\ f(x):=\frac{1}{n}\sum_{i=1}^{n} f_i(x_i) \\ \psi (x) := \frac{1}{2n}\sum_{i=1}^{n} \left \| x_i-\overline{x} \right \| ^2 ♣x1​,...,xn​∈Rdmin​{F(x):=f(x)+λψ(x)}f(x):=n1​i=1∑n​fi​(xi​)ψ(x):=2n1​i=1∑n​∥xi​−x∥2 其中 $\lambda \ge 0$ 是一个惩罚超参， $x_1,...,x_n\in {\mathbb{R}}^d$ 是本地模型参数 ， $x:=(x_1,x_2,...,x_n)\in {\mathbb{R}}^{nd}$ 并且 $\overline{x}:=\frac{1}{n}{\sum }_{i=1}^n{x}_i$ 是所有本地模型的平均值。

文章假设由 $f_i$ 得到的 $F$ 是一个强凸函数。 凸函数是二阶导始终为正（负）的函数，局部最小值即为全局最小值。对于 $◊$ 有一个唯一的解。这个解可以表示为：
$$x(\lambda ):=(x_1(\lambda ),...,x_n(\lambda ))\in {\mathbb{R}}^{nd}$$接着可以计算 $\overline{x}(\lambda ):=\frac{1}{n}{\sum }_{i=1}^n{x}_i(\lambda )$

---

理论逻辑

所提范式 $♣$ 的理论逻辑：

• Local models ($\lambda =0$) ：此时模型退化为局部模型，只需要将本地损失降到最低，即求解$$\underset{x_i\in {\mathbb{R}}^d}{\min }{f}_i(x_i)$$也就是说，$x_i(0)$ 仅基于存储在设备 $i$ 上的数据 $D_i$ 的局部模型。该模型可以由设备 $i$ 计算，而无需任何通信。通常情况下，$D_i$ 不够丰富，无法使用此本地模型。为了学习更好的模型，还必须考虑其他客户的数据。然而，这存在沟通成本。
• Mixed models ($\lambda \in (0,\infty )$)：随着 $\lambda$ 的增加，惩罚 $\lambda \psi (x)$ 的效果越来越明显，需要沟通以确保模型不会太不相似，否则惩罚 $\lambda \psi (x)$ 会增大。
• Global model ($\lambda =\infty$)：现在我们来看 $\lambda \to \infty$ 的极限情况。直观上，这种极限情况应该迫使最优局部模型之间是相同的，同时最小化损失 $f$，即让 $\psi (x)\to 0$ 。$$\psi (x):=\frac{1}{2n}\sum _{i=1}^n{\Vert x_i-\overline{x}\Vert }^2$$此时，这种情况有一个特殊的极限解： min ⁡ { f ( x ) : x 1 , . . . , x n ∈ R d , x 1 = ⋯ = x n } \min\{ f(x):x_1,...,x_n\in \mathbb{R}^d ,x_1=\cdots=x_n \} min{f(x):x1​,...,xn​∈Rd,x1​=⋯=xn​}。可以反证，如果 $\lambda =\infty$ 并且 $x_1=x_2=\cdots =x_n$不成立，那么 $F(x)=\infty$

---

重要假设

对于每一个设备 $i$ ,它的目标函数 $f_i:{\mathbb{R}}^d\to \mathbb{R}$ 是 $L-smooth$ 并且 $\mu -strongly$ 的凸函数。

• $L-smooth$：通常用来描述一个函数的平滑程度。一个函数被称为是 L-smooth 的，如果它的一阶导数（梯度）是 Lipschitz 连续的，即梯度的变化受到了一定的约束。
  如果存在一个常数 $L>0$，使得函数 $f$ 的梯度 $\nabla f(x)$ 对于任意的 $x$ 和 $y$ 满足以下不等式：$$\Vert \nabla f(x)-\nabla f(y)\Vert \le L\Vert x-y\Vert$$$\Vert \cdot \Vert$ 是向量的范数。这个定义表明函数的梯度变化受到了 $L$ 的限制，也就是说在函数曲面上相邻点处的梯度变化是有界的。
• $\mu -strongly$：描述函数的弯曲程度，指的是一个函数在某种程度上比一个凸函数更加强烈地弯曲。如果存在一个常数 $\mu >0$ ,它满足以下不等式：$$f(y)\ge f(x)+⟨\nabla f(x),y-x⟩+\frac{\mu }{2}\Vert y-x{\Vert }^2$$ $⟨\cdot ,\cdot ⟩$ 表示内积运算。这个不等式表明函数 $f$ 在任意点 $x$ 处的曲率至少为 $\mu$，即函数图像在局部区域内弯曲程度足够大。

$L-smooth$ 函数的特性使得在优化问题中的求解更为可行和稳定。因为具有 Lipschitz 连续梯度的函数对于梯度下降等优化算法而言，更容易收敛到局部最优解，避免了梯度变化剧烈导致的震荡或发散。确保收敛

$\mu -strongly$ 函数在局部区域内有一个严格的下界，这种特性使得优化算法能够更快速地收敛到全局最优解。加速收敛

---

解的特性

对于 $♣$ 的最优解，它应该具备以下三个特性：

我们将表征局部和全局的两个函数 $f(x(\lambda ))$ 和 $\psi (x(\lambda ))$ 视作关于变量 $\lambda$ 的函数。

• 特性一： $\psi (x(\lambda ))$ 是非递增的，对于 $\forall \lambda >0$ 有$$\psi (x(\lambda ))\le \frac{f(x(\infty ))-f(x(0))}{\lambda }$$进一步 $f(x(\lambda ))$ 是非递减的，所以 $f(x(\infty ))\ge f(x(\lambda ))$ 。
  
  上述式子表明：随着 $\lambda$ 的增大，惩罚项 $\psi (x(\lambda ))$ 会逐渐减少到 0 ，因此最优的局部模型 $x_i(\lambda )$ 会随着 $\lambda$ 的增长越来越相似。同时根据第二种表述，$f(x(\lambda ))$随 $\lambda$ 增加而增加，但不超过标准FL公式的最优全局损耗 $f(x(\infty ))$ 。
• 特性二：对于 $\forall \lambda >0$ and $1\le i\le n$ 我们可以得到如下最优局部解表示：$$x_i(\lambda )=\bar{x}(\lambda )-\frac{1}{\lambda }\nabla f_i(x_i(\lambda ))$$ 进一步还有$$\begin{gathered} \sum _{i=1}^n\nabla f_i(x_i(\lambda ))=0 \\ \psi (x(\lambda ))=\frac{1}{2{\lambda }^2}||\nabla f(x(\lambda ))|{|}^2 \end{gathered}$$从平均模型中减去局部梯度的倍数，可以得到最优局部模型。在最优状态下，局部梯度的总和总是为零。这对 $\lambda =\infty$ 显然是正确的，但这对 $\forall \lambda >0$ 都不太明显。
• 特性三：最优局部模型以 $O(1/\lambda )$ 的速度收敛于传统的FL解。
  令 $P(z):=\frac{1}{n}{\sum }_{i=1}^n{f}_i(z)$ ，此时 $x(\infty )$ 是 $P$ 的唯一最小值，可以得到：$$||\nabla P(\bar{x}(\lambda ))|{|}^2\le \frac{2L^2}{\lambda }(f(x(\infty ))-f(x(0)))$$

$♣$ 的解 $x(\lambda )$ 到纯局部解 $x(0)$ 和纯整体解 $x(\infty )$ 的距离是 $\lambda$ 的函数。

---