Note_Training time minimization for federated edge learning with optimized gradient quantization and

Training time minimization for federated edge learning with optimized gradient quantization and
bandwidth allocation
Peixi LIU1,3, Jiamo JIANG†‡2, Guangxu ZHU†‡3,LeiCHENG4,5, Wei JIANG1,WuLUO1,Ying DU2, Zhiqin WANG2
Frontiers of Information Technology & Electronic Engineering
CCF C/SCI 3区 刊
2022

一、动机和贡献

动机：在 “同步 FL” 中，如果使用量化技术缓解通信压力，通常会有两个问题：1）若使用低精度量化尽可能减少传输数据量，会减少每轮的执行时间（由于通信时间的减少），但是却会导致量化误差增大，使得总轮数增加；2）若使用高精度量化尽可能保持低量化误差，会减少总轮数，但却相应的增加了通信时间。另一方面，由于client在 “计算能力”（CPU频率） 中有着较大差异，因此对所有client都使用相同的量化精度显然是不合理的。

贡献：这篇文章的主要解决问题是：如何对不同 “计算能力”的client分配相应的通信带宽并确定相应的量化精度，以达到使用最少通信轮数实现指定模型精度的目的。具体贡献如下：

• “训练时间” 的理论分析：难点是如何分析“总执行轮数”，这篇文章分析了给定量化精度下，模型达到指定优化精度所需要的总轮数的上界。但是由于理论结果太严格，因此这篇文章后续使用 “数据/模型驱动拟合法data-and-model-driven fitting method” 来估计总轮数。这样，就可以表示总训练时间；
• “训练时间” 的最小化：这篇文章是要同时优化 “带宽分配”和“量化精度选择”，但是这个优化问题是非凸的，不能直接解决。因此，将上述问题拆分成两个子问题：1）固定量化精度下，带宽的分配问题（使用 双重二分查找解决）；2）固定带宽下，量化精度的选择问题（使用 逐次凸逼近SCA 解决）。

注1：这篇文章将client“通信能力”也作为一个变量（通信带宽），是可以client进行指定的，而不是作为每个client固定的量，这与我们之前的想法是不同的。
注2：量化精度每个client在同一个时间是相同的，即量化精度只是关于时间的函数。

二、算法

2.1 预备知识

无线传输模型：
定义sever和client $k$ 之间的 信道传播系数channel propagation coefficient 为：
$$h_k=\sqrt{{\varphi }_k{\bar{h}}_k},$$其中，${\varphi }_k$ 表示大尺度传播因素（随时间变化保持不变的）；${\bar{h}}_k$ 表示小尺度传播因素，满足标准正太分布${\overline{h}}_k\sim \mathcal{C}\mathcal{N}(0,1)$（随时间变化而变化的）。另外，这篇文章假设 信道传播系数只被sever知晓哦啊。

因此，定义 client $k$ 的 遍历容量ergodic capacity 为：
$$R_k={\mathbb{E}}_{h_k}\left[b_k{\log }_2\left(1+\frac{p_k|h_k{|}^2}{b_k{N}_0}\right)\right]$$其中，$b_k$ 表示分配给client $k$ 的频率带宽，且满足 ${\sum }_{k=1}^K{b}_k=B_0$（这个就是需要优化的带宽）；$p_k$ 表示 client $k$ 的传输能量；$N_0$ 表示噪声。假设有 $S$ bits 数据需要传输，那么时间可以计算为：
$$T=\frac{S}{R_k}.$$

注：在不同client的传输能力方面，本文本质上假设只有带宽 $b_k$ 是可变且需要优化的，其他系数，如 $p_k,h_k,N_0$ 是相同的。

2.2 训练时间分析

假设 $T_d$ 表示一轮更新所需要的时间，那么：
$$T_k^{\mathrm{comp}}+T_k^{\mathrm{comm}}\le T_{\mathrm{d}}$$其中，$T_k^{\mathrm{comp}}$ 表示执行SGD的时间；$T_k^{\mathrm{comm}}$ 表示传输量化后模型的时间。假设 $N_{ϵ}$ 表示达到 $ϵ-$模型精度时所需要的最少轮数，那么总训练时间可以表示为：
$$T=N_{ϵ}{T}_d.$$

