Understanding Outer Optimizers in Local SGD: Learning Rates, Momentum, and Acceleration

2025-09-12

AI Summary

This work investigates the impact of outer-loop optimizers—specifically learning rate, momentum, and acceleration—on the convergence behavior of Local SGD. Addressing the common underestimation of the outer learning rate, we theoretically establish that it can exceed unity to balance optimization error and gradient noise while compensating for suboptimal inner-loop learning rates. We propose, for the first time, an adaptive outer-momentum mechanism and rigorously prove that outer-loop acceleration improves the convergence rate per communication round. Methodologically, we introduce data-dependent analysis to unify the modeling of outer-optimizer effects within the Local SGD framework. Experiments on standard language models demonstrate that proper configuration of outer-loop optimizers significantly accelerates convergence and enhances training stability.

Technology Category

• Machine Learning: Optimization
• Natural Language Processing: Learning & Optimization for NLP
• Search and Optimization: Local Search

Application Category

• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs

Abstract

Modern machine learning often requires training with large batch size, distributed data, and massively parallel compute hardware (like mobile and other edge devices or distributed data centers). Communication becomes a major bottleneck in such settings but methods like Local Stochastic Gradient Descent (Local SGD) show great promise in reducing this additional communication overhead. Local SGD consists of three parts: a local optimization process, an aggregation mechanism, and an outer optimizer that uses the aggregated updates from the nodes to produce a new model. While there exists an extensive literature on understanding the impact of hyperparameters in the local optimization process, the choice of outer optimizer and its hyperparameters is less clear. We study the role of the outer optimizer in Local SGD, and prove new convergence guarantees for the algorithm. In particular, we show that tuning the outer learning rate allows us to (a) trade off between optimization error and stochastic gradient noise variance, and (b) make up for ill-tuning of the inner learning rate. Our theory suggests that the outer learning rate should sometimes be set to values greater than $1$. We extend our results to settings where we use momentum in the outer optimizer, and we show a similar role for the momentum-adjusted outer learning rate. We also study acceleration in the outer optimizer and show that it improves the convergence rate as a function of the number of communication rounds, improving upon the convergence rate of prior algorithms that apply acceleration locally. Finally, we also introduce a novel data-dependent analysis of Local SGD that yields further insights on outer learning rate tuning. We conduct comprehensive experiments with standard language models and various outer optimizers to validate our theory.

Problem

• Analyzing outer optimizer impact on Local SGD convergence and performance
• Investigating outer learning rate tuning for optimization error trade-offs
• Studying momentum and acceleration effects in distributed optimization settings

Innovation

• Tuning outer learning rate for optimization trade-offs
• Using momentum in outer optimizer for convergence
• Applying acceleration to improve communication rounds efficiency