`jax.nn.standardize` returns `nan` when variance is lower than `-epsilon`

[bondquant]

Description

Sirs,

jax.nn.standardize could be easily improved by taking a lower bound on variance instead of adding epsilon. I found a minimal reproducible example and hacked a quick fix (below). Your comment in the code admits lower accuracy, but at least it should not return nans unnecessarily, don't you think?

return jnp.subtract(x, jnp.asarray(mean)) * lax.rsqrt(jnp.asarray(variance) + epsilon)
replace with
return jnp.subtract(x, jnp.asarray(mean)) * lax.rsqrt(jnp.maximum(jnp.asarray(variance), epsilon))

Example:

import jax
from jax import numpy as jnp, random as jrnd
x = -11. * jnp.ones((3,))
noise = jrnd.normal(jrnd.key(0)) * 2e-6
jax.nn.standardize(x + noise)
# returns Array([nan, nan, nan], dtype=float32)

System info (python version, jaxlib version, accelerator, etc.)

jax:    0.5.2
jaxlib: 0.5.1
numpy:  2.0.2
python: 3.11.13 (main, Jun  4 2025, 08:57:29) [GCC 11.4.0]
device info: cpu-1, 1 local devices"
process_count: 1
platform: uname_result(system='Linux', node='5453835f9caa', release='6.1.123+', version='#1 SMP PREEMPT_DYNAMIC Sun Mar 30 16:01:29 UTC 2025', machine='x86_64')

[achintya-p]

Can I take up this issue?

[nandan2003]

PR created!

[jakevdp]

Hi – I think the real issue here is what the comment mentions: the routine uses a less accurate computation for the variance which may return negative values, when mathematically the variance should always be positive.

I worry about changing out + epsilon to max(out, epsilon) without additional testing, because it amounts to a numerics change in a routine that may be used in numerics-sensitive code downstream, in a way that also injects a discontinuity in the gradient.

Better probably would be to correct negative variance directly, with something like jnp.clip(variance, 0). Still, even clipping can lead to poor autodiff behavior, and that may have downstream impact since this is generally used in fitting of models where well-behaved autodiff is important.

Given the uncertainties here, we probably won't accept a simple PR fixing this, unless it's accompanied by some exploration of what the downstream impacts of the fix may be (e.g. by running integration tests from downstream projects to see the effect).

[nandan2003]

@jakevdp I followed up on your concern regarding autodiff behavior and ran a local gradient check on the negative variance edge case.

1. Clipping (jnp.maximum(variance, 0.)): As you suspected, this resulted in dead gradients (all zeros) for inputs triggering negative variance.
2. Absolute Value (lax.abs(variance)): I tested this approach and confirmed it preserves non-zero gradients, allowing the model to recover from these numerical instability states.

I have updated the implementation to use lax.abs(variance) to ensure downstream training stability. This also seems to align with the internal fix synced in #33513.

[nandan2003]

I have also added a regression test and a gradient stability test to tests/nn_test.py.