Loss and learning rate scaling strategies for Tensorflow distributed training when using TF Estimator

Question

For those who don't want to read the whole story:

TL; DR: When using TF Estimator, do we have to scale learning rate by the factor by which we increase batch size (I know this is the right way, I am not sure if TF handles this internally)? Similarly, do we have to scale per example loss by global batch size (batch_size_per_replica * number of replicas)?

Documentation on Tensorflow distributed learning is confusing. I need clarification on below points.

1. It is now understood that if you increase the batch size by a factor of k then you need to increase the learning rate by k (see this and this paper). However, Tensoflow official page on distributed learning makes no clarifying comment about this. They do mention here that learning rate needs to be adjusted. Do they handle the learning rate scaling by themselves? To make matters more complicated, the behavior is different in Keras and tf.Estimator (see next point). Any suggestions on should I increase the LR by a factor of K or not when I am using tf.Estimator?

2. It is widely accepted that the per example loss should be scaled by global_batch_size = batch_size_per_replica * number of replicas. Tensorflow mentions it here but then when illustrating how to achieve this with a tf.Estimator, they either forget or the scaling by global_batch_size is not required. See here, in the code snippet, loss is defined as follows.

loss = tf.reduce_sum(loss) * (1. / BATCH_SIZE)

and BATCH_SIZE to the best of my understanding is defined above as per replica batch size.

To complicate things further, the scaling is handled automatically if you are using Keras (for reasons I will never understand, it would have been better to keep everything consistent).

Answer 1

1. The learning rate is not automatically scaled by the global step. As you said, they even suggest that you might need to adjust the learning rate, but then again only in some cases, so that's not the default. I suggest that you do increase the learning rate manually.

2. If we take a look at a simple tf.Estimator, the tf.estimator.DNNClassifier (link), the default loss_reduction is losses_utils.ReductionV2.SUM_OVER_BATCH_SIZE. If we got to that Reduction (found here), we see that it's a policy for how to combine losses form indiviual samples together. On a single machine, we would just use tf.reduce_mean, but you can't use that in a distributed setting (as mentioned in the next link). The Reduction leads us to here, which shows you 1) an implementation of how you would implement the global step and 2) explains why. As they are telling you you should implement this yourself, this implies it's not handled by tf.Estimator. Also note that yo can find some explanations on the Reduction page they state the differences bewteen Keras and Estimators about these params.