联邦学习中的算子分裂算法

为什么要算子分裂？收敛结果的正确性

已有的联邦学习算法大多处理非 I.I.D. 场景下的收敛率与稳定性问题，而它们收敛到的点的* 正确性* 却长期未被重视。一个简单的最小二乘例子表明， FedAvg 与 FedProx 的收敛结果一般不同于联邦问题的真解 [Pathak and Wainwright[1]]。

考虑用最小二乘拟合线性模型 \(y = mx\) （\(m\) 为待拟合变量），设有两个参与方：参与方 1 拥有数据 \(\{(0, 2), (1, 2)\}\) ，相应的目标函数为

\[f_1(m_1) = \frac{1}{2} \left( 4 + (m_1 - 2)^2 \right);\]

参与方 2 拥有数据 \(\{(2, 0), (2, 1)\}\) ，相应的目标函数为

\[f_2(m_2) = \frac{1}{2} \left( 4 m_2^2 + (2 m_2 - 1)^2 \right).\]

两个参与方单独拟合的结果分别为 \(m_1 = 2\) 与 \(m_2 = \frac{1}{4}\) ； FedAvg 的结果是平均值 \(m = \frac{9}{8}\) ，而 FedProx （极小化带临近项的目标函数）的结果是

(1)\[m = \frac{18 + 4\mu}{16 + 9\mu},\]

其中 \(\mu\) 为临近项系数。然而，把所有数据放在一起拟合的结果是 \(m = \frac{4}{9}\) 。差异的根源在于目标函数关于 \(m\) 是二次的：单独拟合结果的平均值一般不等于整体拟合的结果。

更一般地，[Pathak and Wainwright[1]] 分析了 FedAvg （其简化形式，即子节点上做全量梯度下降，记作 FedGD ）与 FedProx 的收敛点：假设每个 \(f_k\) 是凸的且 \(L\)-光滑，并且算法生成的序列收敛，那么所有子节点的模型参数序列有公共极限 \(\theta^*\) ，且该极限满足如下不动点条件

(2)\[\sum\limits_{i=1}^R \sum\limits_{k=1}^K \nabla f_k\big( \mathcal{G}_k^{i-1}(\theta^*) \big) = 0, \qquad \mathcal{G}_k(\theta_k) := \theta_k - \eta \nabla f_k(\theta_k),\]

对 FedGD 而言（\(R\) 为本地步数）；

(3)\[\sum\limits_{k=1}^K \nabla \mathcal{M}_{f_k, \mu}(\theta^*) = 0, \qquad \mathcal{M}_{f_k, \mu}(\theta) := \inf_{\theta_k} \left\{ f_k(\theta_k) + \frac{\mu}{2} \lVert \theta - \theta_k \rVert^2 \right\},\]

对 FedProx 而言，其中 \(\mathcal{M}_{f_k, \mu}\) 是 \(f_k\) 的 Moreau 包络。把它们应用到最小二乘问题 \(f_k(\theta_k) = \frac{1}{2} \lVert A_k \theta_k + b_k \rVert^2\) （\(A_k\) 列满秩），这些不动点有解析表达式，且一般不同于整体问题的真解 \(\theta^* = \big( \sum_{k=1}^K A_k^{\mathrm{T}} A_k \big)^{-1} \sum_{k=1}^K A_k^{\mathrm{T}} b_k\) 。

上述观察促使我们把联邦问题重新表述成一种等价形式，使其基于算子分裂的解收敛到 正确 的点——这正是下文 FedSplit 与 FedDR 的出发点。

从单调包含问题到分裂方法

从联邦问题的约束形式出发，即 \(\minimize \ F(\Theta) := \frac{1}{K} \sum_{k=1}^K f_k(\theta_k)\) 、约束 \(\Theta = \col(\theta_1, \ldots, \theta_K) \in \mathcal{E}\) （共识集 \(\mathcal{E} = \{ \Theta \mid \theta_1 = \cdots = \theta_K \}\) ），并用示性函数 \(\iota_{\mathcal{E}}\) （在 \(\mathcal{E}\) 上取 0、否则取 \(+\infty\) ）作为罚函数把约束化为目标项。带罚问题 \(\minimize \ F(\Theta) + \iota_{\mathcal{E}}(\Theta)\) 的一阶最优性条件为

