> They don't mention how/whether they tuned the learning rates and batch sizes to optimize for each different device.
All networks were trained with the same hyperparameters. Only the batch size was increased with the amount of parallelization used (so increased 8-fold for TPU and distributed GPU).
> Like they mention, they also use a very small network that isn't something you need the power of a TPU to train quickly and may scale differently than a large network.
I agree completely here. We unfortunately didn't have access to the TPUs long enough to create more useful benchmarks on networks like Resnet-50 and with bigger datasets.
> They also don't post their code so I can't check that their problems with ADAM aren't due to using L2 regularization, which https://arxiv.org/abs/1711.05101 shows leads to worse performance than SGD and you should use weight decay instead.
The code is the same for all devices and in the single-GPU and CPU benchmarks we see Adam performing better than SGD. We did not use any regularization besides a dropout layer so I don't think this explains the bad Adam performance on TPU.
I'll make some effort to clean up the code so it can be shared.