arXiv:2501.05633cs.LGcs.IT2025-01

提出新方法让梯度压缩更稳定,显著提升高压缩率下的收敛效果。

Regularized Top-$k$: A Bayesian Framework for Gradient Sparsification

  • 将梯度压缩建模为贝叶斯推断问题,导出正则化Top-k算法
  • 在分布式线性回归中,高压缩比下仍能收敛到全局最优
  • 适用于模型压缩与分布式训练,适合关注通信效率的研究者

误差累积在分布式场景中有效实现梯度稀疏化:初始未选中的梯度分量会因累积误差超过阈值而被后续选中。该机制本质上对选中分量的更新学习率进行了缩放。尽管此特性可避免分布式梯度下降中横向移动的退化,但在某些场景下会损害收敛性。本文提出一种新型稀疏化方案,用于控制误差累积带来的学习率缩放。首先将梯度稀疏化建模为逆概率(推理)问题,基于拓扑-最大后验估计推导出贝叶斯最优稀疏掩码。利用从Top-k继承的先验分布,导出一种新的稀疏化算法,可解释为正则化版的Top-k。该算法(称作RegTop-k)通过历史聚合梯度评估下一阶段的后验统计,并据此优先选择本地累积梯度分量。数值实验验证了其有效性:在分布式线性回归中,当Top-k始终与全局最优保持固定距离时,RegTop-k可在显著更高的压缩比下收敛至全局最优。进一步在CIFAR-10上用ResNet-18进行分布式训练,以及在ImageNette上微调多个视觉模型,结果表明随着压缩比升高,RegTop-k明显优于Top-k。

原文摘要 · Abstract (English)

Error accumulation is effective for gradient sparsification in distributed settings: initially-unselected gradient entries are eventually selected as their accumulated error exceeds a certain level. The accumulation essentially behaves as a scaling of the learning rate for the selected entries. Although this property prevents the slow-down of lateral movements in distributed gradient descent, it can deteriorate convergence in some settings. This work proposes a novel sparsification scheme that controls the learning rate scaling of error accumulation. The development of this scheme follows two major steps: first, gradient sparsification is formulated as an inverse probability (inference) problem, and the Bayesian optimal sparsification mask is derived as a maximum-a-posteriori estimator. Using the prior distribution inherited from Top-k, we derive a new sparsification algorithm which can be interpreted as a regularized form of Top-k. We call this algorithm regularized Top-k (RegTop-k). It utilizes past aggregated gradients to evaluate posterior statistics of the next aggregation. It then prioritizes the local accumulated gradient entries based on these posterior statistics. We validate our derivation through various numerical experiments. In distributed linear regression, it is observed that while Top-k remains at a fixed distance from the global optimum, RegTop-k converges to the global optimum at significantly higher compression ratios. We further demonstrate the generalization of this observation by employing RegTop-k in distributed training of ResNet-18 on CIFAR-10, as well as fine-tuning of multiple computer vision models on the ImageNette dataset. Our numerical results confirm that as the compression ratio increases, RegTop-k sparsification noticeably outperforms Top-k.

梯度压缩分布式训练贝叶斯方法优化算法

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。