(4)\[ \begin{align}\begin{aligned}\DeclareMathOperator*{\minimize}{minimize} \DeclareMathOperator*{\argmin}{arg\,min} \DeclareMathOperator*{\expectation}{\mathbb{E}} \newcommand{\col}{\operatorname{col}} \newcommand{\R}{\mathbb{R}} \newcommand{\prox}{\mathbf{prox}}\\0 \in \nabla F(\Theta) + \partial \iota_{\mathcal{E}}(\Theta) = \nabla F(\Theta) + \mathcal{N}_{\mathcal{E}}(\Theta),\end{aligned}\end{align} \]

这是一个* 单调包含问题* (Monotone Inclusion Problem, MIP)：\(\nabla F\) 与法锥 \(\mathcal{N}_{\mathcal{E}}\) （在 \(\mathcal{E}\) 上为 \(\mathcal{E}^{\perp}\) 、其余为空）都是* 极大单调算子* (Maximal Monotone Operator) [Ryu and Yin[2]]。分裂此类包含问题 \(0 \in \mathcal{A}(x) + \mathcal{B}(x)\) 的核心工具是* 预解式* (Resolvent) \(\mathcal{R}_{s\mathcal{T}} := (\mathcal{I} + s\mathcal{T})^{-1}\) 与* 反射预解式* （Cayley 算子）:math:mathcal{C}_{smathcal{T}} := 2mathcal{R}_{smathcal{T}} - mathcal{I} 。对梯度或次微分而言，预解式即临近算子 \(\mathcal{R}_{s\partial f} = \prox_{s f}\) ；对示性函数则为投影 \(\mathcal{R}_{s\partial\iota_{\mathcal{E}}} = \Pi_{\mathcal{E}}\) 。关键的等价关系是：求解该包含问题等价于求* 非扩张* 算子复合 \(\mathcal{C}_{s\mathcal{A}} \mathcal{C}_{s\mathcal{B}}\) 的不动点：

(5)\[0 \in \mathcal{A}(x) + \mathcal{B}(x) \quad \Longleftrightarrow \quad \mathcal{C}_{s\mathcal{A}} \mathcal{C}_{s\mathcal{B}}(z) = z, \quad x = \mathcal{R}_{s\mathcal{B}}(z).\]

由于 \(\mathcal{C}_{s\mathcal{A}} \mathcal{C}_{s\mathcal{B}}\) 只具有非扩张性，不动点迭代实际使用其平均算子 \((1 - \tau)\mathcal{I} + \tau\, \mathcal{C}_{s\mathcal{A}} \mathcal{C}_{s\mathcal{B}}\) （\(\tau \in (0, 1]\) ），由此得到分裂方法族

(6)\[\begin{split}\begin{aligned} x^{(t+1/2)} & = \mathcal{R}_{s\mathcal{B}}\left( z^{(t)} \right), \\ z^{(t+1/2)} & = 2 x^{(t+1/2)} - z^{(t)}, \\ x^{(t+1)} & = \mathcal{R}_{s\mathcal{A}}\left( z^{(t+1/2)} \right), \\ z^{(t+1)} & = z^{(t)} + 2\tau \left( x^{(t+1)} - x^{(t+1/2)} \right). \end{aligned}\end{split}\]

\(\tau = 1\) 对应 Peaceman--Rachford 分裂方法 (PRSM)；\(\tau = \frac{1}{2}\) 对应 Douglas--Rachford 分裂方法 (DRSM) [Douglas and Rachford[3], Lions and Mercier[4]]。两者的区别只在最后一步（松弛步）——即 \(\tau\) 的选取——而这一区别是决定性的：与 PRSM 不同，DRSM 对一般的极大单调算子无需单调性之外的任何附加假设即有收敛性保证 [Ryu and Yin[2]]。

FedSplit ：联邦 Peaceman--Rachford 分裂算法

把上述机制应用于 \(\mathcal{A} = \nabla F\) 、\(\mathcal{B} = \partial\iota_{\mathcal{E}} = \mathcal{N}_{\mathcal{E}}\) （并取 \(s = \frac{1}{\mu}\) ）。两个预解式分别为 \(\mathcal{R}_{s\nabla F} = \prox_{F, \mu}\) 与 \(\mathcal{R}_{s\mathcal{N}_{\mathcal{E}}} = \Pi_{\mathcal{E}}\) ，而且——对联邦学习至关重要的是——块可分目标函数的临近算子可分解为子节点本地的临近步，

