联邦学习中的原始对偶算法

联邦学习中的原始对偶算法#

在传统的最优化方法中，原始-对偶算法 (Primal-Dual Algorithms) 也是一类求解具有显式原始-对偶结构问题的常用算法。考虑带显式全局变量约束的共识优化问题（注意它与 算子分裂 与 跳步 两节中共识集格式的细微差别），

(1)#\[ \begin{align}\begin{aligned}\DeclareMathOperator*{\minimize}{minimize} \DeclareMathOperator*{\maximize}{maximize} \DeclareMathOperator*{\argmin}{arg\,min} \newcommand{\col}{\operatorname{col}} \newcommand{\R}{\mathbb{R}}\\\begin{split}\begin{array}{cl} \minimize & F(\Theta) := \sum\limits_{k=1}^K f_k(\theta_k), \\ \text{subject to} & \theta_k = \theta, ~ k = 1, \ldots, K, \end{array}\end{split}\end{aligned}\end{align} \]

其中 \(\Theta = \col(\theta_1, \cdots, \theta_K)\) ，\(\theta, \theta_1, \ldots, \theta_K \in \R^d\) 。约束优化问题 (1) 的* 增广拉格朗日函数* (Augmented Lagrangian, AL) 为

(2)#\[\mathcal{L}(\theta, \Theta, \Lambda) = \sum\limits_{k=1}^K \mathcal{L}_k(\theta, \theta_k, \lambda_k), \qquad \mathcal{L}_k(\theta, \theta_k, \lambda_k) = f_k(\theta_k) + \langle \lambda_k, \theta_k - \theta \rangle + \frac{1}{2s} \lVert \theta_k - \theta \rVert^2,\]

其中 \(\Lambda = \col(\lambda_1, \ldots, \lambda_K)\) 是* 对偶变量* (Dual Variable)，亦称拉格朗日乘子 (Lagrangian Multiplier)，\(s > 0\) 是罚参数。将梯度提升类的算法应用于对偶问题 \(\maximize_\Lambda \ G(\Lambda) := \inf_{\theta, \Theta} \mathcal{L}(\theta, \Theta, \Lambda)\) ，即得到求解这类问题的一般迭代格式

(3)#\[\begin{split}\begin{aligned} \theta_k^{(t+1)} & = \argmin_{\theta_k} \mathcal{L}_k\left( \theta^{(t)}, \theta_k, \lambda_k^{(t)} \right), \\ \lambda_k^{(t+1)} & = \lambda_k^{(t)} + \frac{1}{s} \left( \theta_k^{(t+1)} - \theta^{(t)} \right), \\ \theta^{(t+1)} & = \argmin_{\theta} \sum\limits_{k=1}^K \left\{ \frac{1}{2s} \lVert \theta_k^{(t+1)} - \theta \rVert^2 - \langle \lambda_k^{(t+1)}, \theta \rangle \right\} = \frac{1}{K} \sum\limits_{k=1}^K \left( \theta_k^{(t+1)} + s\lambda_k^{(t+1)} \right). \end{aligned}\end{split}\]

FedPD ：联邦原始对偶算法#

基于这一迭代格式，文献 [Zhang et al.[1]] 设计了联邦原始对偶算法 FedPD ，其交替执行子节点侧的本地更新、子节点侧的对偶更新以及中心节点侧的全局平均。有两点设计值得特别注意。

** 灵活的局部求解器。 ** 对于 (3) 中的本地更新（即局部增广拉格朗日函数 \(\mathcal{L}_k\) 的极小化）， FedPD 允许子节点使用任意的求解器 \(\operatorname{\mathbf{Oracle}}_k\)——可以是普通的随机梯度下降，也可以如 SCAFFOLD 那样结合方差缩减技术。这种灵活性是联邦优化的显著特点（与中心节点侧的 FedOpt 是同一思想）。在使用方差缩减时，还可以进一步将 \(f_k\) 在当前内循环迭代点 \(\theta_k^{(t, r)}\) 处* 线性化* （记局部梯度为 \(g_k^{(t, r)}\) ），此时局部增广拉格朗日函数成为二次函数，其极小点具有解析形式，

(4)#\[\argmin \ \widetilde{\mathcal{L}}(\theta_{k}) = \frac{s}{s+\gamma} \, \theta_k^{(t,r)} + \frac{\gamma}{s+\gamma} \, \theta_{k, 0}^{(t)} - \frac{s\gamma}{s+\gamma}\left( g_k^{(t,r)} + \lambda_k^{(t)} \right),\]

