The goal of the project was to study algorithmic and foundational aspects of training machine learning systems on distributed platforms, in particular using distributed stochastic optimization.  Stochastic optimization refers to optimizing an objective, searching for a good data fit, or training a system, by considering only a few examples at each iteration, as opposed to repeatedly considering the entire data set.  It traces its roots to a classic 1951 paper by Robbins and Monro, has been a staple of machine learning ever since the first learning systems in the 1950s and 1960s, and to this day stochastic optimization techniques, such as Stochastic Gradient Descent (SGD) form the basis of most machine learning training procedures.  A firm theoretical foundation for convex stochastic approximation was provided by the seminal work of Nemirovskii and Yudin in the 1970s, and exact optimal complexities and methods were derived by Lan in 2012 based on Nemirovskii’s work and Nesterov’s 1984 follow-up.  But modern machine learning frequently departs from the sequential settings and relies on parallel and distributed computation in order to handle massive data sets.   This includes parallelization at many scales, ranging from parallelization across cores in a GPU or multiple processors in a single computer, through parallelization across servers in a data center, to distributed computation on edge devices (e.g. individual phones, cars, or sensors) as in Federated Learning.  In this project we provide a firm theoretical framework for stochastic distributed optimization, develop methods with provable guarantees, and establish what is the best that can be done, and how.

The major achievement of the project, recognized by the Best Paper Award at the Conference on Learning Theory (the leading venue for such research) was establishing the “optimal complexity” of stochastic convex distributed optimization.  That is, what is the best that could be achieved, in terms of the required computation and communication, and subject to established assumptions, and what methods achieve this performance.  Previous work done as part of the project built towards this by providing the theoretical grounding, describing the relevant models, and studying the behavior of different methods, including showing that some commonly used methods can actually be problematic in some important regimes.  Understanding what is the best that can be done using the standard assumptions and standard model, also allows us to go further, by allowing us to understand what additional assumptions or stronger types of operations we must use in order to break this barrier and obtain better methods, as we have done in work employing “Hessian-vector-prodcuts”.

Improvements in distributed stochastic optimization directly translate into faster training with less communication in many scientific and engineering applications of machine learning.  E.g. through our interactions with Google, these advancements are quickly finding their way to reducing communication in Federated Learning.  Beyond the reduction in resource consumption, Federated Learning also aims at improving user privacy by maintaining user data on their own devices, which now become nodes in a distributed training/optimization process.  Reductions in communication may thus also lead to improved privacy.

The project resulted in 16 refereed conference and journal publications, including two awards (Best Student Paper and Best Paper) at the Conference on Learning Theory, which is the premier venue for research on the theory of machine learning.  This project was an international collaboration between PI Nathan Srebro at TTI-Chicago and PI Ohad Shamir at the Weizmann Institute of Science in Israel.  The project served as a basis for the PhD thesis of two graduate students, Blake Woodworth at TTI-Chicago and Yossi Arjevani at the Weizmann institute.  In addition, Jialei Wang at the University of Chicago finished his PhD working on initial stages of the project, and Kshitij Patel at TTIC worked on the project during his first two years at the PhD program, and his PhD will be based on extensions to the project.  All four students received mentorship from both PIs, and the project provided opportunities for exposure to different research ideas and environments, at Chicago and at the Weizmann.  Woodworth is now a post-doctoral researcher with Francis Bach at INRIA Paris, further expanding his international exposure, and Arjevani recently started a faculty position at the Hebrew University, after post-doctoral training at NYU.

Last Modified: 06/01/2022
Modified by: Nathan Srebro