完全一阶梯度的联邦随机双层优化

# 全一阶梯度的联邦随机双层优化
来源：https://arxiv.org/html/2609.16350
Rohit Dhaipule机构：Stony Brook University 邮箱：[email protected]
Chiu C. Tan机构：Temple University 邮箱：{yihan.zhang0002
Haibin Ling机构：Stony Brook University 邮箱：chiu.tan
Hongchang Gao*通讯作者 机构：Temple University 邮箱：[email protected]

###### 摘要

联邦随机双层优化因其在机器学习中的广泛应用，近年来受到了积极研究。然而，大多数现有的联邦随机双层优化算法需要计算二阶海森矩阵和雅可比矩阵，这在实践中导致了更长的运行时间。为了应对这些挑战，我们提出了一种新颖的联邦随机方差缩减双层梯度下降算法，该算法仅依赖一阶预言机。具体而言，我们的方法不需要计算二阶海森矩阵和雅可比矩阵，显著减少了运行时间。此外，我们引入了一种新颖的学习率机制，即恒定的单时间尺度学习率，以协调不同变量的更新。我们还提出了一种新的策略来建立我们算法的收敛速率。最后，大量的实验证实了我们所提算法的有效性。

## 1 引言

本文关注以下联邦随机双层优化问题：

minx∈Rdx⁡f⁡\(x,y∗\(x\)\)≜1N∑n=1Nf\(n\)\(x,y∗\(x\)\)\\displaystyle\\min\_\{x\\in\\mathbb\{R\}^\{d\_\{x\}\}\}f\(x,y^\{\*\}\(x\)\)\\triangleq\\frac\{1\}\{N\}\\sum\_\{n=1\}^\{N\}f^\{\(n\)\}\(x,y^\{\*\}\(x\)\)s\.t\.,y∗\(x\)=arg⁡miny∈Rdy⁡g⁡\(x,y\)≜1N∑n=1Ng\(n\)\(x,y\)。\\displaystyle s\.t\.,y^\{\*\}\(x\)=\\arg\\min\_\{y\\in\\mathbb\{R\}^\{d\_\{y\}\}\}g\(x,y\)\\triangleq\\frac\{1\}\{N\}\\sum\_\{n=1\}^\{N\}g^\{\(n\)\}\(x,y\)\\ \.\(1\)
这里，g⁡\(x,y\)≜1N∑n=1Ng\(n\)\(x,y\) g\(x,y\)\\triangleq\\frac\{1\}\{N\}\\sum\_\{n=1\}^\{N\}g^\{\(n\)\}\(x,y\) 表示下层损失函数，其中g\(n\)\(x,y\)=E⁡\[g\(n\)\(x,y,ξ\(n\)\)\] g^\{\(n\)\}\(x,y\)=\\mathbb\{E\}\[g^\{\(n\)\}\(x,y;\\xi^\{\(n\)\}\)\] 表示第n个设备上的下层损失函数；f⁡\(x,y∗\(x\)\)≜1N∑n=1Nf\(n\)\(x,y∗\(x\)\) f\(x,y^\{\*\}\(x\)\)\\triangleq\\frac\{1\}\{N\}\\sum\_\{n=1\}^\{N\}f^\{\(n\)\}\(x,y^\{\*\}\(x\)\) 表示上层损失函数，其中f\(n\)\(x,y∗\(x\)\)=E⁡\[f\(n\)\(x,y∗\(x\),ξ\(n\)\)\] f^\{\(n\)\}\(x,y^\{\*\}\(x\)\)=\\mathbb\{E\}\[f^\{\(n\)\}\(x,y^\{\*\}\(x\);\\xi^\{\(n\)\}\)\] 表示第n个设备上的上层损失函数。在本文中，我们假设上层损失函数关于两个变量都是非凸的，而下层损失函数关于y是强凸的，这是现有文献中常用的假设 [Ghadimi and Wang (2018)](https://arxiv.org/html/2609.16350#bib.bib10); [Ji et al. (2021)](https://arxiv.org/html/2609.16350#bib.bib18); [Chen et al. (2021)](https://arxiv.org/html/2609.16350#bib.bib19); [Tarzanagh et al. (2022)](https://arxiv.org/html/2609.16350#bib.bib4)。

随机双层优化问题近年来引起了越来越多的关注，因为许多机器学习模型属于此类优化问题，例如模型无关元学习 [Finn et al. (2017)](https://arxiv.org/html/2609.16350#bib.bib8)、超参数优化 [Ji et al. (2021)](https://arxiv.org/html/2609.16350#bib.bib18)、神经架构搜索 [Liu et al. (2018)](https://arxiv.org/html/2609.16350#bib.bib16)等。为了解决随机双层优化问题，过去几年在单机环境下已经开发了许多优化算法 [Ghadimi and Wang (2018)](https://arxiv.org/html/2609.16350#bib.bib10); [Ji et al. (2021)](https://arxiv.org/html/2609.16350#bib.bib18); [Chen et al. (2021)](https://arxiv.org/html/2609.16350#bib.bib19); [Hong et al. (2020)](https://arxiv.org/html/2609.16350#bib.bib17); [Khanduri et al. (2021)](https://arxiv.org/html/2609.16350#bib.bib15); [Yang et al. (2021)](https://arxiv.org/html/2609.16350#bib.bib11); [Guo and Yang (2021)](https://arxiv.org/html/2609.16350#bib.bib12); [Dagréou et al. (2022)](https://arxiv.org/html/2609.16350#bib.bib9); [Chu et al. (2024)](https://arxiv.org/html/2609.16350#bib.bib22); [Dagréou et al. (2024)](https://arxiv.org/html/2609.16350#bib.bib23)。然而，为随机双层优化问题启用联邦学习比标准的单层优化问题更具挑战性。具体来说，如公式 (1) 所示，每个设备上的上层问题依赖于全局下层问题的最优解y∗\(x\) y^\{\*\}\(x\)。因此，计算每个设备上层损失函数的随机超梯度需要全局雅可比矩阵和海森矩阵，这给联邦学习中的局部计算和全局通信带来了重大挑战。

近年来，许多工作 [Gao (2022)](https://arxiv.org/html/2609.16350#bib.bib7); [Tarzanagh et al. (2022)](https://arxiv.org/html/2609.16350#bib.bib4); [Li et al. (2024)](https://arxiv.org/html/2609.16350#bib.bib5); [Yang et al. (2024)](https://arxiv.org/html/2609.16350#bib.bib24); [Huang et al. (2023)](https://arxiv.org/html/2609.16350#bib.bib6) 致力于解决FedSBO问题中的独特挑战。例如，[Gao (2022)](https://arxiv.org/html/2609.16350#bib.bib7) 在同构设置下开发了一种带有动量的本地随机双层梯度下降算法，由于数据分布的同质性，可以消除对全局雅可比矩阵和海森矩阵的依赖。[Tarzanagh et al. (2022)](https://arxiv.org/html/2609.16350#bib.bib4) 在异构设置下提出了FedNEST。它在内部循环中采用Neumann级数展开方法来估计海森-逆-向量积。然而，这种方法通信复杂度较高，因为它需要在内部循环中通信每一个中间的海森-逆-向量积。[Li et al. (2024)](https://arxiv.org/html/2609.16350#bib.bib5) 提出了一种单循环算法来解决高通信复杂度的问题。特别是，它引入了一个额外的梯度下降过程来替代Neumann级数展开方法，以估计海森逆向量积。

然而，上述联邦随机双层优化算法仍然存在高计算成本。具体来说，这些算法需要计算二阶雅可比矩阵和海森矩阵。当问题维度较高时，这在计算上是昂贵的。实际上，为了避免计算二阶雅可比矩阵和海森矩阵，[Kwon et al. (2023)](https://arxiv.org/html/2609.16350#bib.bib1); [Shen and Chen (2023)](https://arxiv.org/html/2609.16350#bib.bib25); [Kwon et al. (2024)](https://arxiv.org/html/2609.16350#bib.bib26); [Chen et al. (2023)](https://arxiv.org/html/2609.16350#bib.bib2) 在单机环境下开发了一种全一阶方法。具体来说，[Kwon et al. (2023)](https://arxiv.org/html/2609.16350#bib.bib1) 将双层优化问题转化为单层问题，然后只需一阶梯度即可求解，显著降低了计算成本。更具体地说，[Kwon et al. (2023)](https://arxiv.org/html/2609.16350#bib.bib1) 将下层优化问题转化为约束，然后利用惩罚方法进一步将其转化为无约束的minimax优化问题，该问题可以通过一阶随机梯度下降上升算法求解。受此启发，我们旨在开发一种高效的联邦随机双层优化算法，该算法仅使用一阶梯度来应对上述挑战。然而，为公式 (1) 中的FedSBO问题开发一阶方法面临着独特的挑战，概述如下。

- • 首先，现有的单机环境下的第一阶方法 [Kwon et al. (2023)](https://arxiv.org/html/2609.16350#bib.bib1) 采用与迭代相关的学习率以保证收敛。然而，这种策略在联邦学习中并不实用。具体来说，根据与迭代相关的学习率，参与设备能够推断当前的训练阶段，攻击者可能利用这一点来攻击训练过程。因此，有必要开发一种恒定学习率，以避免此问题同时保证收敛。
- • 其次，现有的单机环境下的第一阶方法 [Kwon et al. (2023)](https://arxiv.org/html/2609.16350#bib.bib1) 采用双时间尺度学习率以保证收敛。具体来说，某些变量的学习率比其他变量大一个数量级。这种学习率对于实际的联邦学习从业者来说难以调整。因此，有必要提出一种单时间尺度学习率，以便于调整同时确保收敛。
- • 第三，单机环境下的第一阶方法需要一个惩罚超参数来保证收敛。然而，目前尚不清楚这个惩罚超参数如何影响联邦学习中的一致性误差。具体来说，尚不清楚惩罚超参数对一致性误差的影响是否会导致更慢的收敛速率。因此，有必要建立理论收敛速率，揭示惩罚超参数如何影响联邦学习中的一致性误差和收敛速率。

为了应对上述独特的挑战，我们开发了一种新颖的、完全依赖一阶预言机的联邦随机方差缩减双层梯度下降算法。具体来说，在算法设计方面，我们开发了一种新颖的单时间尺度恒定学习率，并展示了这种新学习率如何依赖于惩罚超参数。在理论分析方面，我们提出了一种新颖的策略来建立我们算法的收敛速率。特别是，我们为收敛分析开发了一种新颖的势函数，并展示了如何将不同组件与精心设计的系数结合在一起。凭借这些新颖的算法和理论设计，我们的算法可以达到 O⁡\(1Nε5\) O\(\\frac\{1\}\{N\\epsilon^\{5\}\}\) 的收敛速率以获得ε \(\\epsilon\) 精度解，这表明收敛速率相对于设备数量 N \(N\) 具有线性加速。最后，大量的实验证实了我们新算法的有效性。总而言之，我们的论文做出了以下贡献。

- • 我们开发了一种新颖的联邦随机方差缩减双层梯度下降算法，该算法使用完全一阶梯度和单时间尺度恒定学习率。据我们所知，这是首次将单时间尺度恒定学习率应用于一阶双层优化算法。
- • 我们建立了我们算法的收敛速率，其中我们展示了如何使用一个势函数将不同的估计误差与精心设计的系数结合在一起。据我们所知，这是首个为具有单时间尺度恒定学习率的一阶联邦双层优化算法提供收敛保证的工作。
- • 我们在各种任务上的广泛实验结果证实了我们算法的有效性，即我们的新算法比现有的二阶方法在计算上更高效。

## 2 相关工作

### 2.1 随机双层优化

由于上层损失函数依赖于下层优化问题的最优解，上层损失函数的超梯度需要计算∂y∗\(x\)/∂x \\partial y^\{\*\}\(x\)/\\partial x，这给优化此类优化问题带来了独特的挑战。具体来说，计算∂y∗\(x\)/∂x \\partial y^\{\*\}\(x\)/\\partial x 依赖于雅可比矩阵和海森逆矩阵。为了计算它们，最近在单机环境下已经进行了许多工作 [Ghadimi and Wang (2018)](https://arxiv.org/html/2609.16350#bib.bib10); [Ji et al. (2021)](https://arxiv.org/html/2609.16350#bib.bib18); [Chen et al. (2021)](https://arxiv.org/html/2609.16350#bib.bib19); [Hong et al. (2020)](https://arxiv.org/html/2609.16350#bib.bib17); [Khanduri et al. (2021)](https://arxiv.org/html/2609.16350#bib.bib15); [Yang et al. (2021)](https://arxiv.org/html/2609.16350#bib.bib11); [Guo and Yang (2021)](https://arxiv.org/html/2609.16350#bib.bib12); [Dagréou et al. (2022)](https://arxiv.org/html/2609.16350#bib.bib9); [Chu et al. (2024)](https://arxiv.org/html/2609.16350#bib.bib22); [Dagréou et al. (2024)](https://arxiv.org/html/2609.16350#bib.bib23)。例如，[Ghadimi and Wang (2018)](https://arxiv.org/html/2609.16350#bib.bib10) 提出使用Neumann级数展开方法来近似海森-逆-向量积，并在此基础上建立了随机梯度下降的收敛速率。在此方向上，已经开发了几种改进算法。例如，[Ji et al. (2021)](https://arxiv.org/html/2609.16350#bib.bib18) 开发了一种小批量随机梯度下降方法，通过大批量大小提高了收敛速率。[Hong et al. (2020)](https://arxiv.org/html/2609.16350#bib.bib17) 研究了使用双时间尺度学习率的随机梯度下降的收敛速率。[Chen et al. (2021)](https://arxiv.org/html/2609.16350#bib.bib19) 建立了具有交替更新策略的随机梯度下降的收敛速率。为了进一步提高收敛速率，[Khanduri et al. (2021)](https://arxiv.org/html/2609.16350#bib.bib15); [Yang et al. (2021)](https://arxiv.org/html/2609.16350#bib.bib11); [Guo and Yang (2021)](https://arxiv.org/html/2609.16350#bib.bib12) 将方差缩减技术 [Cutkosky and Orabona (2019)](https://arxiv.org/html/2609.16350#bib.bib13); [Fang et al. (2018)](https://arxiv.org/html/2609.16350#bib.bib14); [Nguyen et al. (2017)](https://arxiv.org/html/2609.16350#bib.bib21) 引入随机双层优化，其收敛速率 O⁡\(ε−3\) O\(\\epsilon^\{\-3\}\) 可以匹配单层优化问题的对应速率。除了Neumann级数展开方法外，还有另一种估计超梯度的方法。具体来说，海森-逆-向量积被视为二次优化问题的最优解，然后引入一个额外的梯度下降过程来估计海森-逆-向量积。基于此方法，已经开发了一些算法。例如，[Dagréou et al. (2022)](https://arxiv.org/html/2609.16350#bib.bib9); [Chu et al. (2024)](https://arxiv.org/html/2609.16350#bib.bib22); [Dagréou et al. (2024)](https://arxiv.org/html/2609.16350#bib.bib23) 将其与方差缩减技术相结合，其收敛速率 O⁡\(ε−3\) O\(\\epsilon^\{\-3\}\) 也可以匹配单层优化问题的对应速率。

由于上述方法需要计算二阶雅可比矩阵和海森逆矩阵，[Kwon et al. (2023)](https://arxiv.org/html/2609.16350#bib.bib1); [Shen and Chen (2023)](https://arxiv.org/ht