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
    return torch.where(mask, odd_pow(x, 1), odd_pow(x, 1/3))

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

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

1 Answer 1


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: enter image description 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`
    return torch.pow(tmp + eps, exponent=1 / 3.0)

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.

  • $\begingroup$ The function $\sigma(x) = x^{1/3}$ is not zero for $x \leq 0$—I am taking the real root. $\endgroup$
    – user572780
    Sep 11, 2023 at 13:54
  • 1
    $\begingroup$ 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. $\endgroup$ Sep 11, 2023 at 14:02
  • $\begingroup$ @NeilSlater It could be that torch isn't taking the real roots for the negative values of $x^{-2/3}$. $\endgroup$
    – user572780
    Sep 11, 2023 at 16:17
  • $\begingroup$ @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. $\endgroup$ Sep 11, 2023 at 18:02
  • $\begingroup$ @NeilSlater I implemented and tested the code I provide here in Tensorflow, on a simple MNIST classification problem and it trains without any issue $\endgroup$ Sep 12, 2023 at 13:00

You must log in to answer this question.

Not the answer you're looking for? Browse other questions tagged .