提出新型梯度追踪方法,显著提升分布式学习通信效率与稳定性。
Beyond Scaffold: A Unified Spatio-Temporal Gradient Tracking Method
- 通过时空梯度追踪,同时抑制局部噪声并缓解数据异构影响。
- 强凸问题下实现线性收敛,且通信复杂度随本地更新次数线性加速。
- 适合大规模分布式训练场景,尤其适用于数据分布不均的联邦学习。
在分布式与联邦学习中,为减少通信开销常采用多轮本地更新。然而,节点间数据异构及局部梯度噪声会导致模型偏离全局最优。本文从梯度追踪视角重新审视Scaffold方法,提出统一的时空梯度追踪算法ST-GT,适用于时变图上的分布式随机优化。ST-GT通过跟踪邻接节点的全局梯度以缓解数据异构,同时维持局部梯度的运行平均以大幅抑制噪声,仅需稍高存储开销。无需假设有界数据异构,理论证明其在强凸问题下达到线性收敛,在非凸情况下为次线性收敛。值得注意的是,ST-GT首次在强凸设置下实现通信复杂度关于每轮本地更新数τ的线性加速。相比传统梯度追踪方法,其拓扑相关噪声项由σ²降至σ²/τ,显著提升通信效率。
原文摘要 · Abstract (English)
In distributed and federated learning algorithms, communication overhead is often reduced by performing multiple local updates between communication rounds. However, due to data heterogeneity across nodes and the local gradient noise within each node, this strategy can lead to the drift of local models away from the global optimum. To address this issue, we revisit the well-known federated learning method Scaffold (Karimireddy et al., 2020) under a gradient tracking perspective, and propose a unified spatio-temporal gradient tracking algorithm, termed ST-GT, for distributed stochastic optimization over time-varying graphs. ST-GT tracks the global gradient across neighboring nodes to mitigate data heterogeneity, while maintaining a running average of local gradients to substantially suppress noise, with slightly more storage overhead. Without assuming bounded data heterogeneity, we prove that ST-GT attains a linear convergence rate for strongly convex problems and a sublinear rate for nonconvex cases. Notably, ST-GT achieves the first linear speed-up in communication complexity with respect to the number of local updates per round $τ$ for the strongly-convex setting. Compared to traditional gradient tracking methods, ST-GT reduces the topology-dependent noise term from $σ^2$ to $σ^2/τ$, where $σ^2$ denotes the noise level, thereby improving communication efficiency.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。