google-deepmind / google-deepmind/optax

float32 CTC loss can round below zero for confident valid alignments

Open
#1,771 0 comments 0 reactions 0 assignees View on GitHub
type:bug
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.