How to replace infs to avoid nan gradients in PyTorch

user_9681081 0
Viewed 7607

I need to compute log(1 + exp(x)) and then use automatic differentiation on it. But for too large x, it outputs inf because of the exponentiation:

>>> x = torch.tensor([0., 1., 100.], requires_grad=True)
>>> x.exp().log1p()
tensor([0.6931, 1.3133,    inf], grad_fn=<Log1PBackward>)

Since log(1 + exp(x)) ≈ x for large x, I thought I could replace the infs with x using torch.where. But when doing this, I still get nan for the gradient of too large values. Do you know why this happens and if there is another way to make it work?

>>> exp = x.exp()
>>> y = x.where(torch.isinf(exp), exp.log1p())  # Replace infs with x
>>> y  # No infs
tensor([  0.6931,   1.3133, 100.0000], grad_fn=<SWhereBackward>)
>>> y.sum().backward()  # Automatic differentiation
>>> x.grad  # Why is there a nan and how can I get rid of it?
tensor([0.5000, 0.7311,    nan])
pytorch python
3 Answers

A workaround I've found is to manually implement a Log1PlusExp function with its backward counterpart. Yet it does not explain the bad behavior of torch.where in the question.

>>> class Log1PlusExp(torch.autograd.Function):
...     """Implementation of x ↦ log(1 + exp(x))."""
...     @staticmethod
...     def forward(ctx, x):
...         exp = x.exp()
...         ctx.save_for_backward(x)
...         return x.where(torch.isinf(exp), exp.log1p())
...     @staticmethod
...     def backward(ctx, grad_output):
...         x, = ctx.saved_tensors
...         return grad_output / (1 + (-x).exp())
... 
>>> log_1_plus_exp = Log1PlusExp.apply
>>> x = torch.tensor([0., 1., 100.], requires_grad=True)
>>> log_1_plus_exp(x)  # No infs
tensor([  0.6931,   1.3133, 100.0000], grad_fn=<Log1PlusExpBackward>)
>>> log_1_plus_exp(x).sum().backward()
>>> x.grad  # And no nans!
tensor([0.5000, 0.7311, 1.0000])
user_96810810

But for too large x, it outputs inf because of the exponentiation

This is why x should never be too large. It should be ideally in range [-1, 1]. If this is not the case you should normalize your inputs.

user_58849550

If x >= 20, the function output is approximately x. Use PyTorch method torch.softplus. It helps the issue.

user_123064930
Related