Fix gradient NaN issues in sigmoid_focal_loss for extreme logits

[leochlon]

addresses numerical instability in gradient and Hessian computation when using sigmoid_focal_loss with large logit values and gamma < 2. might fix #1267

the issue occurs because (1 - p_t) approaches zero for extreme logits, causing instability in derivatives of (1 - p_t)^gamma when gamma < 2.

clamping (1 - p_t) to be at least machine epsilon, preventing division by zero and infinite derivatives while maintaining mathematical equivalence for normal cases is a fix

this enables stable use of focal loss with Newton's method, L-BFGS, and other second-order optimization algorithms across the full range of valid gamma values.

test plan:

logits = jnp.array([100.0])
labels = jnp.array([1.0])

Jacobian test (γ=0.5)
grad = jax.grad(lambda x: sigmoid_focal_loss(x, labels, gamma=0.5)[0])(logits)[0]
print(f"Jacobian finite: {jnp.isfinite(grad)}")

Hessian test (γ=1.5)
hess = jax.hessian(lambda x: sigmoid_focal_loss(x, labels, gamma=1.5)[0])(logits)[0, 0]
print(f"Hessian finite: {jnp.isfinite(hess)}")

python -m pytest optax/losses/_classification_test.py::SigmoidFocalLossTest -v

passes

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[Vishal-sys-code]

Excellent fix, unconditionally clamping (1 – pₜ) to ε ensures stable gradients/Hessians for all γ > 0. Only tweak: add a brief docstring note that “(1–pₜ) is clamped to ε to guarantee finite derivatives.” Otherwise, LGTM!

[leochlon]

@Vishal-sys-code I appreciate you taking the time to review the PR so quickly! Thank you for the great feedback, I've added a brief note to the docstring explaining the clamping for finite derivatives as you suggested. Tests are now passing ready to merge :)

[rdyro]

It seems you still have some files in the root directory in the history. You can also run the linter locally using either ruff check . or pre-commit run -a

[leochlon]

Thanks @rdyro, I'll note this for next time! :)

[leochlon]

@vroulet , please let me know what you think of the implementation and if it matches what you were thinking.

[vroulet]

Ok @leochlon you're vibe coding?
How do you expect us maintainers (mostly volunteers) to interact with you then?
(It's really an open question, not a rhetoric one.)

[vroulet]

To be clear, I like this contribution, and you understood pretty well the mathematics behind this I think.
But a library is supposed to provide clear, concise, clean code. You provided 500 lines of code.
Do you expect us to review all 500 lines? How do you picture this to work?

Again this is an open question, and your answer will truly be appreciated. You won't be the last one to vibe code and you may be the first one to actually care and produce something interesting for the library.

[leochlon]

@vroulet

Thanks for the feedback - let me clarify my process and intentions:

On the 500 lines: A significant portion is the automated stability test, which I included as a reproducible proof of the numerical properties but didn't intend for the final merge. I wanted to demonstrate that the implementation handles the edge cases correctly, but I realise dumping it all in one PR makes review difficult.

I spent about 2 days working through the mathematics and implementation. I do use Copilot for debugging and code cleanup, but the core logic and mathematical understanding is mine. I can explain any part of the algorithm and respond to technical feedback.

[leochlon]

@vroulet

What I'd actually propose for the library:

The core log-space implementation (~20-30 lines)

Basic correctness tests + a simple regression test for numerical stability

If you'd like, I can split this into a focused PR with just the core changes + basic tests, with the stability analysis as supporting documentation.

Would you prefer the core implementation as a clean, minimal PR with the analysis referenced separately?

[vroulet]

@leochlon thank you for the answer! Again I appreciate this contribution, it clears up a problem that has been lingering for some time in a very clean way.

Yes, ideally just the core log-space implementation and associated tests would be the best.

About supporting documentation for future PRs. The best practice would be to add it comments to the PR as a reproducible code example. That way maintainers can easily run the example without having to carefully read each line. Careful reviews are really important for good maintenance and so all code that is in your PR should generally be carefully read by someone (like me for example) so it's best to keep it short for the sake of the maintainers.

Thank you again. Sorry if I was very direct but some previous users did not care as much as you and simply returned llm generated code without even communicating (you did and I sincerely appreciate it).

[leochlon]

@vroulet I really appreciate your dedication keeping this repo high quality, there's so much cool stuff to build with optax and I've been working on a bunch of other improvements, so hopefully we can work together more going forward.

Makes total sense about keeping PRs focused and reviewable. I'll restructure this as:

Clean core implementation + tests in the PR
Put the extensive validation stuff in comments as reproducible examples

I'm pretty obsessive about reproducibility and including comprehensive tests to back up why I change things, but now I know the right way to present them without creating review burden (again sorry about that haha).

Thanks for being welcoming and explaining the process! I'll clean this up to just the core log-space fix and move the stability analysis to comments. Looking forward to contributing more here.

[vroulet]

That looks really great. I also appreciate a lot the additional documentation.
I'm leaving you a few comments, once resolved we should be able to merge.
Looking forward for more contributions :)

[vroulet]

Do you still need the clamping or is it really for unreasonable cases (like logits up 10**6)?

[leochlon]

So when we compute exp(γ * log(1 - p_t)), the gradient is:

∂/∂logits [exp(γ * log(1 - p_t))] = γ * (1 - p_t)^(γ-1) * ∂(1-p_t)/∂logits

When γ < 1: γ-1 < 0, so (1 - p_t)^(γ-1) = (1 - p_t)^{-(1-γ)} → ∞ as p_t → 1
When γ ≥ 1: γ-1 ≥ 0, so (1 - p_t)^(γ-1) → 0 as p_t → 1 (bounded)

