本文章是对论文链接: FedADMM:A Robust Federated Deep Learning Framework with Adaptivity to System Heterogeneity作的笔记。
前言
在今天,数据以分布式的方式大量生成,数据的隐私性的保护尤为重要,由此联邦学习(FL)应运而生。
尽管FedAvg是目前最流行的FL算法,但是它没有考虑系统异质性和统计异质性,具体来说,没有解决异构网络的掉线问题(鲁棒性),也不适用于非IID数据分布情况。
本文章提出了一种基于原始-对偶算法的新的FL算法FedADMM,能够实现以下功能(有以下特点):
- 通过使用局部增广拉格朗日和对偶变量,FedADMM同时解决了系统和统计异质性问题;
- 通过引入一般的服务器更新步长,实现收敛速度和抗震荡性间的良好平衡;
- 在比现有算法限制更小的条件下,实现算法收敛速度为 O ( 1 T ) + O ( 1 ) \mathcal{O}(\frac{1}{T})+\mathcal{O}(1) O(T1)+O(1),其中 T T T表示迭代轮数,这个收敛速度在较宽松的假设下达到了与现有技术相当的最佳复杂度下界;
- 还有一个重要特点是对客户端的激活方案没有严格限制,只需确保客户端在足够多的轮次中至少参与一次即可;
- 在IID与非IID数据集试验下,证明FedADMM在更少的轮次下获得了更高的准确度,在数据异构分布的大规模网络中最为明显。
现有的基于原始-对偶算法的FL算法有FedHybrid和FedPD。
- FedHybrid是一个同步算法,要求客户端在每一轮次都更新。将客户端分为两组来处理异质性:一组进行梯度更新,一组进行牛顿更新。由于牛顿法不适用于大规模神经网络,且收敛性是在强凸假设下成立的,因此该算法并不能完全适用于FL。
- FedPD使用梯度下降(GD)或随机梯度下降(SGD)进行训练,并使用对偶变量来捕捉局部模型和全局模型之间的差异。虽然该算法允许客户端的工作量可变,但是要求客户端与服务器通信频率是固定的,且要求所有客户端都参与(即所有客户端同时与服务器通信),这不仅带来了大量的通信开销,且对网络的性能要求更高,鲁棒性差。
相比之下FedADMM:
- 允许部分客户端参与通信,而无需全部客户端每个轮次都参与;
- 每个轮次都更新全局模型,而不仅仅等待所有客户端都完成更新;
- 全局模型更新规则中包含跟踪机制,有助于加快收敛速度并平滑路径。
算法内容
为更好适应联邦学习的设置,原始的损失最小化问题被重写为一个共识形式的问题:
min
w
i
,
θ
∈
R
d
∑
i
=
1
m
f
i
(
w
i
)
,
s.t.
w
i
=
θ
,
∀
i
∈
[
m
]
.
\mathop{\text{min}}\limits_{w_i,\theta\in\mathbb{R}^d} \sum_{i=1}^mf_i(w_i),\quad \text{s.t.}\quad w_i=\theta,\quad \forall i\in[m].
wi,θ∈Rdmini=1∑mfi(wi),s.t.wi=θ,∀i∈[m].
其中,
w
i
w_i
wi可解释为客户端
i
i
i的本地局部模型,而
θ
\theta
θ是服务器持有的全局模型,
f
i
(
w
i
)
f_i(w_i)
fi(wi)用来捕获加权的局部训练损失。
服务器端算法流程
- 在每一轮的开始(这里记为第 t t t轮),服务器选择一个客户端子集 S t S^t St并将全局模型 θ t \theta^t θt发给这些客户端;
- 接收到全局模型后,客户端执行本地训练并更新它们的模型参数 w i w_i wi和对偶变量 y i y_i yi;
- 客户端将更新后的模型差分 Δ i t \Delta ^t_i Δit发回到服务器;
- 服务器聚合从所选择的客户端接收到的消息,更新全局模型: θ t + 1 ← θ t + η ∣ S t ∣ ∑ i ∈ S t Δ i t \theta^{t+1}\leftarrow\theta^t+\frac\eta{|S^t|}\sum_{i\in S^t}\Delta_i^t θt+1←θt+∣St∣ηi∈St∑Δit其中 η \eta η为学习(速)率。
客户端更新
- 输入本地迭代次数 E i E_i Ei和学习率 η i \eta_i ηi;
- 创建数据批次 B \mathcal{B} B;
- 对每个批次
b
b
b计算批次梯度
∇
f
i
(
w
i
,
b
)
\nabla f_i(w_i,b)
∇fi(wi,b)并更新模型参数
w
i
w_i
wi:
w i ← w i − η i ( ∇ f i ( w i , b ) + y i + ρ ( w i − θ ) ) w_i\leftarrow w_i-\eta_i\big(\nabla f_i(w_i,b)+y_i+\rho(w_i-\theta)\big) wi←wi−ηi(∇fi(wi,b)+yi+ρ(wi−θ)) - 更新对偶变量
y
i
y_i
yi:
y i ← y i + ρ ( w i − θ ) y_i\leftarrow y_i+\rho(w_i-\theta) yi←yi+ρ(wi−θ)
算法全部内容:
对偶变量的作用
对偶变量 y i y_i yi不仅用于惩罚局部模型与全局模型之间的差异,还提供了量化衡量让本地局部模型与全局模型不同的好处。在训练初期,让局部模型与全局模型有差异是有好处的,因为不同客户端数据存在异质性,对偶变量帮助捕捉了这种异质性,使FedADMM能够自适应数据的变化,而无需进行超参数调整。
算法的理论分析和实践验证
文章对FedADMM算法的收敛性分析是建立在任意客户端参与的条件下的,且允许只有部分客户端参与每次的通信轮次,因此可以说明FedADMM能够适应联邦学习环境的常见的系统异质性问题,特别是客户端参与的不确定性。
文章同时证明了,在已有的方法中,只有FedProx和FedPD与FedADMM一样,在解决方案准确性上的依赖性是相同的,即它们都具有
O
(
1
/
ε
)
\mathcal{O}(1/\varepsilon)
O(1/ε)的收敛率。但是另两个算法都有各自的缺陷FedProx要求选择的活跃客户端数量
S
S
S必须大于数据差异度量
B
2
B^2
B2,而FedPD则要求所有客户端在同一时间通信。FedADMM没有这两方面的问题,便有更广的应用性。
文章中,参与更新
w
i
w_i
wi和
y
i
y_i
yi的系数
ρ
\rho
ρ会影响算法的收敛速度和稳定性,而FedADMM对
ρ
\rho
ρ的选择可以在没有任何关于系统规模、参与客户端数量和客户端统计差异的先验知识的情况下进行。这使得FedADMM在实践中非常实用。
经过实验验证,得出:
- FedADMM在IID和非IID环境中都能在更少的轮次中获得很高的精度,这意味着,可以节省大量通信开销;
- FedADMM对统计和系统异质性都表现出了很好的鲁棒性,并且不需要进行超参数调整;
- 相对于已有方法,FedADMM在所有测试场景下,均是最优的,尤其是在处理大规模系统重尤为明显。