2 karma · joined January 21, 2019
Still, since Adam performs worse even for non-distributed TPU (where the batch size and learning rate are the same for all devices) this wouldn't explain what we're seeing. Anyway, I'll try to find some time tomorrow to post the code so you can all have a look at it and hopefully see what's up.
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.