[论文解读] MARINA: Faster Non-Convex Distributed Learning with Compression

Eduard Gorbunov, Konstantin Burlachenko+2 more|arXiv (Cornell University)|Feb 15, 2021

Stochastic Gradient Optimization Techniques参考文献 45被引用 15

一句话总结

MARINA 提出了一种用于非凸分布式学习的通信高效一阶方法，采用梯度差值压缩技术，通过一种新颖的有偏梯度估计器，实现了更优的理论与实证性能。其通信复杂度降低至 $\mathcal{O}\left(\frac{1 + \omega / \sqrt{n}}{\varepsilon^2}\right)$，在寻找 $\varepsilon$-驻点方面，优于以往方法在通信与查询复杂度上的表现。

ABSTRACT

We develop and analyze MARINA: a new communication efficient method for non-convex distributed learning over heterogeneous datasets. MARINA employs a novel communication compression strategy based on the compression of gradient differences that is reminiscent of but different from the strategy employed in the DIANA method of Mishchenko et al. (2019). Unlike virtually all competing distributed first-order methods, including DIANA, ours is based on a carefully designed biased gradient estimator, which is the key to its superior theoretical and practical performance. The communication complexity bounds we prove for MARINA are evidently better than those of all previous first-order methods. Further, we develop and analyze two variants of MARINA: VR-MARINA and PP-MARINA. The first method is designed for the case when the local loss functions owned by clients are either of a finite sum or of an expectation form, and the second method allows for a partial participation of clients -- a feature important in federated learning. All our methods are superior to previous state-of-the-art methods in terms of oracle/communication complexity. Finally, we provide a convergence analysis of all methods for problems satisfying the Polyak-Lojasiewicz condition.

研究动机与目标

• 解决分布式非凸优化中的通信瓶颈，特别是在联邦学习与异构学习设置中。
• 开发一种在有效压缩下保持高收敛速度的同时最小化通信成本的方法。
• 通过引入专为压缩通信设计的有偏梯度估计器，克服现有方法的局限性。
• 与先前最先进方法相比，实现更优的理论通信与查询复杂度边界。
• 支持实际部署场景，如部分客户端参与以及有限和或期望形式的本地目标函数。

提出的方法

• 提出 MARINA，一种使用基于压缩梯度差值的有偏梯度估计器的一阶方法，与 DIANA 的无偏方法不同。
• 采用量化算子 $\mathcal{Q}$，其期望密度为 $\zeta_{\mathcal{Q}}$，在控制偏差的同时实现梯度压缩。
• 在 Polyak-Łojasiewicz (PL) 条件下推导收敛保证，确保非凸问题的线性收敛。
• 提出两种变体：VR-MARINA 用于有限和与期望形式的局部损失，PP-MARINA 用于联邦学习中的部分客户端参与。
• 使用李雅普诺夫函数 $\Phi_k = f(x^k) - f(x^*) + \frac{\gamma}{p}\|g^k - \nabla f(x^k)\|^2$ 分析收敛性并控制误差传播。
• 优化步长 $\gamma$ 以平衡收敛速度与误差累积，确保 $\mathbf{E}[\Phi_{k+1}] \leq (1 - \gamma\mu)\mathbf{E}[\Phi_k]$。

实验结果

研究问题

• RQ1 有偏梯度估计器是否能在非凸分布式学习中实现比无偏方法更优的通信效率？
• RQ2 在非凸分布式优化中，寻找 $\varepsilon$-驻点的最优通信复杂度是多少？
• RQ3 在压缩的联邦学习中，部分客户端参与如何影响收敛性与通信成本？
• RQ4 梯度差值压缩是否能实现优于标准梯度压缩的理论边界？
• RQ5 PL 条件是否能为压缩的分布式方法提供更紧致的收敛边界？

主要发现

• MARINA 实现了 $\mathcal{O}\left(\frac{1 + \omega / \sqrt{n}}{\varepsilon^2}\right)$ 的通信复杂度，优于 DIANA 的 $\mathcal{O}\left(\frac{1 + (1+\omega)\sqrt{\omega/n}}{\varepsilon^2}\right)$。
• VR-MARINA 通过在压缩更新中利用方差减少，降低了有限和与期望形式问题的查询复杂度。
• PP-MARINA 支持每轮采样 $r$ 个客户端，实现 $\mathcal{O}\left(dn + \max\left\{dn, \frac{L}{\mu}\left(\zeta_{\mathcal{Q}}r + \sqrt{(1+\omega)\zeta_{\mathcal{Q}}(dn - \zeta_{\mathcal{Q}}r)}\right)\right\} \log \frac{\Delta_0}{\varepsilon}\right)$ 的通信成本。
• 在 PL 条件下，MARINA 确保线性收敛，满足 $\mathbf{E}[f(x^K) - f(x^*)] \leq \varepsilon$，所需轮数为 $K = \mathcal{O}\left(\frac{dn}{\zeta_{\mathcal{Q}}r} \frac{L}{\mu} \left(1 + \sqrt{\frac{1+\omega}{r}\left(\frac{dn}{\zeta_{\mathcal{Q}}r} - 1\right)}\right)\right)$。
• 该方法的有偏估计器在压缩误差与收敛速度之间实现了更优的权衡，理论与实践中均优于无偏方法（如 DIANA 和 FedCOMGATE）。
• 实验结果证实，MARINA 在深度学习基准测试中展现出更高的通信效率与更快的收敛速度，尤其在高比例压缩下表现更优。