可能更好的阅读体验 Jeefy's Blog: FedProx
对于联邦学习，有一个非常重要的假设：No-IID 的数据分布。所以对于 FedAvg 其实不容易收敛（泛用性很差），直接平均可能震荡不收敛。所以提出了 FedProx，引入正则化参数，其基础还是 FedAvg
其中引入的两个放缩：
- \(\gamma\)-inexact solution：允许设备根据算力随时停止，解决算力不均衡问题
- \(B\)-local dissimilarity假设：量化数据异构程度
和部分工作分析收敛性的方法是相似的，例如 SCAFFOLD
基本定义：\(f\) 表示中心 LLM，\(F_k\) 表示边缘 SLM
在论文中，基础问题建模为：
\[\min_{\omega} f(\omega) = \sum p_k F_k(x) = {\mathbb E}_k [F_k(\omega)]
\]
其中 \(p_k\) 表示边缘设备的数据占比 \(\frac {n_k} n\)。
其核心就是加入了正则项的损失函数：
\[\min_\omega {\mathbb E} [F_k(\omega) + \frac \mu 2 ||\omega - \omega^*||^2]
\]
利用参数正则约束偏离程度。
Definition 1：同时，为了解决算力和资源不均衡的情况（有慢设备），其引入了 \(\gamma\text{-inexact solution}\) 这个概念：\(||\nabla h_k(\omega^*;\omega_t)|| \le \gamma^t_k ||\nabla h_k(\omega_t; \omega_t)|| \implies\) \(\omega^*\) 是一个 \(\gamma^t_k\) 不精确解。即当前梯度 \(\le \gamma \times\) 初始梯度即可。需要注意的是：\(h\) 表示
\[h_k(\omega; \omega_t) = F_k(\omega) + \frac \mu 2\| \omega - \omega_t \|^2
\]
这样可以有四点好处：
- 自然的将正则引入其中，可以让 Client 的更新自然的更靠近原始模型。
- 允许 Client 根据自身算力随时停止，减少任务超时而被放弃的情况。建模了不同设备不同状态时，完成任务的判断标准。
- 由于 \(\gamma\) 依据设备动态计算，所以这将很多具体情况抽象成了一个独立参数，使得该建模更具鲁棒性。
- 合理的设置 \(\mu\) 的大小，可以将 \(F_k\) 利用后文提到的二阶梯度下有界，将 \(h_k\) 转变为一个完全的凸函数（二阶梯度 \(>0\)）
Assumption 1：在分析阶段，引入了 \(B\)-local dissimilarity，用于衡量模型的相似程度。具体来说：
\[{\mathbb E}_k [||\nabla F_k(\omega)||^2] \le || \nabla f(\omega)||^2 B^2 \iff F_k \text{ is B-locally dissimilar at }\omega
\]
其中 \(\mathbb E\) 部分表示每个 Client SLM 的梯度，而 \(f\) 表示聚合后的梯度。对于 IID 的数据来说，\(\mathbb E\) 部分趋近于 \(0\)，对于 No-IID（越是 Heterogeneous ），\(B\) 取值就会越大。由于有限，所以可以定义 \(B(\omega)\) 表示这个放缩的界 \(\sqrt{\frac {{\mathbb E}_k [||\nabla F_k(\omega)||^2]}{|| \nabla f(\omega)||^2}}\)。也就是说 \(B(\omega)\) 可以衡量不相似的程度。如果 \(B \approx 1\)，说明 Client 数据分布在共有数据集上非常接近。
\(\oplus\) 假设：\(B(\omega)\) 是有界的，具体来说，如果说 \(||\nabla f(\omega)||^2 > \epsilon\)，那么总有 \(B(\omega) \le B_\epsilon\)
L-smooth Assumption：参考 L-smooth & Lipschitz 连续性假设 约束了 \(\| \nabla^2 F(x) \| \le L\) ，这里约束了二阶导的下界，也就是说 \(\nabla^2 F(x) \succeq -L_{\_} I\)，这里方便，将 \(L_\_\) 直接记为 \(L\) 。也就是说，这里约束整个模型虽然可以上凸，但是不能无限上凸，上凸是有界的。
记 \(\bar{\mu} = \mu - L > 0\)，也就是这里设置了“超参” \(\mu\)，将 \(h_k(x)\) 限制为了二阶梯度 \(\succeq \bar{\mu} I\)。
这里 \(\mu\) 的作用是抑制 Non-IID 数据带来的模型偏移，\(\mu\) 越高，抑制越强，但是收敛速度可能变慢。对于不同的数据集，其设计的值不同。
原文 Figure 3 展示了一种启发式策略：Loss 上升时增大 \(\mu\)，Loss 连续下降时减小 \(\mu\)。
证明收敛的流程大抵分为三步：
- 利用正则项控制凸性后，利用梯度估计误差和模型偏移
- 不考虑随机采样，估计理想下降
- 进行随机采样的偏差修正
现在，我们展开来证明模型收敛。首先参考 L-smooth & Lipschitz 连续性假设 明确根据泰勒展开，我们可以知道：
\[f(\omega^{t + 1}) \le f(\omega^t) + \langle \nabla f(\omega^t), \bar\omega^{t + 1} - \omega^t \rangle + \frac L 2 \| \bar\omega^{t + 1} - \omega^t\|^2
\]
于是我们的终极目的是通过神秘的放缩，将 \(f(\omega^t)\) 之外的项放缩为一个负数。
我们可以定义新一轮通信与上一轮的一个差距，但其实就是 \(\nabla_{\omega_k^{t + 1}} h_k(\omega_k^{t + 1}; \omega_k^t)\)：
\[\begin{aligned}
e^{t+1}_k = \nabla F_k(\omega_k^{t + 1}) + \mu (\omega_k^{t + 1} - \omega^t)
\end{aligned} {\tag 1}
\]
于是根据 \(\gamma\) 定义（这里 \(\omega_k^t - \omega_k^t\) 变成 \(0\) 了）：
\[\| e_k^{t + 1} \| \le \gamma_k^t \| \nabla h_k(\omega_t;\omega_t) \| = \gamma_k^t \| \nabla F_k(\omega_k^{t + 1}) \|
\]
定义平均聚合参数：\(\bar{\omega}^{t + 1} = {\mathbb E}_k[\omega_k^{t + 1}]\)，对于 \((1)\) 式两边取期望：
\[\bar{\omega}^{t + 1} - \omega^t = \frac 1 {\mu} {\mathbb E}[e_k^{t + 1}] - \frac 1 {\mu} {\mathbb E} [\nabla F_k(\omega_k^{t + 1}) ]
\]
对于新训练的这一轮，我们期望得到的东西是 \(\hat{\omega}_k^{t + 1} = \arg \min_\omega h_k(\omega; \omega^t)\)，而实际上得到的是不精确解 \(\omega_k^{t + 1}\)：
\[\| \nabla h_k(\omega_k^{t + 1}) \| \le \gamma_k^{t} \| \nabla F_k(\omega^t) \|
\]
根据 \(\bar{\mu}\) 强凸性，\(h_k\) 一定存在全局最小值，并且离最小值越远，\(\| \nabla h_k \|\) 越大。也就是 \(\hat{\omega}_k^{t + 1}\) 唯一确定，并且 \(\| \nabla h_k(\hat{\omega}_k^{t + 1}) \| = 0\)。对于强凸性，其存在性质：
\[\langle \nabla h_k(x) - \nabla h_k(y), x - y \rangle \ge \bar{\mu} ||x - y||^2
\]
那么带入 \(x = \hat{\omega}_k^{t + 1}, y = \omega_k^{t + 1}\)，那么有 \(\nabla h_k(x) = 0\)，得：
\[\bar{\mu} \| \hat{\omega}_k^{t + 1} - \omega_k^{t + 1} \|^2 \le \langle - \nabla h_k(y), x - y \rangle \le \| \nabla h_k(\omega_k^{t + 1}) \| \cdot \| \hat{\omega}_k^{t + 1} - \omega_k^{t + 1} \|
\]
就有了：
\[\| \hat{\omega}_k^{t + 1} - \omega_k^{t + 1} \| \le \frac 1 {\bar{\mu}} \| \nabla h_k(\omega_k^{t + 1}) \| \le \frac {\gamma_k^{t}}{\bar{\mu}} \| \nabla F_k(\omega^t) \|
\]
根据单纯的凸性：
\[\|\omega_k^{t} - \hat{\omega}_k^{t + 1} \| \le \frac 1 {\bar{\mu}} \| \nabla h_k(\omega_k^{t}) \| = \frac 1 {\bar{\mu}} \| \nabla F_k(\omega^t) \|
\]
于是利用三角不等式：
\[\| \omega_k^t - \omega_k^{t + 1} \| \le \|\omega_k^{t} - \hat{\omega}_k^{t + 1} \| + \| \hat{\omega}_k^{t + 1} - \omega_k^{t + 1} \| \le \frac {1 + \gamma_k^{t}}{\bar{\mu}} \| \nabla F_k(\omega^t) \|
\]
于是继续利用三角不等式：
\[\| \omega^t - \bar\omega^{t + 1} \| \le {\mathbb E}_k[\| \omega_k^t - \omega_k^{t + 1} \|] \le \frac {1 + \gamma^{t}}{\bar{\mu}} {\mathbb E}_k[ \| \nabla F_k(\omega^t) \| ]
\]
然后利用 \(B\) 以及 \(E[x^2] \ge E^2[x]\)，进一步得到：
\[\| \omega_k^t - \bar\omega_k^{t + 1} \| \le \frac {1 + \gamma^{t}}{\bar{\mu}} \sqrt{{\mathbb E}_k[ \| \nabla F_k(\omega^t) \|^2 ]} \le \frac {B(1 + \gamma)}{\bar\mu} \| \nabla f(\omega^t) \|
\]
现在，二阶项就可以放缩了，问题是一阶项，于是构造：
\[\bar\omega^{t + 1} - \omega^t = \frac {-1} \mu (\nabla f(\omega^t) + M_{t + 1})
\]
利用前面已经推导过的一些不等式和变量代替：
\[M^{t + 1} = {\mathbb E}_k [\nabla F_k(\omega^{t + 1}) - \nabla F_k(\omega^t) - e_k^{t + 1}]
\]
于是利用三角不等式放缩，然后利用 L-smooth 的性质：
\[\| M^{t + 1} \| \le {\mathbb E}_k [\|\nabla F_k(\omega^{t + 1}) - \nabla F_k(\omega^t) \|] + {\mathbb E} [ \| e_k^{t + 1}\| ] \le {\mathbb E} [L || \omega^{t + 1} - \omega^t \|] + {\mathbb E}[\| e_k^{t + 1}\|]
\]
然后就可以随便带入得：
\[\| M^{t + 1} \| \le \left(\frac {L(1 + \gamma)} {\bar\mu} + \gamma \right) B \| \nabla f(\omega^t) \|
\]
于是对于一阶项，利用柯西不等式放缩（这里需要注意正负性）：
\[\begin{aligned}
\langle \nabla f(\omega^t), \bar\omega^{t + 1} - \omega^t \rangle &= \langle \nabla f(\omega^t), \frac {-1} \mu (\nabla f(\omega^t) + M^{t + 1}) \rangle \\
&\le \frac {-1} \mu \| \nabla f(\omega^t)\|^2 + \frac {-1} \mu \langle \nabla f(\omega^t), M^{t + 1} \rangle
\\
&\le \frac {-1} \mu \| \nabla f(\omega^t)\|^2 + \frac {1} \mu \| \nabla f(\omega^t) \| \cdot \| M^{t + 1} \| \\
&\le \left(- \frac {1 + \gamma B} \mu + + \frac {L(1 + \gamma)}{\mu\bar\mu} B \right) \| \nabla f(\omega^t) \|^2
\end{aligned}
\]
于是将整体带入：
\[f(\bar\omega^{t + 1}) \le f(\omega^t) + (- \frac {1 + \gamma B} \mu + + \frac {L(1 + \gamma)}{\mu\bar\mu} B + \frac {LB^2(1 + \gamma)^2} {2\bar\mu^2}) \| \nabla f(\omega^t) \|^2
\]
于是如果说，让括号内的那部分小于 \(0\)，就可以期望 \(f(\bar\omega^{t + 1})\) 收敛了。
但是新的问题来了，这里 \(\bar\omega\) 是全体训练，但是实际上每一轮是随机采样训练，也就是只用了 \(K\) 个设备。所以我们使用局部 Lipschitz 连续性假设（因为这里可以认为 \(\omega\) 和 \(\bar\omega\) 差距不大），使用 \(L_0\) 表示其一阶阶梯度的界（其实是说，感觉不是这样用的）：
\[f(\omega^{t + 1}) \le f(\bar\omega^{t + 1}) + L_0 \| \omega^{t + 1} - \bar\omega^{t + 1} \| \tag 2
\]
现在考虑放缩一下 \(L_0\)，利用 L-smooth & Lipschitz 连续性假设（L-smooth 的梯度是 Lipschitz 连续的），这里先绕一绕：
对于限制在 \(\omega = \lambda \omega^{t + 1} + (1 - \lambda) \bar\omega^{t + 1}\) 上的常数 \(L_0\)，根据微分中值：
\[f(\omega^{t + 1}) - f(\bar\omega^{t + 1}) = \nabla f(\xi) \cdot (\omega^{t + 1} - \bar\omega^{t + 1}) \implies L_0 = \sup \|\nabla f(\xi) \|
\]
这里 \(L_0\) 其实取啥都可以，我们的目的只是给 \(L_0\) 一个可控的界，并不是严格的遵循局部 Lipschitz 连续性。
将 \(\nabla f(\xi)\) 通过 L-smooth 展开：
\[\| \nabla f(\xi) \| \le \| \nabla f(\omega^t) \| + L \| \xi - \omega^t\|
\]
注意到 \(\xi = \lambda \omega^{t + 1} + (1 - \lambda) \bar\omega^{t + 1} \implies \xi - \omega^t = \lambda (\omega^{t + 1} - \omega^t) + (1 - \lambda) (\bar\omega^{t + 1} - \omega^t)\)，所以 \(\| \xi - \omega^t\| \le \max (\| \omega_k^t - \omega_k^{t + 1} \| + \| \omega^t - \bar\omega^{t + 1} \|)\)
于是
\[\begin{aligned}
L_0 &\le \| \nabla f(\omega^t) \| + L \max (\| \omega_k^t - \omega_k^{t + 1} \| + \| \omega^t - \bar\omega^{t + 1} \|) \\ &\le \| \nabla f(\omega^t) \| + L (\| \omega_k^t - \omega_k^{t + 1} \| + \| \omega^t - \bar\omega^{t + 1} \|)
\end{aligned}
\]
那么接下来对于 \((2)\) 两边取期望：
\[\mathbb{E}_{S_t} [f(\omega^{t + 1})] \le f(\bar\omega^{t + 1}) + Q_t
\]
\[\begin{aligned}
Q_t &= \mathbb{E}_{S_t} [ L_0 \| \omega^{t + 1} - \bar\omega^{t + 1} \|] \\
& \le \mathbb{E}_{S_t} \Big[ {\Big (}L(\| \omega_k^t - \omega_k^{t + 1} \| + \| \omega^t - \bar\omega^{t + 1} \| ) +\| \nabla f(\omega^{t}) \| \Big) \| \omega^{t + 1} - \bar\omega^{t + 1} \| \Big]
\end{aligned}
\]
由于采样仅和 \(\omega^{t + 1}_k\) 有关，所以可以把无关的东西拉出来：
\[\begin{align*}
Q_t &\le \mathbb{E}_{S_t} \left[ \left( \|\nabla f(w^t)\| + L(\|\bar{w}^{t+1} - w^t\| + \|w^{t+1} - w^t\|) \right) \times \|w^{t+1} - \bar{w}^{t+1}\| \right] \\
&\le \left( \|\nabla f(w^t)\| + L\|\bar{w}^{t+1} - w^t\| \right) \mathbb{E}_{S_t} [\|w^{t+1} - \bar{w}^{t+1}\|] + L \mathbb{E}_{S_t} [\|w^{t+1} - w^t\| \cdot \|w^{t+1} - \bar{w}^{t+1}\|] \\
&\le \left( \|\nabla f(w^t)\| + 2L\|\bar{w}^{t+1} - w^t\| \right) \mathbb{E}_{S_t} [\|w^{t+1} - \bar{w}^{t+1}\|] + L \mathbb{E}_{S_t} [\|w^{t+1} - \bar{w}^{t+1}\|^2]
\end{align*}
\]
考虑到 \(Var(X) = E[X^2] - E^2[X] \ge 0\)，那么有下式：
\[\mathbb{E}_{S_t} [\|w^{t+1} - \bar{w}^{t+1}\|] \le \sqrt{\mathbb{E}_{S_t} [\|w^{t+1} - \bar{w}^{t+1}\|^2]}
\]
然后就是采样方差分解，也就是将偏差转化为整体方差的 \(\frac 1 K\) 倍。其背后的核心是，两者期望相同，看作独立采样 \(K\) 个设备，那么对于每个设备，有：
\[{\mathbb E}_{S_t}[ (x - \bar x)^2] = {\mathbb E}_{S_t} \left[\frac 1 {K^2} \sum (x_k - \bar x)^2 \right] = \frac 1 {K^2} \cdot K \cdot {\mathbb E}_k[(x_k - \bar x)^2]
\]
由于这里是不放回采样，所以需要乘以修正因子 \(\frac {N - K}{N - 1} \le 1\)，所以可以这样放缩：
\[\begin{align*}
\mathbb{E}_{S_t} [\|w^{t+1} - \bar{w}^{t+1}\|^2] &\le \frac{1}{K} \mathbb{E}_k [\|w_k^{t+1} - \bar{w}^{t+1}\|^2] \\
&\le \frac{2}{K} \mathbb{E}_k [\|w_k^{t+1} - w^t\|^2], \quad (\text{as } \bar{w}^{t+1} = \mathbb{E}_k [w_k^{t+1}]) \\
&\le \frac{2}{K} \frac{(1+\gamma)^2}{\bar{\mu}^2} \mathbb{E}_k [\|\nabla F_k(w^t)\|^2] \quad (\text{from } (6)) \\
&\le \frac{2B^2}{K} \frac{(1+\gamma)^2}{\bar{\mu}^2} \|\nabla f(w^t)\|^2,
\end{align*}
\]
剩下的就好办了。
浙公网安备 33010602011771号