注：原文中认为 $T_k^{\mathrm{comp}}$ 表示SGD时间，但是这里应该也包括将全精度梯度量化所需要的时间。

定义 $T_k^{\mathrm{comp}}$：
$$T_k^{\mathrm{comp}}=\frac{\nu }{f_k},$$其中，$\nu$ 表示执行SGD的处理轮数；$f_k$ 表示client $k$ 的CPU主频。

定义 $T_k^{\mathrm{comm}}$：
假设 $S$ 表示量化后需要传输的总比特数，$q$ 表示随机量化精度，$d$ 表示梯度维数（$\mathbf{g}\in {\mathbb{R}}^d$），那么有如下关系：
$$S=(1+{\log }_2(q+1))d.$$因此，可以得到通信时间为：
$$T_k^{\mathrm{comm}}=\frac{S}{R_k}.$$

2.2.1 最少轮数 $N_{ϵ}$ 的估计

首先我们需要知道：$N_{ϵ}$ 的值只和量化精度 $q$ 有关，与其他量无关。

原文给出定理：
Theorem 1：在量化精度 $q\le 2$ 的条件下，当FL模型执行 $N$ 轮更新后，损失函数的精度误差为：
$$\mathbb{E}\left[F\left({\mathbf{w}}^{(N)}\right)\right]-F({\mathbf{w}}_{\ast })\le \frac{\alpha \kappa }{N+2\alpha \kappa -1}\left(L{\Vert {\mathbf{w}}^{(0)}-{\mathbf{w}}_{\ast }\Vert }^2+\frac{2\Gamma }{\mu }\right),$$其中，$\alpha =\frac{\sqrt{d}}{qK}+1,\kappa =\frac{L}{\mu },F_{\delta }=F({\mathbf{w}}_{\ast })-\frac{1}{K}{\sum }_{k=1}^K{F}_k^{\ast },F_k^{\ast }={\min }_{\mathbf{w}}{F}_k(\mathbf{w}),\Gamma =2L{F}_{\delta }+\frac{1}{K}{\sum }_{k=1}^K{\delta }_k^2$，且 ${\mathbf{w}}^{(0)}$ 表示初始权重。

虽然 Theorem 1 给出了理论上界，但是却有两个缺点：1）由于证明过程中进行太多放缩，导致上界是偏大的，也就意味着实际中可能陷入局部最优；2）一些实际的模型参数，例如$\mu ,L,F_{\delta },{\mathbf{w}}_{\ast },\Gamma$ 是不容易获取的。

因此这篇文章采用了 “数据/模型驱动法” 来估计 $N_{ϵ}$ 的值，核心思想是：进行一定程度的预训练，得到估计值。具体做法如下：

1. 首先根据 Theorem 1 定义 “上界函数” ：
   $$\hat{U}(N)=\frac{\alpha \kappa \left(L{\Vert {\mathbf{w}}_0-{\mathbf{w}}_{\ast }\Vert }^2+\frac{2\Gamma }{\mu }\right)}{N+2\alpha \kappa -1}.$$因为只关心 $N$与$q$（即$N$与$\alpha$）的关系，上式有两个特点：1）$\hat{U}(N)$ 是递减函数，且当 $N\to \infty$ 有 $\hat{U}(N)\to 0$；2）$\hat{U}(N)$ 总体是分数形式，且分子分母都随 $\alpha$ 做线性增加。因此重新定义 $\hat{U}(N)$ 为：
   $$\mathbb{E}\left[F({\mathbf{w}}^{(N)})\right]-F({\mathbf{w}}_{\ast })=\frac{\alpha A+D}{n+\alpha B+C}\triangleq U(N)$$其中，$A>0,B>0,C\ge 0,D\ge 0$。因此，只需要通过预训练估计得出 $A,B,C,D$ 的值，就可以得出 $N_{ϵ}$（即 $U(N)$）。

