# Why does my activation function cause NaNs?

I implemented the following activation function: $$\sigma(x) = x^{1/3},$$ which returns NaNs after some epochs. I think this is due to the derivative exploding close to $$0$$. To fix this issue, I implemented a modified activation function: $$\sigma'(x) = \begin{cases} x, & \text{if } -1 \leq x \leq 1, \\ x^{1/3}, & \text{otherwise.} \end{cases}$$ Unfortunately, $$\sigma'$$ suffers from the same issue. Why is that?

Here is my code:

def odd_pow(input, exponent):
return input.sign() * input.abs().pow(exponent)

def custom_activation(x):
# This will create a mask where values are between -1 and 1
mask = (x >= -1) & (x <= 1)
# For values between -1 and 1, apply x
# For other values, apply x^1/3
# Use where to combine results based on the mask


Note: Both ReLU and Tanh don't return NaNs, but I am looking for homogeneous activation functions $$\sigma(ax) = g(a)\sigma(x)$$.

• Does the derivative cause the NaNs, or the normal activation function? Commented Sep 11, 2023 at 13:13
• @RobinvanHoorn I think the derivative is what is causing the NaNs. Commented Sep 11, 2023 at 13:13
• You need to use a small constant to prevent exploding gradients. Please, see my answer. Commented Sep 13, 2023 at 8:21
• Please look at the gradients being backpropagated to see where the issue comes from. Commented Sep 13, 2023 at 14:10

Update: the issue is that your odd_pow() is numerically unstable, and so to fix the exploding gradients, which cause a NaN loss, you simply need to add a small constant when computing the power:

def odd_pow(input, exponent, eps=1e-4):
return input.sign() * (input.abs() + eps).pow(exponent)


this ensures that the pow is not computed at zero, avoiding NaNs.

Note: for reference I also keep my old answer, since also the activation $$\sigma(x)$$ suffers the same problem.

The issue is that the function $$\sigma(x) = x^{1/3}$$ is zero for $$x\le 0$$, and so its derivative becomes NaN for negative $$x$$ and infinity for $$x=0$$ as shown here: where the blue curve is the derivative of $$\sigma(x)$$, which is depicted in green.

To avoid either NaN or exploding gradients, you need to add a small value (usually called an epsilon) to make the computation of the derivative more numerically stable. The following should work:

def custom_activation(x, eps=1e-3):
# set negative values to zero
tmp = torch.nn.functional.relu(x)
# add the small epsilon eps

Notice that as you lower the eps, e.g. at 1e-5, the derivative would be larger. If you really need a small epsilon, then you should also consider to clip the individual gradients (there should be a clip_grad functionality related to the optimizer) to ensure stable training.
• The function $\sigma(x) = x^{1/3}$ is not zero for $x \leq 0$—I am taking the real root. Commented Sep 11, 2023 at 13:54
• The zero at $x=0$ is not the issue. It's the derivative is $\frac{1}{3}x^{-\frac{2}{3}}$ which is undefined at zero. With their intended $\sigma'$ function, which should simply have derivative of $1$ at $x=0$, the OP looks to have an implementation problem here, not a theoretical one. Commented Sep 11, 2023 at 14:02
• @NeilSlater It could be that torch isn't taking the real roots for the negative values of $x^{-2/3}$. Commented Sep 11, 2023 at 16:17
• @user572780: I would have thought that the .abs() would have stopped that issue. I cannot see clearly what is wrong, although I suspect one or more sub-components of your custom_activation function is confusing the auto-differentiation. Commented Sep 11, 2023 at 18:02