ICML2023论文学习_Dynamic Regularized SAM in FL: Approaching Global Consistency and Smooth Landscape

文章链接：Link

• 摘要：在联邦学习(FL)中，由于多次本地更新和non-i.i.d.数据集，客户端容易过度拟合到局部最优，偏离全局目标，影响性能。以往的研究大多从优化的角度出发，只注重增强局部目标与全局目标的一致性(FedLin, FedDyn…)，以缓解client drift。文章提出了一种新的通用算法FedSMOO，将优化目标和泛化目标结合起来，有效地提高了FL的性能。FedSMOO采用动态正则化器来保证局部最优点向全局目标靠近，同时通过全局锐度感知最小化(sharp Aware Minimization, SAM)优化器对其进行修正，以搜索一致的平坦最小值。理论分析表明，FedSMOO具有较快的$\mathcal{O}(\frac{1}{T})$收敛速度和较低的泛化界。

• Main Contributions:
  (1) 相较于FedSAM只寻找local flatness，FedSMOO通过迭代更新局部扰动$s_i$促进global generality
  (2) 类似于FedDyn，增加了一个正则化项来保证global consistency

• Motivation:
  (1) Sharp Aware Minimization, SAM: 旨在优化一个极小-极大问题
  ${\min }_w\left\{f_s(w)\triangleq {\max }_{\Vert s\Vert \le r}f(w+s)\right\}$
  SAM采用近似方式求解此问题，通过在$w$处进行一阶泰勒展开，此时$w$处最优的$s$取值易得
  $s^{\ast }(w)\approx \underset{\Vert s\Vert \le r}{\mathrm{argmax}}\left\{f(w)+s^{\top }\nabla f(w)\right\}=r\cdot \nabla f(w)/\Vert \nabla f(w)\Vert$
  故SAM可被视作用于寻找当前损失"flat landscape"的基于GD/SGD的改进optimizer，在集中式设置下基于SGD的更新流程如下($t$-round)：
  Perform a SGD step with batch ${\xi }_t$ – Find the parameter $w+r\frac{\nabla f(w)}{\Vert \nabla f(w)\Vert }$ – Calculate the SGD gradient then update

  (2) SAM拓展到联邦设置(FedSAM)时存在的问题
  当拓展到联邦设置下，由于non-i.i.d.数据，局部的pertubation的聚合并不等于全局(基于所有设备本地数据分布)的pertubation，即局部的flat minima聚合后不一定对于全局损失的flat minima，使得模型聚合后得到的并不一定对应全局平坦值，泛化能力的提升是不确定的
  $\frac{1}{m}{\sum }_{i=1}^m{s}_i\ne s$