2. 接下来通过 “数据/模型驱动法”得出 $A,B,C,D$。假设 $\tilde{N}$表示预设的预训练轮数，且： X i = α i A + D , Y i = α i B + C , Z = F ( w ∗ ) , ∀ i ∈ { 1 , 2 } X_i=\alpha_iA+D,Y_i=\alpha_iB+C, Z=F(\boldsymbol{w}_{*}),\forall i\in\{1,2\} Xi​=αi​A+D,Yi​=αi​B+C,Z=F(w∗​),∀i∈{1,2}，那么根据 $\hat{U}(N)$ 定义有：
   $$F_{i,n}-Z\approx \frac{X_i}{n+Y_i},n\in [1,\tilde{N}],$$其中，${\alpha }_i=\frac{\sqrt{d}}{q_{i}K}+1$。（注：这里 i ∈ { 1 , 2 } i\in\{1,2\} i∈{1,2}是为了联立解出$A,B,C,D$的值）。为了得到合适的 $X_i,Y_i,Z$ 的值，一个有效的做法是解决如下非线性回归问题：
   $$\underset{X_i,Y_i,Z}{\min }\sum _{i=1}^2\sum _{n=1}^{\tilde{N}}{\left((F_{i,n}-Z)(n+Y_i)-X_i\right)}^2.$$注意到上述优化问题在 固定 $Z$ 的条件下 可以转化为两个线性回归问题：
   $$\underset{X_i,Y_i}{\min }\sum _{n=1}^{\tilde{N}}{\left((F_{i,n}-Z)(n+Y_i)-X_i\right)}^2.$$根据线性回归中的“最小二乘法”解出：
   $$\begin{gathered} X_i=\frac{{\sum }_{n=1}^{\tilde{N}}{\chi }_{i,n}{\sum }_{n=1}^{\tilde{N}}{\psi }_{i,n}^2-{\sum }_{n=1}^{\tilde{N}}{\chi }_{i,n}{\psi }_{i,n}{\sum }_{n=1}^{\tilde{N}}{\psi }_{i,n}}{N{\sum }_{n=1}^{\tilde{N}}{\psi }_{i,n}^2-{\left({\sum }_{n=1}^{\tilde{N}}{\psi }_{i,n}\right)}^2}, \\ {Y}_i=\frac{{\sum }_{n=1}^{\tilde{N}}{\chi }_{i,n}{\sum }_{n=1}^{\tilde{N}}{\psi }_{i,n}-N{\sum }_{n=1}^{\tilde{N}}{\chi }_{i,n}{\psi }_{i,n}}{N{\sum }_{n=1}^{\tilde{N}}{\psi }_{i,n}^2-{\left({\sum }_{n=1}^{\tilde{N}}{\psi }_{i,n}\right)}^2}. \end{gathered}$$其中，${\chi }_{i,n}=(F_{i,n}-Z)n,{\psi }_{i,n}=F_{i,n}-Z$。对于 $Z$是变量的情况，可以在 $F_{i,\tilde{N}}$ 周围进行 一维搜索得到。这样，$A,B,C,D$ 的值可以由 { X i } , { Y i } \{X_i\},\{Y_i\} {Xi​},{Yi​} 表示为：
   $$\{\begin{array}{l}A=\frac{X_1-X_2}{{\alpha }_1-{\alpha }_2},\\ B=\frac{Y_1-Y_2}{{\alpha }_1-{\alpha }_2},\\ C=\frac{{\alpha }_2{Y}_1-{\alpha }_1{Y}_2}{{\alpha }_2-{\alpha }_1},\\ D=\frac{{\alpha }_2{X}_1-{\alpha }_1{X}_2}{{\alpha }_2-{\alpha }_1}.\end{array}$$

3. 因此，可以得到 $N_{ϵ}$ 的估计值为：
   $$N_{ϵ}=\lceil \left(\frac{\sqrt{d}}{qK}+1\right)\left(\frac{A}{ϵ}-B\right)+\frac{D}{ϵ}-C\rceil .$$

2.3 训练时间最小化

本文通过优化 “量化精度”（与每轮训练时间和总训练轮数都相关） 和 “带宽分配”（只和每轮训练时间相关） 最小化训练时间。

上述问题可以表示为：
( P 1 ) min ⁡ q ∈ Z + , { b k } T d N ϵ s . t .   T k c o m p + T k c o m m ≤ T d ( ∀ k ∈ [ K ] ) , ∑ k = 1 K b k = B 0 , q ≥ 2. \begin{gathered} (\mathrm{P}1)\quad\operatorname*{min}_{q\in\mathbb{Z}^{+},\{b_{k}\}}T_{\mathrm{d}}N_{\epsilon} \\ \mathrm{s.t.~}T_{k}^{\mathrm{comp}}+T_{k}^{\mathrm{comm}}\leq T_{\mathrm{d}}(\forall k\in[K]), \sum_{k=1}^{K}b_{k}=B_{0}, q \geq2. \end{gathered} (P1)q∈Z+,{bk​}min​Td​Nϵ​s.t. Tkcomp​+Tkcomm​≤Td​(∀k∈[K]),k=1∑K​bk​=B0​,q≥2.​由于上述优化问题涉及到两个控制变量（$q$ 和 { b k } \{b_k\} {bk​}），因此直接解决是困难的。这篇文章中，作者将 问题(P1) 转化为两个子问题，具体来说是：

• 固定量化精度 $q$，优化带宽分配 { b k } \{b_k\} {bk​}：借用 双重二分查找技术
• 固定带宽分配 { b k } \{b_k\} {bk​}，优化量化精度 $q$：借用 逐次凸逼近技术SCA

2.3.1 带宽分配优化

在固定带宽的条件下，问题(P1) 转化为：
( P 2 ) min ⁡ { b k } , T d T d s . t .   T k c o m p + T k c o m m ≤ T d ( ∀ k ∈ [ K ] ) , ∑ k = 1 K b k = B 0 . \begin{gathered}(\text{P}2)\quad\min_{\{b_k\},T_{\mathrm{d}}}T_{\mathrm{d}}\\\mathrm{s.t.~}T_k^\mathrm{comp}+T_k^\mathrm{comm}\leq T_\mathrm{d}(\forall k\in[K]),\sum_{k=1}^Kb_k=B_0.\end{gathered} (P2){bk​},Td​min​Td​s.t. Tkcomp​+Tkcomm​≤Td​(∀k∈[K]),k=1∑K​bk​=B0​.​注意到 $T_k^{\mathrm{comm}}=\frac{S}{R_k}$，因此有：
( P 2 ) min ⁡ { b k } , T d T d s . t .   T k c o m p + S R k ≤ T d ( ∀ k ∈ [ K ] ) , ∑ k = 1 K b k = B 0 . . \begin{gathered}(\text{P}2)\quad\min_{\{b_k\},T_{\mathrm{d}}}T_{\mathrm{d}}\\\mathrm{s.t.~}T_k^\mathrm{comp}+\frac{S}{R_{k}}\leq T_\mathrm{d}(\forall k\in[K]),\sum_{k=1}^Kb_k=B_0.\end{gathered}. (P2){bk​},Td​min​Td​s.t. Tkcomp​+Rk​S​≤Td​(∀k∈[K]),k=1∑K​bk​=B0​.​.

根据 $R_k$ 定义，可以得到：这部分没有看懂
$$R_k=-\frac{b_k}{\ln 2}{\mathrm{e}}^{b_k{\theta }_k}\mathrm{Ei}(-b_k{\theta }_k)$$其中，${\theta }_k=\frac{N_0}{p_k{\varphi }_k},\mathrm{Ei}(x)={\int }_{-\infty }^x\frac{{\mathrm{e}}^{\rho }}{\rho }\mathrm{d}\rho$。可以证明，函数 $R_k(b_k)$ 单调递增。

原文得到引理：
Lemma 1：问题(P2) 中的第一个约束条件可以代替为：
$$T_k^{\mathrm{comp}}+T_k^{\mathrm{comm}}=T_{\mathrm{d}},\forall k\in [K]$$因此，可以得到：
$$R_k(b_k)=\frac{S}{T_{\mathrm{d}}-T_k^{\mathrm{comp}}},$$由于函数 $R_k(b_k)$ 的 单调性，可以使用 “二分查找”（双重） 进行求解：

• 外层二分查找：搜索 $T_d$ 在范围 $[T_d^{-},T_d^{+}]$，这里 T d − = max ⁡ k { T k c o m p } , T d + = max ⁡ k { T k c o m p + R k ( B 0 / K ) } T_{\mathrm{d}}^{-}=\operatorname*{max}_{k}\{T_{k}^{\mathrm{comp}}\},T_{\mathrm{d}}^{+}=\max_{k}\{T_{k}^{\mathrm{comp}}+R_{k}(B_{0}/K)\} Td−​=maxk​{Tkcomp​},Td+​=maxk​{Tkcomp​+Rk​(B0​/K)}; 注：这里我认为有问题，应该是 T d + = max ⁡ k { T k c o m p + R k ( S R k ( B 0 ) ) } T_{\mathrm{d}}^{+}=\max_{k}\{T_{k}^{\mathrm{comp}}+R_{k}(\frac{S}{R_k(B_0)})\} Td+​=maxk​{Tkcomp​+Rk​(Rk​(B0​)S​)}
• 内层二分查找：根据公式 $R_k=-\frac{b_k}{\ln 2}{\mathrm{e}}^{b_k{\theta }_k}\mathrm{Ei}(-b_k{\theta }_k)$搜索 $b_k$，搜索范围 $[0,B_0]$

2.3.2 量化精度优化

这里先假设 $p$ 是正实数进行优化。

在固定带宽的条件下，问题(P1) 可以转化为：
$$\begin{array}{rl}(\text{P3})& \underset{q}{\min }\tilde{T}\\ \text{s.t.}& (T_k^{\mathrm{comp}}+T_k^{\mathrm{comm}})\left(\frac{\sqrt{d}}{qK}{H}_1+H_2\right)\le \tilde{T},q\ge 2.\end{array}$$其中，引入 $\tilde{T}$ 是为了消除 $N_{ϵ}$ 的取上界操作；$H_1=\frac{A}{ϵ}-B,H_2=\frac{A+D}{ϵ}-B-C$。

使用 SCA 技术将院非凸问题，转化为局部凸问题进行求解。核心思路是：改变第一个约束条件。根据 $T_k^{\mathrm{comm}}=\frac{(1+{\log }_2(1+q))d}{R_k}$，则第一个约束条件为：
$$\left(T_k^{\mathrm{comp}}+\frac{\left(1+{\log }_2(1+q)\right)d}{R_k}\right)\left(\frac{\sqrt{d}}{qK}{H}_1+H_2\right)\le \tilde{T}.$$对等式两边取对数有：（对数的优势是可以将乘法变加法）
$$J_k(q)-\ln (qK)-\ln \tilde{T}\le 0,$$其中，$J_k(q)=\ln \left(T_k^{\mathrm{comp}}+\frac{(1+{\log }_2(1+q))d}{R_k}\right)+\ln \left(qK{H}_2+H_1\sqrt{d}\right)$。可以证明 函数$J_k(q)$ 是凹函数，那么根据一阶泰勒公式有：
$$J_k(q)\le J_k\left(q^{(r)}\right)+J_k^{\prime }\left(q^{(r)}\right)\left(q-q^{(r)}\right)\triangleq {\hat{J}}_k(q)$$其中，导数 $J_k^{\prime }\left(q^{(r)}\right)$ 为：
$$\begin{array}{r}{J}_k^{\prime }\left(q^{(r)}\right)=\frac{K{H}_2}{q^{(r)}K{H}_2+H_1\sqrt{d}}+\frac{1}{\ln 2\left[{\log }_2(1+q^{(r)})+\frac{1}{d}{R}_k{T}_k^{\mathrm{comm}}+1\right](1+q^{(r)})}\end{array}$$因此，问题(P3) 可以是：
$$\begin{array}{r}(P3.1)q^{(r+1)}=\arg \underset{q}{\min }\tilde{T}\\ \text{s.t.>}{\hat{J}}_k(q)-\ln (qK)-\ln (\tilde{T})\le 0(\forall k\in [K]),q\ge 2.\end{array}$$这是一个 凸问题，可以通过 CVXPY 技术进行求解。

这里对凸问题还是凹问题的判断不是很清楚

2.3.3 联合算法

由于之前假设 $q$ 是正实数，因此需要找到合适的正整数赋值给 $q$，这篇文章使用如下方法得到：
q ∗ = arg ⁡ min ⁡ q ∈ { ⌈ q ^ ⌉ − 1 , ⌈ q ^ ⌉ } T ( q , { b k ∗ } ) q^{*}=\arg\min_{q\in\{\lceil\hat{q}\rceil-1,\lceil\hat{q}\rceil\}}T(q,\{b_{k}^{*}\}) q∗=argq∈{⌈q^​⌉−1,⌈q^​⌉}min​T(q,{bk∗​})

三、讨论

这篇文章作者写出的缺点（未来方向）有：

• 没有涉及到 client 采样。在这种情况下如何惊醒 带宽分配；
• 没有涉及到误差补偿，这种技术可以减少量化精度。

我认为的优点是：

• 对通讯时间进行了表征（$R_k$），而不是单单使用时间开估算
• 考虑了一轮训练时间和总训练轮数之间的 tradeoff

我认为的缺点有：

• 本文聚焦于 “减少总训练时间”，而对模型精度问题没有进行考虑
• 需要进行一定程度的预训练

四、定理证明

4.1 L-Smooth 和 Strongly convex 的相关定义和推论

4.1.1 L-Smooth

L-Smooth定义： 假设 $L>0$，如果一个在在集合 $D\subset \mathbb{E}$ 上可微的函数 $f:\mathbb{E}\to (-\infty ,+\infty ]$ 是 L-Smooth，那么函数将满足：
$$\Vert \nabla f(\mathbf{x})-\nabla f(\mathbf{y})\Vert \le L\Vert \mathbf{x}-\mathbf{y}\Vert \text{for all }\mathbf{x},\mathbf{y}\in D$$其中，$L$被称为 smoothness parameter，$\Vert \cdot \Vert$表示 2范数。

Lemma 1(descent lemma)：
$$f(\mathbf{y})\le f(\mathbf{x})+⟨\nabla f(\mathbf{x}),\mathbf{y}-\mathbf{x}⟩+\frac{L}{2}\Vert \mathbf{x}-\mathbf{y}{\Vert }^2$$
Lemma 2：
$$f(\mathbf{y})-f(\mathbf{x})-⟨\nabla f(\mathbf{x}),\mathbf{y}-\mathbf{x}⟩\ge \frac{1}{2L}\Vert \nabla f(\mathbf{x})-\nabla f(\mathbf{y}){\Vert }_{}^2$$
Lemma 3：
$$⟨\nabla f(\mathbf{x})-\nabla f(\mathbf{y}),\mathbf{x}-\mathbf{y}⟩\le L\Vert \mathbf{x}-\mathbf{y}{\Vert }^2$$
Lemma 4：
$$⟨\nabla f(\mathbf{x})-\nabla f(\mathbf{y}),\mathbf{x}-\mathbf{y}⟩\ge \frac{1}{L}\Vert \nabla f(\mathbf{x})-\nabla f(\mathbf{y}){\Vert }^2$$

4.1.2 Strongly convex

将在L-Smooth 部分得到的 所有不等式变换方向 并且 将 $L$ 替换为 $\mu$ 就得到了Strongly Convex 函数满足的性质：
$$\begin{array}{rl}f(\mathbf{y})-f(\mathbf{x})-⟨\nabla f(\mathbf{x}),\mathbf{y}-\mathbf{x}⟩& \ge \frac{\mu }{2}\Vert \mathbf{x}-\mathbf{y}{\Vert }^2\\ ⟨\nabla f(\mathbf{x})-\nabla f(\mathbf{y}),\mathbf{x}-\mathbf{y}⟩& \ge \mu \Vert \mathbf{x}-\mathbf{y}{\Vert }^2\\ \Vert \nabla f(\mathbf{x})-\nabla f(\mathbf{y}){\Vert }_{}& \ge \mu \Vert \mathbf{x}-\mathbf{y}\Vert \\ f(\mathbf{y})-f(\mathbf{x})-⟨\nabla f(\mathbf{x}),\mathbf{y}-\mathbf{x}⟩& \le \frac{1}{2\mu }\Vert \nabla f(\mathbf{x})-\nabla f(\mathbf{y}){\Vert }_{}^2\\ ⟨\nabla f(\mathbf{x})-\nabla f(\mathbf{y}),\mathbf{x}-\mathbf{y}⟩& \le \frac{1}{\mu }\Vert \nabla f(\mathbf{x})-\nabla f(\mathbf{y}){\Vert }_{}^2\end{array}$$

1.https://zhuanlan.zhihu.com/p/369961290 点击打开知乎博客