联邦学习中的跳步算法

联邦学习中的跳步算法#

最早的联邦学习算法 FedAvg [McMahan et al.[1]] 在子节点的内循环中执行多步梯度下降——从某种意义上说，这已经是一种* 跳步* (Skipping) 更新：通信轮以概率 \(1 - p = 0\) 被跳过。 ProxSkip 算法 [Mishchenko et al.[2], Zhang and Loizou[3]]（原文献中命名为 Scaffnew ）系统地研究了跳步技术在联邦学习优化算法中的应用，使通信轮以* 非零* 的概率被跳过。将 ProxSkip 视作 SCAFFOLD 结合跳步更新技术的改进算法，是不无启发的。

ProxSkip/Scaffnew#

从最优化的视角出发，将联邦学习问题建模为共识集 \(\mathcal{E} = \{ \Theta \mid \theta_1 = \cdots = \theta_K \}\) 上的共识优化问题

(1)#\[ \begin{align}\begin{aligned}\DeclareMathOperator*{\minimize}{minimize} \newcommand{\col}{\operatorname{col}} \newcommand{\R}{\mathbb{R}} \newcommand{\prox}{\mathbf{prox}}\\\begin{split}\begin{array}{cl} \minimize & F(\Theta) = \frac{1}{K} \sum\limits_{k=1}^K f_k(\theta_k), \\ \text{subject to} & \Theta \in \mathcal{E}, \\ \text{where} & \Theta = \col(\theta_1, \cdots, \theta_K) \in \R^{Kd}. \end{array}\end{split}\end{aligned}\end{align} \]

正如 算子分裂一节 所述，可以借助共识集的示性函数 \(\iota_{\mathcal{E}}\) 将该约束问题转化为无约束问题，并对其一阶最优性条件 \(0 \in \nabla F(\Theta) + \partial \iota_{\mathcal{E}}(\Theta)\) 进行分裂。更一般地，把 \(\iota_{\mathcal{E}}\) 换成任意 CCP 函数 \(g\) ，相应的不动点求解方法即* 临近梯度下降* (Proximal Gradient Descent, ProxGD)

(2)#\[\Theta^{(t+1)} = \prox_{s g}\left( \Theta^{(t)} - s \nabla F(\Theta^{(t)}) \right),\]

对问题 \(\minimize \ F(\Theta) + g(\Theta)\) 而言，ProxGD 的实际困难在于 CCP 函数 \(g\) 的临近算子 \(\prox_{s g}\) 通常求解代价高昂 [Friedlander and Goh[4]]。对应到联邦学习中，这一临近步* 正是* 通信本身：先由子节点向中心节点聚合模型参数，再由中心节点向子节点广播。以一定概率跳过这一性能瓶颈步骤便是一种自然的想法——这正是 ProxSkip 所做的事情。

关于控制变量 (Control Variates) 的说明：与 SCAFFOLD 不同， ProxSkip 并未（显式地）使用梯度方差缩减技术。控制变量 \(c_k\) 的首要作用是保证整体迭代的* 正确收敛性* ：可以证明在原问题的解点 \(\theta^*\) 处 \(c_k \to \nabla f_k(\theta^*)\) ，从而迭代的稳定点即为原问题的解点——这是 ProxSkip 的关键所在。 ProxSkip 整体上并不维护一个在中心节点与子节点之间通信的全局控制变量；相反，每当通信发生时，子节点上的局部控制变量 \(c_k\) 会由 \(\frac{p}{\gamma} \left( \theta^{(t+1)} - \theta_k^{(t + \frac{1}{2})} \right)\) * 隐式* 地更新，后者起到了全局控制变量的作用。这一点类似于 SCAFFOLD 中局部控制变量的 Option II 更新规则：好处是省去了中心节点与子节点之间额外的参数传递，代价是损失了控制变量更新方法的灵活性。伪代码总结如下。

\begin{algorithm}
\caption{PseudoCode for \texttt{ProxSkip}}
\begin{algorithmic}
\REQUIRE learning rate $\gamma,$ probability $p \in [0, 1)$

\STATE {\bfseries Initiation:}
\STATE $\hspace{1.3em}$ {\bfseries Server init:} global model parameters $\theta^{(0)} \in \mathbb{R}^d$
\STATE $\hspace{1.3em}$ {\bfseries Clients init:} local model parameters $\theta_k^{(0)} \in \mathbb{R}^d,$ control variates $c_k^{(0)} \in \mathbb{R}^d, ~ \forall k \in [K]$

\FOR{each round $t = 0, 1, \cdots, T-1$}
    \FOR{each client $k \in [K]$ {\bfseries in parallel}}
        \STATE $\theta_k^{(t, 0)} \gets \theta_k^{(t)}$
        \FOR{local step $r = 0, 1, \cdots, R-1$}
            \COMMENT{$R$ iterates of SGD with variance reduction}
            \STATE compute mini-batch gradient $g_k^{(t, r)}$ of $\nabla f_k(\theta_k^{(t, r)})$
            \STATE $\theta_k^{(t, r+1)} \gets \theta_k^{(t, r)} - \gamma \left( g_k^{(t, r)} - c_k^{(t)} \right)$
        \ENDFOR
        \STATE $\theta_k^{(t+\frac{1}{2})} \gets \theta_k^{(t, R)}$
    \ENDFOR
    \STATE with probability $1 − p$ do global communication
    \STATE $\hspace{1.3em}$ client $k$ send $\theta_{k}^{(t+\frac{1}{2})}$ to server $\forall k \in [K]$
    \STATE $\hspace{1.3em}$ {\bfseries Server Update:} $\theta^{(t+1)} \gets \frac{1}{K} \sum\limits_{k=1}^K \theta_{k}^{(t+\frac{1}{2})}$ \COMMENT{compute global average}
    \STATE $\hspace{1.3em}$ server broadcast $\theta^{(t+1)}$ to clients $k \in [K]$
    \STATE $\hspace{1.3em}$ on client $k:$ $\theta^{(t+1)}_{k} \gets \theta^{(t+1)}, ~ c_k^{(t+1)} \gets c_k^{(t)} + \frac{p}{\gamma}(\theta^{(t+1)}_{k} - \theta^{(t+\frac{1}{2})}_{k}), ~ \forall k \in [K]$
    \STATE with probability $p$ skip global communication:
    \STATE $\hspace{1.3em}$ {\bfseries Client Update:} $\theta^{(t+1)}_{k} \gets \theta_{k}^{(t+\frac{1}{2})}, ~ c_k^{(t+1)} \gets c_k^{(t)}, ~ \forall k \in [K]$
    \STATE $\hspace{1.3em}$ \COMMENT{on server, $\theta^{(t+1)} \gets \theta^{(t)}$}
\ENDFOR
\end{algorithmic}
\end{algorithm}

参考文献

目录