Can We Learn Communication-Efficient Optimizers? (2312.02204v1)

Abstract: Communication-efficient variants of SGD, specifically local SGD, have received a great deal of interest in recent years. These approaches compute multiple gradient steps locally, that is on each worker, before averaging model parameters, helping relieve the critical communication bottleneck in distributed deep learning training. Although many variants of these approaches have been proposed, they can sometimes lag behind state-of-the-art adaptive optimizers for deep learning. In this work, we investigate if the recent progress in the emerging area of learned optimizers can potentially close this gap while remaining communication-efficient. Specifically, we meta-learn how to perform global updates given an update from local SGD iterations. Our results demonstrate that learned optimizers can substantially outperform local SGD and its sophisticated variants while maintaining their communication efficiency. Learned optimizers can even generalize to unseen and much larger datasets and architectures, including ImageNet and ViTs, and to unseen modalities such as LLMing. We therefore demonstrate the potential of learned optimizers for improving communication-efficient distributed learning.

• The paper introduces learned optimizers that reduce communication costs by integrating local SGD with adaptive learning strategies.
• It presents two architectures—worker-aware and worker-invariant—that manage updates from multiple nodes for optimized performance.
• Results show that these optimizers generalize well to large datasets and diverse models, improving efficiency in distributed deep learning.

In the field of distributed deep learning, there's a pivotal challenge that often goes unmentioned: the communication bottleneck. As neural networks grow and data explodes, it becomes cumbersome for systems to frequently synchronize model parameters across numerous computational nodes. This synchronization is essential when using the popular Stochastic Gradient Descent (SGD) algorithm, which updates models incrementally as new data is processed.

Local SGD offers one solution. By conducting several gradient steps locally on individual workers before a concerted model update, this variant of SGD reduces how often nodes must communicate. However, despite being less demanding on communication, local SGD sometimes struggles to keep pace with the latest adaptive optimizers, which can navigate the complex optimization landscapes of deep learning more adeptly.

Intriguingly, recent advances suggest that optimizers themselves can be learned. Essentially, instead of pre-establishing how a model should adapt over time, we can train an algorithm to figure out the best adaptation process as it goes. This paper explores the combination of local SGD's communication efficiency with the dynamism of learned optimizers.

The paper introduces two architectures for these novel optimizers: one that's aware of individual workers and one that isn't. The worker-aware kind, quite intuitively, has direct access to updates from each worker node, allowing it to potentially make more informed and complex decisions when aggregating these updates. On the flip side, the worker-invariant kind deals with a single average update from all nodes and is more versatile, as it isn't constrained by the number of workers.

In practice, these learned optimizers displayed not only proficiency in performing well with the seen task and dataset but also an impressive ability to generalize to entirely new datasets and models, including vast datasets like ImageNet and architectures like Vision Transformers and LLMs. This generalizability is particularly significant—it implies that an optimizer learned on one problem might apply to a variety of others, which is a coveted trait in machine learning systems.

Overall, the paper lays a strong foundation for considering learned optimizers in the quest for communication-efficient deep learning. These results suggest that learned optimizers could substantially improve the efficiency of distributed deep learning while maintaining, and potentially even enhancing, model performance. By efficiently navigating the trade-offs between computation, communication, and learning rates, learned optimizers might soon become an indispensable tool in distributed AI systems.