Log-space computation solves numerical representation problems:

log(1 - p_t) is representable even when 1 - p_t underflows
γ * log(1 - p_t) is representable even when (1 - p_t)^γ underflows

but this doesn't fix gradient explosion. Even if you can represent the values, the gradients can still explode when γ < 1, hessian explodes when γ < 2.

Clamping prevents gradient explosion by keeping p_t away from 1.

[vroulet]

It is pretty stable to me without the clamping (see code below).
Moving to logspace is sufficient for numerical stability. Of course from a mathematical viewpoint if log(1-pt) tends to -infty we would have issues. But numerically log(1-pt) (for e.g. label=0) is equivalent to -logits so unless we go to logits equal to huge numbers we are good.

import jax
import jax.numpy as jnp
import chex
import optax

def sigmoid_focal_loss(
    logits: chex.Array,
    labels: chex.Array,
    alpha = None,
    gamma: float = 2.0,
) -> chex.Array:
  chex.assert_type([logits], float)
  labels = jnp.astype(labels, logits.dtype)

  # Cross-entropy loss
  ce_loss = optax.sigmoid_binary_cross_entropy(logits, labels)

  # Compute log(1-p_t) using logsumexp unconditionally
  log_p = jax.nn.log_sigmoid(logits)
  log_q = jax.nn.log_sigmoid(-logits)

  log_one_minus_p_t = jax.scipy.special.logsumexp(
      jnp.stack([log_p, log_q], axis=-1),
      axis=-1,
      b=jnp.stack([1 - labels, labels], axis=-1)
  )

  # Focal weight and final loss
  focal_weight = jnp.exp(gamma * log_one_minus_p_t)
  loss = ce_loss * focal_weight

  # Alpha weighting
  if alpha is None:
      return loss
  else:
      weighted = (alpha * labels + (1 - alpha) * (1 - labels)) * loss
      return weighted

labels = jnp.asarray(0)
logits = jnp.asarray(-10.**4)
grad_stable = jax.grad(sigmoid_focal_loss)(logits, labels, gamma=0.5)
grad_unstable = jax.grad(optax.sigmoid_focal_loss)(logits, labels, gamma=0.5)
print(f'{grad_unstable=}, {grad_stable=} at {logits=}')
hess_stable = jax.grad(jax.grad(sigmoid_focal_loss))(logits, labels, gamma=0.5)
hess_unstable = jax.grad(jax.grad(optax.sigmoid_focal_loss))(logits, labels, gamma=0.5)
print(f'{hess_unstable=}, {hess_stable=} at {logits=}')

[leochlon]

You're right!!

The explosion in (1 - p_t)^(γ-1) is exactly canceled by the vanishing of ∂(1-p_t)/∂logits, resulting in a product that goes to 0!

For case $y=0$, $\text{logits} \to -\infty$ (the problematic regime):

The focal weight is $w = (\sigma(\text{logits}))^{\gamma}$ and its gradient is $$\frac{\partial w}{\partial \text{logits}} = \gamma \cdot (\sigma(\text{logits}))^{\gamma-1} \cdot \sigma(\text{logits})(1-\sigma(\text{logits}))$$
$$= \gamma \cdot (\sigma(\text{logits}))^{\gamma} \cdot (1-\sigma(\text{logits}))$$

While $(\sigma(\text{logits}))^{\gamma-1} \to \infty$ when $\gamma &lt; 1$, the full gradient contains the product $(\sigma(\text{logits}))^{\gamma-1} \cdot \sigma(\text{logits}) = (\sigma(\text{logits}))^{\gamma}$.

when we take the limit $\text{logits} \to -\infty$: $\sigma(\text{logits}) \sim e^{\text{logits}} \to 0$ and $(\sigma(\text{logits}))^{\gamma} \sim e^{\gamma \cdot \text{logits}} \to 0$ exponentially fast.

so $\lim_{\text{logits} \to -\infty} \frac{\partial w}{\partial \text{logits}} = 0$ for any $\gamma &gt; 0$.

The explosion in $(\sigma(\text{logits}))^{\gamma-1}$ is exactly canceled by the vanishing of $\sigma(\text{logits})$, and so clamping is actually unnecessary in the logarithmic regime.

new commit removes the clamping =]

[vroulet]

I would have expected the results to be more accurate. Is it because what is it tested against does not catch up the accuracy you gained?

[leochlon]

The failing test expected the focal loss to match a very specific mathematical formula for extreme logits with gamma=2. When we switched to the log-space implementation implementation, we went from from using (1 - p_t) ** gamma directly to exp(gamma * log(1 - p_t)) , and our epsilon clamping prevents underflow, slightly changing results for extreme values
logsumexp precision. Finally, JAX's logsumexp has different floating-point behaviour than manual computation. For extreme logits, these have slightly different numerical behaviour log-space version is more numerically stable

[vroulet]

Great. Some additional comments on the tests.
Also check that it matches the sigmoid_binary_cross_entropy() for gamma =0

[vroulet]

Go bigger, like 100.
Have some non integer labels too.

[leochlon]

Done

[vroulet]

Testing with a single gamma<1 (like gamma=0.5) is sufficient

[leochlon]

Done

[vroulet]

Could test with hessian too.

[leochlon]

There was a Hessian test in the same function, did you mean a different kind of test?

Test Hessians for numerical stability

hessian = jax.hessian(loss_fn)(extreme_logits)
self.assertTrue(jnp.all(jnp.isfinite(hessian)),
f"Hessians should be finite for gamma={gamma}")

[vroulet]

You can remove the test with alpha. It won't pose problems

[leochlon]

done

[rdyro]

Could you squash the commits please, when you're ready with the PR?

[leochlon]

@rdyro done.