(7)\[\prox_{F, \mu}(\Theta) = \argmin_{\omega} \left\{ \sum_{k=1}^K f_k(\omega_k) + \frac{\mu}{2} \sum_{k=1}^K \lVert \omega_k - \theta_k \rVert^2 \right\} = \left( \prox_{f_k, \mu}(\theta_k) \right)_{k=1}^K,\]

而投影 \(\Pi_{\mathcal{E}}\) 恰是中心节点上的平均聚合。于是整个分裂在子节点间完全解耦合，天然契合中心节点--子节点的通信模式。将 PRSM（\(\tau = 1\) ）应用于此即得 FedSplit 算法 [Pathak and Wainwright[1]]——分裂方法在联邦学习领域的首次应用——其中 \(\texttt{prox\_update}_k\) 是子节点上 \(\prox_{f_k, \mu}\) 的非精确求解方法（例如随机梯度下降）。

FedSplit 解决了前文的正确性问题：其迭代的固定点给出联邦问题真正的解，而非 FedAvg/FedProx 那样的偏移解。不过它的收敛性继承了 PRSM 的短板：对一般的非扩张 Cayley 算子，PRSM 没有收敛性保证，需要附加假设（\(f_k\) 强凸且 \(L\)-光滑），此时收敛率为

(8)\[\lVert \theta^{(t+1)} - \theta^* \rVert \leqslant \left( 1 - \frac{2}{\sqrt{\kappa}+1} \right)^t \frac{\lVert \theta^{(1)} - \theta^* \rVert}{\sqrt{K}} + (\sqrt{\kappa}+1)\, b,\]

其中 \(\ell_* = \min_k \ell_k\) 为强凸模量，\(L^* = \max_k L_k\) ，\(\kappa = L^* / \ell_*\) ，\(b\) 为非精确临近求解的一致误差界。伪代码总结如下。

\begin{algorithm}
\caption{PseudoCode for \texttt{FedSplit}}
\begin{algorithmic}
\REQUIRE {\bfseries proximal solvers} $\texttt{prox\_update}_k: \mathbb{R}^d \to \mathbb{R}^d$

\STATE {\bfseries Initiation:} parameters $\theta^{(0)} \in \mathbb{R}^d$

\FOR{each round $t = 0, 1, \cdots$}
    \STATE broadcast $\theta^{(t)}$ to clients $k \in [K]$
    \FOR{each client $k = 1, \cdots, K$ {\bfseries in parallel}}
        \STATE $\theta_k^{(t)} \gets \theta^{(t)}$
        \STATE $\theta_k^{(t+1/2)} \gets$ $\texttt{prox\_update}_k(2\theta^{(t)} - \theta_k^{(t)})$ \COMMENT{local prox step}
        \STATE $\theta_k^{(t+1)} \gets$ $\theta_k^{(t)} + 2(\theta_k^{(t+1/2)} - \theta^{(t)})$ \COMMENT{local centering step}
        \STATE send $\theta_k^{(t+1)}$ to server
    \ENDFOR
    \STATE {\bfseries Server Update:}
    \STATE $\hspace{1.3em}$ $\theta^{(t+1)} \gets \frac{1}{K} \sum\limits_{k=1}^K \theta_k^{(t+1)}$ \COMMENT{global averaging}
    \IF{meet convergent criteria}
        \STATE $\theta^* \gets \theta^{(t+1)}$
        \STATE {\bfseries break}
    \ENDIF
\ENDFOR
\end{algorithmic}
\end{algorithm}

FedDR ：联邦 Douglas--Rachford 分裂算法

FedDR 算法 [Tran-Dinh et al.[5]] 用具有收敛性保证的 DRSM（\(\tau = \frac{1}{2}\) ）替换 PRSM，并借此把可处理的问题类扩大到* 非凸* 的光滑目标函数，同时在模型中显式引入正则项，

(9)\[\minimize \quad F(\theta) = f(\theta) + g(\theta) = \frac{1}{K} \sum_{k=1}^K f_k(\theta) + g(\theta),\]

其中 \(f_k\) 是 \(L\)-光滑的（可以非凸），\(g\) 是闭的真是凸 (CCP) 正则项。为 \(g\) 再引入一个块 \(\theta_{K+1}\) ，并照旧用 \(\iota_{\mathcal{E}}\) 罚化共识约束 \(\theta_1 = \cdots = \theta_{K+1}\) ，一阶条件即成为三算子的单调包含

(10)\[0 \in \nabla f(\Theta) + \partial \left( g + \iota_{\mathcal{E}} \right)(\Theta),\]