其中 \(\gamma\) 是线性化（临近）系数，\(\theta_{k, 0}^{(t)}\) 是本轮开始时局部模型的取值。

跳过全局步。 (3) 中的外层（中心节点）步仅以概率 \(1 - p\) 执行。当它被跳过时，子节点用自身的 \(\widetilde{\theta}^{(t)} = \theta_k^{(t-1)} + s \lambda_k^{(t-1)}\) 代替全局模型，

(5)#\[\begin{split}\theta^{(t+1)} = \begin{cases} \frac{1}{K} \sum\limits_{k=1}^K \left( \theta_k^{(t+1)} + s\lambda_k^{(t+1)} \right), & \text{with probability } 1 - p, \\ \theta^{(t)}, & \text{with probability } p, \end{cases}\end{split}\]

于是代价高昂的中心节点轮次——恰恰是通信瓶颈——以概率 \(p\) 被跳过。更进一步，跳步概率还可以做成* 异质感知* 的：[Zhang et al.[1]]（TABLE III）给出了可达的 \(p\) 与子节点间数据分布异质程度 \(G\) （通过梯度差异有界性刻画）、(ε-)稳定点的目标精度 \(\varepsilon\) 以及 \(f, f_k\) 的 Lipschitz 常数 \(L\) 之间的关系，其中步长取 \(s = \frac{\sqrt{5}-1}{8L}\) ：

(6)#\[\begin{split}p \approx \begin{cases} \sqrt{\frac{1}{36s} \cdot \frac{\varepsilon}{G^2}}, & \text{if } 0 \leqslant p < \frac{1-2Ls}{1+Ls}, \\ 1 - 2 \Big/ \log\left(\frac{1}{42s} \cdot \frac{\varepsilon}{G^2}\right), & \text{if } \frac{1-2Ls}{1+Ls} < p < 1. \end{cases}\end{split}\]

子节点间的数据分布越均匀（\(G\) 越小），通信轮就可以越激进地被跳过。伪代码总结如下。

\begin{algorithm}
\caption{PseudoCode for \texttt{FedPD}}
\begin{algorithmic}
\REQUIRE step size $s = \frac{1}{\mu} > 0,$ skip probability $p \in [0, 1)$

\STATE {\bfseries Initiation:}
\STATE $\hspace{1.3em}$ {\bfseries Init server:} global parameters $\theta^{(0)} \in \mathbb{R}^d$
\STATE $\hspace{1.3em}$ {\bfseries Init clients:} local parameters$\theta_{0,k}^{(0)} \in \mathbb{R}^d,$ dual variables $\lambda_k^{(0)} \in \mathbb{R}^d,$ $\forall k \in [K]$

\FOR{each round $t = 0, 1, \cdots$}
    \FOR{each client $k = 1, \cdots, K$ {\bfseries in parallel}}
        \STATE $\theta_k^{(t+1)} \gets \textbf{Oracle}_k(\mathcal{L}_k(\theta_{k, 0}^{(t)}, \theta_k, \lambda_k^{(t)}), \theta_k^{(t)})$ \COMMENT{$\textbf{Oracle}_k$ can be SGD, etc.}
        \STATE $\lambda_k^{(t+1)} \gets \lambda_k^{(t)} + \frac{1}{s} (\theta_k^{(t+1)} - \theta_{k, 0}^{(t)})$ \COMMENT{dual update step}
        \STATE $\theta_{k, 0}^{(t+\frac{1}{2})} \gets \theta_k^{(t+1)} + s \lambda_k^{(t+1)}$
    \ENDFOR
    \STATE with probability $1 - p$ do global communication
    \STATE $\hspace{1.3em}$ client $k$ send $\theta_{k, 0}^{(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, 0}^{(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,0} \gets \theta^{(t+1)}, ~ \forall k \in [K]$
    \STATE with probability $p$ skip global communication:
    \STATE $\hspace{1.3em}$ {\bfseries Client Update:} $\theta^{(t+1)}_{k,0} \gets \theta_{k, 0}^{(t+\frac{1}{2})}, ~ \forall k \in [K]$
    \STATE $\hspace{1.3em}$ \COMMENT{On server, $\theta^{(t+1)} \gets \theta^{(t)}$}
\ENDFOR
\end{algorithmic}
\end{algorithm}

参考文献

目录