How to get the second moment of the gradient

Viewed 1157

In the OpenAI Five paper it is mentioned that the "Gradients are additionally clipped per parameter to be within between ±5√v where v is the running estimate of the second moment of the (unclipped) gradient.". This is something I would like to implement in my project, but I am not sure how to do it neihter in theory nor in practice.

From wikipedia I found out that the "The second central moment is the variance. The positive square root of the variance is the standard deviation [...]". My best guess regarding the "running estimate" is that it is the Exponential Moving Average. The gradients of a network can be accessed as this comment suggests.

From these I would assume that √v is the Exponential Running Average of the standard dev. of the gradients and could be calculated via: estimate = alpha * torch.std(list(param.grad for param in model.parameters())) + (1-alpha) * estimate

Is my theory correct? Is there a better way to do it? Thanks in advance.

Edit: fixed gradient gathering after Mr. For Example"s answer.

1 Answers

I think you are on the right path, my guess is basically as same as yours, just with little different.

First, what is moment?

N-th moment of a random variable is defined as the expected value of that variable to the power of n. More formally:

enter image description here

m — moment, X — random variable

So the first moment is mean, and the second moment is uncentered variance (meaning we don’t subtract the mean during variance calculation), intuitively, clipping the gradients by moving average of its standard deviation wrt zero make sense.

Second, what is the correct code?

list(network.parameters()) only give you the parameters, for get gradient of each parameters you need [param.grad for param in network.parameters()]

Given all the things we know above, the correct code should be (you can try to optimize it by all means):

grads_square = torch.FloatTensor([torch.square(param.grad) for param in network.parameters()])
estimate = alpha * torch.sqrt(torch.mean(grads_square)) + (1-alpha) * estimate
Related