• Methodology
  1.为解决global flat landscape问题，即保证local pertubation趋向于global pertubation，文章首先重构优化问题：（这一步相当于解决inner level的极大化问题，即促进global generality）
  $\begin{array}{r}{\min }_w\left\{\mathcal{F}(w)=\frac{1}{m}{\sum }_{i\in [m]}{\mathcal{F}}_i(w)\right\},\\ {\mathcal{F}}_i(w)\triangleq {\max }_{\Vert s\Vert \le r}{f}_i(w+s),\end{array}$
  引入约束，上述问题可进一步重写为
  min ⁡ w i = w { F = 1 m ∑ i ∈ [ m ] F i ( w , w i , s , s i ) } , F i ( w , w i , s , s i ) ≜ max ⁡ ∥ s i ∣ ≤ r , s i = s f i ( w i + s i ) . \begin{aligned} & \min _{w_i=w}\{\mathcal{F}\left.=\frac{1}{m} \sum_{i \in[m]} \mathcal{F}_i\left(w, w_i, s, s_i\right)\right\}, \\ & \mathcal{F}_i\left(w, w_i, s, s_i\right) \triangleq \max _{\substack{\| s_i \mid \leq r, s_i=s}} f_i\left(w_i+s_i\right) . \end{aligned} ​wi​=wmin​{F=m1​i∈[m]∑​Fi​(w,wi​,s,si​)⎭ ⎬ ⎫​,Fi​(w,wi​,s,si​)≜∥si​∣≤r,si​=s​max​fi​(wi​+si​).​
  为求解此问题，首先采用同SAM一致的一阶泰勒展开，随后定义增广拉格朗日函数，通过一阶和二阶项惩罚$s_i=s$
  ${\mathcal{L}}_i^s:f_i\left(w_i\right)+s_i^{\top }\nabla f_i\left(w_i\right)+{\mu }_i^{\top }\left(s-s_i\right)+\frac{1}{2\alpha }{\Vert s-s_i\Vert }^2,s.t.\Vert s_i\Vert \le r$
  即定义了一个modified local loss function，优化变量变为 { s i , μ i , s } \{ s_i, \mu_i, s \} {si​,μi​,s}，随后采用ADMM算法进行交替更新：
  $\begin{array}{rl}{\hat{s}}_i& =\underset{\Vert s_i\Vert \le r}{\mathrm{argmax}}\left\{s_i^{\top }\left(\nabla f_i\left(w_i\right)-{\mu }_i\right)+\frac{1}{2\alpha }{\Vert s_i-s\Vert }^2\right\}\\ & =\underset{\Vert s_i\Vert \le r}{\mathrm{argmax}}\left\{\frac{1}{2\alpha }{\Vert s_i+{\bar{s}}_i\Vert }^2\right\},\end{array}$
  ${\bar{s}}_i=\alpha \left(\nabla f_i\left(w_i\right)-{\mu }_i\right)-s,{\hat{s}}_i=r{\bar{s}}_i/\Vert {\bar{s}}_i\Vert ,{\mu }_i={\mu }_i+\frac{1}{\alpha }\left({\hat{s}}_i-s\right)$
  同时，模型聚合得到相应的modified global loss function也通过此方式优化
  $\begin{array}{rl}\hat{s}& =\underset{\Vert s\Vert \le r}{\mathrm{argmax}}\frac{1}{m}\sum _{i\in [m]}\left\{s^{\top }{\mu }_i+\frac{1}{2\alpha }{\Vert s-{\hat{s}}_i\Vert }^2\right\}\\ & =\underset{\Vert s\Vert \le r}{\mathrm{argmax}}\left\{\frac{1}{2\alpha }\frac{1}{m}\sum _{i\in [m]}{\Vert s+\alpha {\mu }_i-{\hat{s}}_i\Vert }^2\right\}.\end{array}$
  $s=\frac{1}{m}{\sum }_{i\in [m]}\left(\alpha {\mu }_i-{\hat{s}}_i\right),\hat{s}=rs/\Vert s\Vert$
  2.为解决client drift问题，文章采用FedDyn方式引入动态正则化项：（这一步相当于解决outer level的极小化问题，即促进global consistency）
  $\mathcal{L}:\frac{1}{m}{\sum }_i\left\{{\mathcal{F}}_i+{\lambda }_i^{\top }\left(w^t-w_i\right)+\frac{1}{2\beta }{\Vert w^t-w_i\Vert }^2\right\}$
  将上述问题分布到client端逐轮求解，即
  $w_{i,K}^t=\underset{w_i}{\mathrm{argmin}}\left\{{\mathcal{F}}_i-{\lambda }_i^{\top }{w}_i+\frac{1}{2\beta }{\Vert w^t-w_i\Vert }^2\right\}$
  其中${\lambda }_i={\lambda }_i-\frac{1}{\beta }\left(w_{i,K}^t-w^t\right)$
  FedDyn链接：Link

• Algorithm Flow
  

• Summary&Rethinking

  1. 事实上$\begin{array}{r}\frac{1}{m}{\sum }_{i=1}^m\nabla f_i\left(w_t,{\xi }_i\right)\ne \nabla f\left(w_t,\xi \right)where{\xi }_i\sim D_i\xi \sim {\sum }_{i=1}^M{D}_i\end{array}$，离global越近一定越好吗
  2. 现在Stochastic Non-convex setting不推出来个$\mathcal{O}(1/T)$都不太好说自己Fast
  3. 从Experiment看，FedSMOO前期收敛较慢是符合直觉的
  4. FedDyn的引入的必要性没有被很好的展示