Adaptive Consensus Gradients Aggregation for Scaled Distributed Training
This work addresses communication bottlenecks in distributed training for large-scale machine learning, offering an incremental improvement over existing methods.
The paper tackles the problem of suboptimal gradient aggregation in distributed deep learning by formulating it as an objective-aware subspace optimization, resulting in improved performance over standard averaging on MLPerf tasks with maintained efficiency.
Distributed machine learning has recently become a critical paradigm for training large models on vast datasets. We examine the stochastic optimization problem for deep learning within synchronous parallel computing environments under communication constraints. While averaging distributed gradients is the most widely used method for gradient estimation, whether this is the optimal strategy remains an open question. In this work, we analyze the distributed gradient aggregation process through the lens of subspace optimization. By formulating the aggregation problem as an objective-aware subspace optimization problem, we derive an efficient weighting scheme for gradients, guided by subspace coefficients. We further introduce subspace momentum to accelerate convergence while maintaining statistical unbiasedness in the aggregation. Our method demonstrates improved performance over the ubiquitous gradient averaging on multiple MLPerf tasks while remaining extremely efficient in both communicational and computational complexity.