Thread by @zacharynado

A thread on our latest optimizers work! We tune Nesterov/Adam to match performance of LARS/LAMB on their more commonly used workloads. We (@jmgilmer, Chris Shallue, @_arohan_, @GeorgeEDahl) do this to provide more competitive baselines for large-batch training speed measurements

We are **not** trying to prove that any optimizer is better than any other (more on that later). However, we believe that well tuned baselines are very important, especially in optimization where there are so many confounding factors.

In general, we found that LR schedules were the most impactful part of the pipeline to tune (that is, once we got all the model and dataset details to be the same!).

LARS has been tuned for years in MLPerf, so tuning the LR schedules for other algorithms is critical. We found different LR schedule shapes work better for different optimizers, meaning just using the same schedule in the baseline could give suboptimal/misleading results

Similarly, after fixing bugs in existing Adam baselines for BERT-Large and re-tuning hyperparameters including schedule, we set new stronger baselines for batch sizes up to 65k

We tuned *a lot*, and report excruciating amounts of detail in appendices. We recommend others release these intermediate results (regardless of how silly) as it is hard to a priori determine which choices one should consider when tuning, and these results show our decision path

We’ve also released #JAX ResNet-50 code at github.com/google/init2wi…, to more easily reproduce our results. We believe this is important because without the MLPerf reference code we would never have realized how impactful several implementation differences are.

One consideration of optimizer studies we did not consider was “ease of tuning”. We believe this is still an open question, which will require extremely careful consideration and study, that we leave for future work 😀

Comparing optimizers is hard, but in collaboration with researchers in academia and industry, we’re developing a competition to measure optimization training speeds as a new MLCommons working group for algorithmic efficiency mlcommons.org/en/groups/rese…

A link to the arxiv preprint! arxiv.org/abs/2102.06356