Why does pow return nan during the backward pass?

Question

I am trying to evaluate the following expression as a loss function:

This function describes a family of closed surfaces. I am using a neural network to try to predict the surface parameters e1, e2, ax, ay, az that minimize the least squares distance from a given set of points. The quantity F(p) - 1, though not proportional, gives a metric of how far a point lies with respect to the surface. This is my code:

def eval_inout_loss(self, p, x):
    #x: BxNx3 tensor -> x, y, z
    #p: Bx5 tensor -> e1, e2, ax, ay, az

    ###########################################################################################
    # INSIDE - OUTSIDE FUNCTION LOSS
    ###########################################################################################
    # Evaluating F(x,y,z) and calculating Σ (F(x,y,z) - 1)^2
    
    #x^(2/e2), y^(2/e2), z^(2/e1)
    xs = torch.pow(torch.abs(x[:,:,0]), 2/p[:,1].unsqueeze(-1))
    ys = torch.pow(torch.abs(x[:,:,1]), 2/p[:,1].unsqueeze(-1))
    zs = torch.pow(torch.abs(x[:,:,2]), 2/p[:,0].unsqueeze(-1))

    s = torch.pow(xs + ys, p[:,1].unsqueeze(-1))
    s = torch.pow(s, 1/p[:,0].unsqueeze(-1)) + zs
    s = torch.pow(s, p[:,0].unsqueeze(-1))
    s = torch.abs(s - 1)

    #multiplying by axayaz to prioritize smaller volumes    
    volparams = p[:,2:].prod(-1)
    s = volparams * s.sum(-1)
    l1 = s.mean()

I understand that low values of e1, e2 would cause this expression to be numerically unstable, so I have restrained the network output to the range [0.2, 1.9] by applying a sigmoid function at the output of the network:

x = 1.7 * torch.nn.functional.sigmoid(self.mlp6(x)) + 0.2

An obvious suspect would be trying to raise a negative number to a fractional exponent, thereby producing nan as the output. However this cannot be the case here, since all numbers are squared and therefor strictly non-negative. To further enhance the numerical stability I read in a relevant paper that it's useful to raise F(p)^e1, and I achieve this with the line s = torch.pow(s, p[:,0].unsqueeze(-1)) shown above.

Setting autograd anomaly detection tool to true it gives me the following error:

  File "/home/..../Desktop/code/file.py", line 391, in eval_inout_loss
  s = torch.pow(xs + ys, p[:,1].unsqueeze(-1))
  (function print_stack)                                                                                                                                                          
 
Traceback (most recent call last):
  File "/home/...../Desktop/code/samplenet.py", line 477, in 
  loss.backward()
  File "/home/..../.local/lib/python3.8/site-packages/torch/tensor.py", line 185, in backward
  torch.autograd.backward(self, gradient, retain_graph, create_graph)
  File "/home/vlassis/.local/lib/python3.8/site-packages/torch/autograd/__init__.py", line 125, in backward
  Variable._execution_engine.run_backward(
  RuntimeError: Function 'PowBackward1' returned nan values in its 0th output.

This can happen during the first iteration or several hundred iterations later, but it always happens. The output of the function doesn't seem to be particularly abnormal when this happens. For example, a possible sequence goes something like this: l1 = 0.2560 -> l1 = 0.2458 -> l1 = nan. I have tried disabling the anomaly detection tool to print the values of the tensors before nan values appear:

e1 min max mean tensor(nan, device='cuda:0', grad_fn=) tensor(nan, device='cuda:0', grad_fn=) tensor(nan, device='cuda:0', grad_fn=)
e2 min max mean tensor(nan, device='cuda:0', grad_fn=) tensor(nan, device='cuda:0', grad_fn=) tensor(nan, device='cuda:0', grad_fn=)
x min max mean tensor(-1.8769, device='cuda:0') tensor(1.6930, device='cuda:0') tensor(0.0006, device='cuda:0')
y min max mean tensor(-1.4985, device='cuda:0') tensor(1.5802, device='cuda:0') tensor(0.0108, device='cuda:0')
z min max mean tensor(-1.6468, device='cuda:0') tensor(1.8991, device='cuda:0') tensor(-0.0091, device='cuda:0')
xs min max tensor(nan, device='cuda:0', grad_fn=) tensor(nan, device='cuda:0', grad_fn=)
ys min max tensor(nan, device='cuda:0', grad_fn=) tensor(nan, device='cuda:0', grad_fn=)
zs min max tensor(nan, device='cuda:0', grad_fn=) tensor(nan, device='cuda:0', grad_fn=)

What are other possible causes for nan to appear?

Why does pow return nan during the backward pass?

0 Answers