其中 \(\mathcal{B} = \nabla f\) ，\(\mathcal{A} = \partial (g + \iota_{\mathcal{E}})\) 。\(\mathcal{B}\) 的预解式为 \(\mathcal{R}_{s\mathcal{B}} = \prox_{K s f}\)——按块分解即子节点本地的临近步；\(\mathcal{A}\) 的预解式则在中心节点上一并处理正则项* 与* 共识约束。DRSM 迭代（(6) 取 \(\tau = \frac{1}{2}\) ，松弛系数记作 \(\alpha\) ）为

(11)\[\begin{split}\begin{aligned} \omega_k^{(t+1)} & = \omega_k^{(t)} + \alpha \left( \overline{\theta}^{(t)} - \theta_k^{(t)} \right), & k \in [K], \\ \theta_k^{(t+1)} & = \prox_{f_k, \mu} \left( \omega_k^{(t+1)} \right), & k \in [K], \\ \widehat{\theta}_k^{(t+1)} & = 2\theta_k^{(t+1)} - \omega_k^{(t+1)}, & k \in [K], \\ \widetilde{\theta}^{(t+1)} & = \frac{1}{K}\sum_{k=1}^K \widehat{\theta}_k^{(t+1)}, & \\ \omega^{(t+1)} & = \omega^{(t)} + \alpha \left( \overline{\theta}^{(t)} - \omega^{(t)} \right), & \\ \overline{\theta}^{(t+1)} & = \prox_{g, \frac{K+1}{Ks}} \left( \frac{K}{K+1} \widetilde{\theta}^{(t+1)} + \frac{1}{K+1} \omega^{(t+1)} \right). & \end{aligned}\end{split}\]

前三行是子节点的本地工作；最后一行是中心节点——它平均反射块并对正则项 \(g\) 执行临近步。

关于收敛性 [Tran-Dinh et al.[5]]：设 \(f_k\) 都是 \(L\)-光滑的，记 \(e_k^{(t)}\) 为第 \(t\) 轮子节点非精确临近求解的误差：(i)（经正则项修正的）全局模型与各局部模型的差被局部模型相邻迭代步的差以及临近误差所控制；(ii) 全局目标函数的梯度映射 \(\mathcal{G}_{\beta}(\theta) := \frac{1}{\beta} \left( \theta - \prox_{\beta g}\left( \theta - \beta \nabla f(\theta) \right) \right)\) 在 \(\overline{\theta}^{(t)}\) 处的模长平方被同样的量（以及辅助变量 \(\omega^{(t)}\) 的漂移）所控制；(iii) 梯度映射模长平方的平均值具有 \(\mathcal{O}(1/T)\) 界，

(12)\[\frac{1}{T} \sum_{t=0}^{T-1} \expectation \left[ \lVert \mathcal{G}_{\frac{Ks}{K+1}} \left( \overline{\theta}^{(t)} \right) \rVert^2 \right] \leqslant \frac{C_1 \left( F(\theta^{(0)}) - F^* \right)}{T} + \frac{1}{KT} \sum_{t=0}^{T-1} \sum_{k=1}^K \left( C_2 \varepsilon_{k,t}^2 + C_3 \varepsilon_{k,t+1}^2 \right),\]

其中 \(F^* = \inf_\theta F(\theta)\) ；特别地，经过

(13)\[T = \left\lfloor \frac{C_1 \left[ F(\theta^{(0)}) - F^* \right] + (C_2 + C_3) M}{\varepsilon^2} \right\rfloor, \qquad M \geqslant \frac{1}{K} \sum_{t=0}^{T-1} \sum_{k=1}^K \varepsilon_{k,t}^2,\]

轮迭代后，输出 \(\overline{\theta}^{(T)}\) 即为 (9) 的 \(\varepsilon\)-稳定点。伪代码总结如下。

\begin{algorithm}
\caption{PseudoCode for \texttt{FedDR}}
\begin{algorithmic}
\REQUIRE step size $s = \frac{1}{\mu} > 0$, $\alpha \in (0, 2)$, error bounds $\varepsilon_{k,0} \geqslant 0$

