google-deepmind / google-deepmind/optax
float32 CTC loss can round below zero for confident valid alignments
- Dominant language
- Python
- Stars
- 2.3k
- Forks
- 369
- Avg merge
- 10h 15m
- Merged PRs (30d)
- 7
Description
I had training runs regularly produce nans and then deteriorate late into training.
After a lot more time than I'd like to admit I've found the culprit. I observed mean CTC losses near `-1.9e-5` and `-3.4e-6`.
In my multi task training setup a downstream adaptive loss balancer computed `log(loss + 1e-8)` and then generated NaNs from otherwise finite training data and gradients.
I didn't check whether this is specific to CTC loss or whether there are other loss functions that can return negative results.
## Reproduction
Observed with Optax 0.2.8 and JAX/jaxlib 0.11.1. No training or mixed precision is required. Run on CPU with x64 disabled:
```python
import jax
import jax.numpy as jnp
import optax
jax.config.update("jax_enable_x64", False)
loss = optax.ctc_loss(
logits=jnp.array([[[0.0, 17.0], [0.0, 17.0]]], dtype=jnp.float32),
logit_paddings=jnp.zeros((1, 2)),
labels=jnp.array([[1]], dtype=jnp.int32),
label_paddings=jnp.zeros((1, 1)),
blank_id=0,
)
print(loss)
```
Observed loss: approximately `[-8.279875e-8]`.
There are three valid paths: `(1, 1)`, `(blank, 1)`, `(1, blank)`.
Writing `p` for the blank probability at each frame, their combined probability is `1 - p**2`;
the true loss `-log1p(-p**2)` is approximately `1.713908e-15`.
The loss should therefore be nonnegative.
With x64 enabled and float64 logits, the same example yields approximately `1.857496e-15`.
## Workaround
I now clamp each sequence's CTC loss to zero before averaging and guard the downstream balancer separately just in case.
## Expected Behavior
I think Optax should either enforce CTC's nonnegative mathematical range or document that callers must handle small negative results.
Contributor guide
Research direction
Start at the optax.ctc_loss entry point and run the provided CPU float32 reproduction with x64 disabled. Trace the CTC loss computation and its numerical behavior for confident valid alignments; done means the reproduced loss no longer falls below zero, with regression coverage for this case.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 70/100