Update EMA decay rate to follow TF: https://github.com/tensorflow/tensorflow/blob/v2.9.1/tensorflow/python/training/moving_averages.py#L469
Update EMA decay rate to follow TF:
https://github.com/tensorflow/tensorflow/blob/v2.9.1/tensorflow/python/training/moving_averages.py#L469

This should speed up convergence in the beginning of training since it gives lower weights to the initial weights.

PiperOrigin-RevId: 458130203

google/paxml,google/paxml