[论文解读] ProxSkip: Yes! Local Gradient Steps Provably Lead to Communication Acceleration! Finally!

Konstantin Mishchenko, Grigory Malinovsky+2 more|arXiv (Cornell University)|Feb 18, 2022

Stochastic Gradient Optimization Techniques被引用 10

一句话总结

ProxSkip 是一种新型优化方法，通过在大多数迭代中跳过昂贵的近端算子计算，可证明地加速分布式学习中的通信。它在 $\epsilon$-准确解下仅需 $\epsilon(\kappa\log(1/\epsilon))$ 次迭代，且仅需 $\epsilon(\sqrt{\kappa}\log(1/\epsilon))$ 次近端计算，显著降低了联邦学习和去中心化设置下的通信成本，且无需异构性假设。

ABSTRACT

We introduce ProxSkip -- a surprisingly simple and provably efficient method for minimizing the sum of a smooth ($f$) and an expensive nonsmooth proximable ($ψ$) function. The canonical approach to solving such problems is via the proximal gradient descent (ProxGD) algorithm, which is based on the evaluation of the gradient of $f$ and the prox operator of $ψ$ in each iteration. In this work we are specifically interested in the regime in which the evaluation of prox is costly relative to the evaluation of the gradient, which is the case in many applications. ProxSkip allows for the expensive prox operator to be skipped in most iterations: while its iteration complexity is $\mathcal{O}\left(κ\log \frac{1}{\varepsilon} ight)$, where $κ$ is the condition number of $f$, the number of prox evaluations is $\mathcal{O}\left(\sqrtκ \log \frac{1}{\varepsilon} ight)$ only. Our main motivation comes from federated learning, where evaluation of the gradient operator corresponds to taking a local GD step independently on all devices, and evaluation of prox corresponds to (expensive) communication in the form of gradient averaging. In this context, ProxSkip offers an effective acceleration of communication complexity. Unlike other local gradient-type methods, such as FedAvg, SCAFFOLD, S-Local-GD and FedLin, whose theoretical communication complexity is worse than, or at best matching, that of vanilla GD in the heterogeneous data regime, we obtain a provable and large improvement without any heterogeneity-bounding assumptions.

研究动机与目标

• 为解决分布式优化中近端算子带来的高通信成本，特别是在联邦学习中梯度平均操作代价高昂的问题。
• 开发一种在保持收敛保证的同时，大幅减少近端算子计算次数的方法。
• 在非同质数据环境下，实现对 FedAvg 和 SCAFFOLD 等现有本地梯度方法的通信复杂度的可证明改进。
• 提供一个理论坚实、可跳过近端步骤而不损失收敛速度或准确性的框架。

提出的方法

• 提出 ProxSkip，一种近端梯度下降的变体，允许在大多数迭代中跳过近端算子的计算。
• 采用随机机制：以概率 $p$ 执行近端步骤，否则跳过，并引入校正项以维持收敛性。
• 采用随机逼近框架，利用基于前一时刻迭代值的控制变量对更新方向进行校正。
• 推导出迭代序列的收敛速率 $\mathcal{O}(\kappa \log(1/\varepsilon))$，同时将近端计算次数保持在 $\mathcal{O}(\sqrt{\kappa} \log(1/\varepsilon))$。
• 将该方法应用于基于一致性约束的分布式问题，将问题重述为最小化 $f(x) + \psi(\mathbf{L}x)$，其中 $\psi$ 为一致性约束的指示函数。
• 推导出使用混合矩阵 $\mathbf{W}$ 的去中心化变体，其中通信对应于矩阵-向量乘法 $\mathbf{L}x$，并证明收敛速率为 $\tilde{\mathcal{O}}(\kappa + 1/(p^2\delta))$。

实验结果

研究问题

• RQ1 是否可以使用本地梯度步骤在不牺牲收敛保证的前提下降低分布式优化中的通信成本？
• RQ2 是否可以可证明地减少昂贵的近端算子计算次数，同时保持相同的收敛速率？
• RQ3 在非同质数据设置下，ProxSkip 是否相比 FedAvg 和 SCAFFOLD 实现了更好的通信复杂度？
• RQ4 在收敛速度方面，本地步数与通信频率之间的最优权衡是什么？
• RQ5 该方法能否扩展到使用一般混合矩阵的去中心化设置，并仍能实现最优收敛？

主要发现

• ProxSkip 仅需 $\mathcal{O}(\sqrt{\kappa} \log(1/\varepsilon))$ 次近端算子计算，即可实现 $\mathcal{O}(\kappa \log(1/\varepsilon))$ 的迭代复杂度，从而实现显著的通信节省。
• 该方法在非同质数据环境下，可证明地优于原始近端梯度下降法及其他本地方法（如 FedAvg 和 SCAFFOLD），通信复杂度得到显著改善。
• 在去中心化设置中，收敛速率为 $\tilde{\mathcal{O}}(\kappa + 1/(p^2\delta))$，最优概率 $p = 1/\sqrt{\delta\kappa}$ 可在本地步数与通信频率之间实现平衡。
• 当网络连接较差时（$\delta \leq 1/\kappa$），设置 $p=1$（即每轮都通信）为最优选择，与已知下界一致。
• 在 w8a 数据集上的逻辑回归实验表明，Scaffnew（ProxSkip 的变体）在使用理论超参数和最优 $p \approx 1/\sqrt{\kappa}$ 时，性能优于其他方法。
• 在随机设置下，随着客户端数量增加，该方法实现了线性加速，甚至超过理论预期。