\STATE {\bfseries Initiation:}
\STATE $\hspace{1.3em}$ {\bfseries Init server:} global model parameters $\theta^{(0)} \in \text{dom}(f)$, $\overline{\theta}^{(0)} = \widetilde{\theta}^{(0)} = \omega^{(0)} = \theta^{(0)} \in \mathbb{R}^d$
\STATE $\hspace{1.3em}$ {\bfseries Init clients:} $\omega_k^{(0)} = \theta^{(0)}$, $\theta_k^{(0)} \approx \mathbf{prox}_{f_k, \mu}(\omega_k^{(0)})$, $\widehat{\theta}_k^{(0)} = 2\theta_k^{(0)} - \omega_k^{(0)}, ~ \forall k \in [K]$

\FOR{$t = 0, 1, \cdots, T-1$}
    \STATE $\mathcal{S}^{(t)} \gets$ (random set of clients) $\subseteq [K]$
    \STATE each client $k \in \mathcal{S}^{(t)}$ receives $\overline{\theta}^{(t)}$ from server \COMMENT{communication}
    \FOR{each client $k \in \mathcal{S}^{(t)}$ {\bfseries in parallel}}
        \STATE choose $\varepsilon_{k,t+1} \geqslant 0$
        \STATE update $\omega_k^{(t+1)} \gets \omega_k^{(t)} + \alpha(\overline{\theta}^{(t)} - \theta_k^{(t)})$
        \STATE $\theta_k^{(t+1)} \approx \mathbf{prox}_{f_k, \mu}(y_k^{(t+1)})$ \COMMENT{inexact local prox step with error bound $\varepsilon_{k,0}$}
        \STATE $\widehat{\theta}_k^{(t+1)} \gets 2\theta_k^{(t+1)} - \omega_k^{(t+1)}$
        \STATE send $\Delta \widehat{\theta}_k^{(t)} = \widehat{\theta}_k^{(t+1)} - \widehat{\theta}_k^{(t)}$ to server
    \ENDFOR
    \STATE {\bfseries Server Update:}
    \STATE $\hspace{1.3em}$ $\omega^{(t+1)} \gets \omega^{(t)} + \alpha (\overline{\theta}^{(t)} - \omega^{(t)})$
    \STATE $\hspace{1.3em}$ $\widetilde{\theta}^{(t+1)} \gets \widetilde{\theta}^{(t)} + \frac{1}{K}\sum_{k \in \mathcal{S}^{(t)}} \Delta \widehat{\theta}_k^{(t)}$
    \STATE $\hspace{1.3em}$ $\overline{\theta}^{(t+1)} \gets \mathbf{prox}_{g, \frac{K+1}{Ks}} \left( \frac{K}{K+1} \widetilde{\theta}^{(t+1)} + \frac{1}{K+1} \omega^{(t+1)} \right)$
\ENDFOR
\end{algorithmic}
\end{algorithm}

小结

算子分裂法是* 分解算法* (Decomposition Algorithms) 的一大类：把大规模问题拆成若干分别处理的小问题（在这里，由子节点与中心节点分别承担），使每个子问题更易求解、更易并行。在理论上，它把众多联邦算法纳入统一框架——并且与 FedAvg/FedProx 不同，其极限保证是正确的（FedSplit ），或逼近真实目标的稳定点（FedDR ）。代价是簿记：除模型参数外，每个节点（中心与子节点）还要维护辅助变量（反射块 \(z\) 、混合拷贝 \(\omega\) ），带来额外的存储与计算负担。

参考文献

[1] (1,2,3)

Reese Pathak and Martin J Wainwright. FedSplit: An Algorithmic Framework for Fast Federated Optimization. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, 7057–7066. Curran Associates, Inc., 2020.

[2] (1,2)

Ernest K Ryu and Wotao Yin. Large-Scale Convex Optimization: Algorithms & Analyses via Monotone Operators. Cambridge University Press, 11 2022. doi:10.1017/9781009160865.

[3]

Jim Douglas and H. H. Rachford. On the Numerical Solution of Heat Conduction Problems in Two and Three Space Variables. Transactions of the American Mathematical Society, 82(2):421–439, 1956. doi:10.1090/s0002-9947-1956-0084194-4.

[4]

P. L. Lions and B. Mercier. Splitting Algorithms for the Sum of Two Nonlinear Operators. SIAM Journal on Numerical Analysis, 16(6):964–979, 12 1979. doi:10.1137/0716071.

[5] (1,2)

Quoc Tran-Dinh, Nhan Pham, Dzung T. Phan, and Lam M. Nguyen. FedDR – Randomized Douglas-Rachford Splitting Algorithms for Nonconvex Federated Composite Optimization. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems. 2021.