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.