[论文解读] ATOMO: Communication-efficient Learning via Atomic Sparsification

Hongyi Wang, Scott Sievert+4 more|arXiv (Cornell University)|Jun 11, 2018

Stochastic Gradient Optimization Techniques参考文献 5被引用 129

一句话总结

ATomo 引入一种通用框架，通过原子分解对随机梯度进行稀疏化，统一 QSGD 和 TernGrad，并展示在对神经网络梯度应用基于 SVD 的稀疏化时的加速。

ABSTRACT

Distributed model training suffers from communication overheads due to frequent gradient updates transmitted between compute nodes. To mitigate these overheads, several studies propose the use of sparsified stochastic gradients. We argue that these are facets of a general sparsification method that can operate on any possible atomic decomposition. Notable examples include element-wise, singular value, and Fourier decompositions. We present ATOMO, a general framework for atomic sparsification of stochastic gradients. Given a gradient, an atomic decomposition, and a sparsity budget, ATOMO gives a random unbiased sparsification of the atoms minimizing variance. We show that recent methods such as QSGD and TernGrad are special cases of ATOMO and that sparsifiying the singular value decomposition of neural networks gradients, rather than their coordinates, can lead to significantly faster distributed training.

研究动机与目标

• 为降低通信成本，通过在任意原子分解中对随机梯度进行稀疏化来正式化并提出一般的稀疏化方法的动机。
• 证明现有方法如 QSGD 和 TernGrad 是 ATomo 框架的特殊情形。
• 证明对层梯度的奇异值分解进行稀疏化可在相同通信预算下实现更快的训练。
• 在不同原子分解下，对基于稀疏性的方差最小化提供理论保证。
• 在标准数据集和架构上对比评估 spectral-ATomo、QSGD 和 TernGrad，以量化加速和准确性影响。

提出的方法

• 定义一个单位范数原子 a_i 的原子分解 g = sum_i lambda_i a_i。
• 提出无偏原子稀疏化：g_hat = sum_i (lambda_i t_i / p_i) a_i，其中 t_i ~ Bernoulli(p_i) 且固定稀疏预算 sum_i p_i = s。
• 在稀疏性约束下，求解最小化 E||g_hat||^2 的闭式优化，得到依赖于 lambda_i 和所选分解的最优 p_i。
• 证明 1-bit QSGD 和 TernGrad 在 ATomo 的特定 p_i 下是特殊情形，并在这些子域中证明方差最优性。
• 将 ATomo 扩展到矩阵，比较逐元素稀疏化与谱（SVD）稀疏化，提出 Spectral-ATomo。
• 给出在相同通信成本下，谱稀疏化相较逐元素稀疏化能带来更低方差的条件。

实验结果

研究问题

• RQ1 在固定通信预算下，是否存在一个统一的原子分解框架实现无偏梯度稀疏化的最优方差？
• RQ2 在何种原子分解（逐元素、SVD、傅里叶等）可以在通信成本和收敛方差之间获得最佳权衡？
• RQ3 在不牺牲精度的前提下，对层梯度的 SVD 进行稀疏化是否能比逐元素稀疏化方法更快地训练？
• RQ4 何时 QSGD 和 TernGrad 成为 ATomo 的精确实例，以及在何种情形下它们是最优的？
• RQ5 在标准 CNN 因子训练上，Spectral-ATomo 的实际性能收益如何？

主要发现

• ATomo 提供了跨任意原子分解的方差最小化的无偏梯度估计器的闭式解。
• 1-bit QSGD 和 TernGrad 是 ATomo 的特殊情形，在各自的参数子域中是最优的。
• 对层梯度的奇异值分解（Spectral-ATomo）进行稀疏化在相同通信预算下可以实现比逐元素稀疏化更快的训练。
• Spectral-ATomo 与 QSGD 在实证评估中在训练时间上可达到对 QSGD 的约 2x 加速、对 TernGrad 的约 3x 加速（论文中给出上下文）。
• 在 CIFAR-10、CIFAR-100、SVHN 的 ResNet-18 和 VGG-11-BN 上的实验表明，谱化的 ATomo 能在保持精度的同时带来显著的加速，具体取决于稀疏性/量化位宽配置。