对于联邦学习，有一个非常重要的假设：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 连续性假设#^2f66a6（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*} $$
剩下的就好办了。