【联邦学习论文阅读】FedProx（2018）Federated Optimization in Heterogeneous Networks
【FedProx】论文链接：https://arxiv.org/abs/1812.06127
摘要
联邦学习面临两个关键挑战：系统异构性；统计异构性。本文的FedProx可以解决联邦学习的异构性，可看作FedAvg的泛化和再参数化。理论上保证了当数据为非独立同分布时模型的收敛性，同时通过允许每个参与设备执行可变的工作量，保证了设备级别的系统约束。实验证明，广义FedProx框架相对于FedAvg，在异质网络中更具稳健性和稳定性。
一、介绍
为了应对异构性和高通信成本，允许局部更新和低参与度的优化方法已经成为联邦学习的流行方法（McMahan等人FedAvg，2017；Smith等人，2017）。虽然FedAvg可用于解决异构性问题，但它不允许参与设备根据底层系统约束执行可变的本地计算，而是简单丢弃无法完成指定计算轮数的设备。
在联邦学习中，系统异构性和统计异构性相互作用，无论是丢弃落后者，还是合并来自落后者的部分信息，都会隐式地增加统计异构性，不利于收敛。为缓解这一问题，在目标函数中增加一个近端项来提高稳定性。
二、相关工作
联邦学习需要解决隐私、异构数据和设备以及大规模分布式计算网络的新挑战
- 过去十年中，大规模机器学习（特别是在数据中心设置下）的研究，推动了分布式优化方法的发展。然而随着计算基元的能力和普及率的增长，直接通过分布式设备网络学习统计模型变得越来越有吸引力。
- 针对联邦环境中的挑战提出的优化方法，相比传统的分布式方法（如ADMM或mini-batch方法）有了显著改进。它们允许不精确的局部更新以平衡大型网络中的通信和计算，并且允许在任何通信回合中仅使用部分设备。
如Smith等人（2017）提出的一种通信效率高的原始-对偶优化方法，通过多任务学习框架学习每个设备的独立但相关的模型，尽管该方法具有理论上的保证和实际的有效性，但这种方法不能推广到非凸问题，例如深度学习（不再保证强对偶性）；在非凸情况下的FedAvg
FedAvg是一种基于原始数据的平均局部随机梯度下降进行更新的启发式方法
- 它在实际中表现中确实有效，但由于其本地更新方案、每轮使用的设备很少且数据经常以异构形式分布的问题，其分析非常具有挑战性。
- 最近的工作已经朝着在更简单的非联邦设置中分析FedAvg迈出了一步（如在IID设置中研究的并行SGD和相关变体），它们在IID设置下研究了类似于FedAvg的本地更新。但结果依赖于“每个局部求解器是同一随机过程的副本”这一前提，不适用于异构环境。
解决统计异构性的工作：
- 一些工作研究了统计异构环境下的收敛保证，但提出了限制性假设，即所有设备都参与每一轮通信，这在现实的联邦网络中通常是不可行的。与本文中提出的求解器无关框架相比，它们要求每个设备上使用特定求解器（SGD或GD），并且在分析中增加了额外的凸性假设或一致有界梯度假设。
- 也有一些方法旨在通过共享本地设备数据或一些服务器端代理数据来解决统计异构性问题。然而，这些方法向服务器发送本地数据，不仅增加了网络带宽负担，还违背了联邦学习的密钥隐私假设；需向所有设备发送全局共享代理数据，而生成或收集这样的辅助数据较为困难。
系统异构性：
- 联邦网络中每个设备的存储、计算和通信能力可能因硬件、网络连接和电力的可变性而有所不同，这加剧了减少掉队者和容错等挑战。
- 简单忽略掉队者可能对收敛产生负面影响；若掉队设备有特定数据特征，还会导致设备采样偏差。
受FedAvg启发，探索了一个更广泛的框架FedProx：
- 能够处理异构联邦数据，同时保持类似的隐私和计算优势；
- 分析了框架的收敛性，不仅考虑局部函数间存在统计异构表征，还考虑到实际的系统约束；
- 对统计异构的处理方法，受到求解线性方程组的随机化Kaczmarz方法的启发，该方法的类似假设已被用于分析其他情况下的SGD变种；
- 所提框架在异构联邦网络中有更好的鲁棒性和稳定性。
本文工作中的近端项，和分析中使用的有界不同假设，在其他文献中已经被研究过（动机不同），在附录B中进一步说明了相关背景工作。
三、联邦优化：算法
联邦学习方法用于处理收集数据的多个设备和协调整个网络全局学习目标的中央服务器，目标是最小化：
min w f ( w ) = ∑ k = 1 N p k F k ( w ) = E k [ F k ( w ) ] \min _{w} f(w)=\sum_{k=1}^{N} p_{k} F_{k}(w)=\mathbb{E}_{k}\left[F_{k}(w)\right] wminf(w)=k=1∑NpkFk(w)=Ek[Fk(w)]
N是设备数量；pk≥0且和为1（系数），可设置为nk/n（nk为设备k上的样本数量，n为总样本数量）； F k ( w ) : = E x k ∼ D k [ f k ( w ; x k ) ] F_{k}(w):=\mathbb{E}_{x_{k} \sim \mathcal{D}_{k}}\left[f_{k}\left(w ; x_{k}\right)\right] Fk(w):=Exk∼Dk[fk(w;xk)] 为局部经验风险，Dk指数据分布（不同设备上可能不同），本文考虑Fk(w)可能为非凸。
减少通信的一种常用技术是，在每个设备上使用基于设备数据的局部目标函数作为全局目标函数的代理。在每次外部迭代中，选择一个设备子集，并使用局部求解器优化每个选定设备上的局部目标函数。然后，这些设备将其局部模型更新传递给中央服务器，中央服务器聚合并更新全局模型。允许灵活性能的关键是可以不精确地解决每个局部目标，考虑根据执行的局部迭代次数调整局部计算量与通信量。
定义1（ γ γ γ-不精确解）：有函数 h ( w ; w 0 ) = F ( w ) + μ 2 ∥ w − w 0 ∥ 2 h\left(w ; w_{0}\right)=F(w)+\frac{\mu}{2}\left\|w-w_{0}\right\|^{2} h(w;w0)=F(w)+2μ∥w−w0∥2（则 ∇ h ( w ; w 0 ) = ∇ F ( w ) + µ ( w − w 0 ) ∇h(w; w_0) = ∇F(w) + µ(w − w_0) ∇h(w;w0)=∇F(w)+µ(w−w0)），
若 ∥ ∇ h ( w ∗ ; w 0 ) ∥ ≤ γ ∥ ∇ h ( w 0 ; w 0 ) ∥ \left\|\nabla h\left(w^{*} ; w_{0}\right)\right\| \leq \gamma\left\|\nabla h\left(w_{0} ; w_{0}\right)\right\| ∥∇h(w∗;w0)∥≤γ∥∇h(w0;w0)∥， γ ∈ [ 0 , 1 ] γ∈[0, 1] γ∈[0,1]，则 w ∗ w^∗ w∗ 是 m i n w h ( w ; w 0 ) \underset{w}{min}\ h(w;w_0) wmin h(w;w0) 的γ-不精确解。
3.1 FedAvg
具体可见【联邦学习论文阅读】FedAvg（2016）Communication-Efficient Learning of Deep Networks from Decentralized Data
FedAvg中超参数的设置很重要，特别是局部更新次数 E：执行更多的局部更新会增加局部计算量并减少通信，从而提高通信受限网络的整体收敛速度；但对异构的局部目标函数，较多局部更新不但会影响收敛，甚至偏离全局目标，还可能使更多的设备无法在给定时间内完成要求次数的更新。
随着本地数据和可用系统资源的变化，最佳局部更新次数会有所不同，需要增加收敛鲁棒性，且尽可能增大局部更新次数。据此提出FedProx，形式化表示局部更新次数随网络特征的变化。
3.2 FedProx
在FedAvg（在每一轮中选择一个设备子集，执行局部更新，然后对这些更新进行平均以形成全局更新）的基础上做了以下更改：
容忍部分工作
允许局部设备根据其可用系统资源，执行可变数量的工作（而不是丢弃落后设备），即为不同设备和不同迭代，提供可变的γ，并由定义1扩展为定义2。
定义2（ γ k t γ_k^t γkt-不精确解）：有函数 h k ( w ; w t ) = F k ( w ) + μ 2 ∥ w − w t ∥ 2 h_k\left(w ; w_{t}\right)=F_k(w)+\frac{\mu}{2}\left\|w-w_{t}\right\|^{2} hk(w;wt)=Fk(w)+2μ∥w−wt∥2（则 ∇ h k ( w ; w t ) = ∇ F k ( w ) + µ ( w − w t ) ∇h_k(w; w_t) = ∇F_k(w) + µ(w − w_t) ∇hk(w;wt)=∇Fk(w)+µ(w−wt)）），
若 ∥ ∇ h k ( w ∗ ; w t ) ∥ ≤ γ k t ∥ ∇ h k ( w t ; w t ) ∥ \left\|\nabla h_k\left(w^{*} ; w_{t}\right)\right\| \leq \gamma_k^t\left\|\nabla h_k\left(w_{t} ; w_{t}\right)\right\| ∥∇hk(w∗;wt)∥≤γkt∥∇hk(wt;wt)∥， γ ∈ [ 0 , 1 ] γ∈[0, 1] γ∈[0,1]，则 w ∗ w^∗ w∗ 是 m i n w h k ( w ; w t ) \underset{w}{min}\ h_k(w;w_t) wmin hk(w;wt) 的 γ k t γ_k^t γkt-不精确解。
近端项
允许跨设备执行可变工作量可以缓解系统异构的影响，但因为底层数据异构，过多的局部更新仍会导致分歧。为限制局部更新的影响，在局部问题中加入近似项，则改写局部函数 F k ( ⋅ ) F_k(·) Fk(⋅) 为：
min w h k ( w ; w t ) = F k ( w ) + μ 2 ∥ w − w t ∥ 2 \min _{w} h_{k}\left(w ; w^{t}\right)=F_{k}(w)+\frac{\mu}{2}\left\|w-w^{t}\right\|^{2} wminhk(w;wt)=Fk(w)+2μ∥∥w−wt∥∥2
有两个优点：限制局部更新，使其更接近初始化的全局模型，且不用手动设置局部更新次数，缓解了统计异构性问题；可以和针对系统异构性的可变工作量方法合并。
近端项与之前工作的不同在于：
- 用于处理联邦网络中的异构性
- 可用于分布式设置中：非独立同分布数据；任意局部求解器；跨设备不精确更新；每轮选择活动设备子集
FedProx步骤总结如下：
FedProx更为通用，FedAvg可看作FedProx的一个特例（μ=0，局部求解器为SGD，γ固定）
四、收敛性分析
近端项的使用使得FedProx更适合于理论分析（即局部目标的表现可能更好）。特别是，如果相应地选择μ，hk的海森矩阵可能是半正定的。因此，当Fk为非凸时，hk为凸；当Fk为凸时，hk为μ-强凸。
提出一个专门衡量局部函数之间差异性的度量标准（4.1），并在此假设下分析FedProx（4.2）
- FedAvg和FedProx本质上是随机算法：在每一轮中，只随机选择小部分设备来执行更新，并且在各设备上执行的更新可能不精确。
- 为了使随机方法收敛到一个平稳点，需要减小步长；而非随机方法（如梯度下降法）可以采用恒定步长。
- 为了分析常步长方法的收敛性，需要量化局部目标函数之间的差异程度。这可以通过假设数据是IID来实现（即跨设备的同构性），但在现实的联邦网络中，这种假设是不切实际的。
4.1 客户端异构
引入一个衡量联邦网络中设备之间差异性的方法，来证明收敛性；也可以通过更简单的梯度有界方差假设（推论8）来满足（见后续实验）。
定义3（B-局部差异）：若 E k [ ∥ ▽ F k ( w ) ∥ 2 ] ≤ ∥ ▽ f ( w ) ∥ 2 B 2 \mathbb{E}_k[\left\| \bigtriangledown F_k(w)\right\|^2]\leq \left\|\bigtriangledown f(w) \right\|^2B^2 Ek[∥▽Fk(w)∥2]≤∥▽f(w)∥2B2，则称 F k F_k Fk在w为B-局部差异，
进一步可得 B ( w ) = E k [ ∥ ▽ F k ( w ) ∥ 2 ] ∥ ▽ f ( w ) ∥ 2 B(w)=\sqrt{\frac{\mathbb{E}_k[\left\| \bigtriangledown F_k(w)\right\|^2]}{\left\|\bigtriangledown f(w) \right\|^2}} B(w)=∥▽f(w)∥2Ek[∥▽Fk(w)∥2]， ∥ ▽ f ( w ) ∥ ≠ 0 \left\|\bigtriangledown f(w) \right\|\neq 0 ∥▽f(w)∥=0
定义3允许统计异构性，为IID假设下有界差异的推广。当所有局部函数相同时，B(w)=1，其余B>1。B(w)越大，局部函数越不相似。要求定义3中的差异有界，可得到收敛性分析用到的正式差异假设。
假设1：有界差异
精度过高会导致过拟合；虽然联邦学习不是IID样本，但并非完全无关。据此合理假设局部函数间的差异有界。
4.2 FedProx分析
定理4：非凸FedProx收敛性（B-局部差异）
由执行一次FedProx，目标值的预期减少量（根据假设1），可进一步推导出收敛速度。使用定义3中的差异性，来确定每次迭代时目标值足够减小。
备注5：定理4中参数设置及原因
定理6：FedProx的收敛速度
推论7：凸情况下的收敛性
精确度增高时，增大µ会使收敛更稳定；可推得实现某精度需要的步骤数，借此分析局部函数不同时，FedProx和类似方法的性能。
备注8：与SGD比较
推论9：应用可变γ的收敛性
前面的分析假设γ不变，这里考虑系统异构性，扩展为γ可变。
五、实验
所有代码、数据和实验都可以在github.com/litian96/FedProx上获取。
5.1 实验细节
- 合成数据：采用与Shamir等人(2014)类似的设置，并增加设备异构性（α控制局部模型间差异，β控制局部数据间差异）。
- 真实数据：如下表所示，前两个为凸，后两个为非凸。
- 实现：在Tensorflow中实现FedAvg和FedProx（也采用SGD），抽样方法不同且发现本文提出的会使性能更稳定。
- 超参数和评估指标：在所有数据集上使用相同的学习率，每次选择10个设备；对比时，固定所有运行中随机选择的设备、落后者和小批量顺序；基于全局目标函数、通信回合数
5.2 系统异构：容忍部分工作
系统异构性模拟：
- 各设备根据全局时钟和系统约束确定局部工作量（对应 γ k t γ_k^t γkt）；
- 模拟时设置全局局部更新次数E，并使一些设备执行比E少的更新，即在[1,E]中随机选择局部更新次数，分配给所选设备的0%、50%、90%（0%时，没有设备执行比E少的更新，没有系统异构；90%时系统高度异构）
- 用于对比的FedAvg直接删除0%、50%、90%的落后设备，而FedProx聚合所有所选设备（部分）的更新。
设置E=20，synthetic(1,1)，μ=0且0%落后设备的FedProx等同于FedAvg；
对比FedAvg和μ=0的FedProx：发现系统异构性越大，收敛性越差，与放弃受限设备相比，加入可变工作量有利于更快更稳定的收敛；
对比μ=0和μ>0的FedProx：(表明增加近似项的优势）发现不论系统异构性如何，适当的μ可以提高稳定性，使发散的方法收敛；在大多数情况下μ>0的准确性更高（特别是高度异构环境中）
对5个数据集，最佳μ值分别为1、1、1、0.001和0.01
还设置了两个不太异构的环境：
- E最大等于1，其他设置类似，发现允许部分工作仍可提高收敛性；
- 使用没有统计异构性的合成数据，发现FedAvg很健壮且允许部分工作未导致重大改进。
5.3 统计异构：近端项
随着统计异构性的增加，收敛更困难。
第一行：在不存在系统异构的情况下，统计异构性对收敛的影响（固定E=20）；随着数据异构性的增加，μ=0（FedAvg）收敛性变差，而μ>0可帮助解决此问题；
第二行：使用B-局部差异度量方法（可衡量统计异构性），观察每个设备上的梯度变化；增加近端项会减小局部函数间差异性；差异度量与训练损失趋势一致，异构性越小，收敛性越好，可通过设置μ来改进收敛性。
影响性能的关键参数为E和μ：增大E会使局部模型偏离全局初始点过远，导致分歧；而合适的μ可以约束局部更新轨迹，使之更接近全局模型，保证收敛性。
大μ会减慢收敛速度，小μ可能不会产生任何影响；
自适应选择μ：连续5轮损失增加/减少时，设置μ增加/减少0.1；对于完全IID的数据，初始化μ=1，对synthetic(1,1)，初始化μ=0
六、结论
FedProx允许跨设备执行可变数量的工作，并通过增加近端项（限制本地训练时模型对全局模型的偏离）来稳定该方法。个人理解是稳定通用，但存在收敛速度慢的缺点。