Improve stability of BCE loss calculation for input probabilities close to or exactly 0 or 1

[plpxsk]

Proposed changes

Predicted probabilities that are close to, or exactly 0 or 1, may result in
numerical instabilities or undefined values in a Binary Cross-Entropy (BCE)
loss calculation.

Currently, our BCE function will return nan in such scenarios.

If we clip the logs in the BCE loss calculation, we can always return a finite
loss value.

This is indeed the PyTorch approach [1]:

Our solution is that BCELoss clamps its log function outputs to be greater
than or equal to -100. This way, we can always have a finite loss value and a
linear backward method.

[1] https://pytorch.org/docs/stable/generated/torch.nn.BCELoss.html

This PR suggests and implements this log clipping in the BCE formula.

Implementation

Apply mx.clip(..., a_min=-100) to the mx.log() calculations in
the BCE loss function.

Following pytorch, and to always obtain a finite loss, clipping is always
done (and uses the same pytorch clip value of -100).

But if desired, we can add a parameter like clip=True

Note: both before and after this implementation, a nan will still be returned
for negative input probabilites.

Example

Current behavior:

TINY_PROB = 1e-59
probs = mx.array([0, TINY_PROB, 1-TINY_PROB, 1])
targets = mx.array([0, 0, 1, 1])

# array([nan, nan, nan, nan], dtype=float32)
nn.losses.binary_cross_entropy(probs, targets, with_logits=False, reduction="none")

# array(nan, dtype=float32)
nn.losses.binary_cross_entropy(probs, targets, with_logits=False, reduction="mean")

# array(nan, dtype=float32)
nn.losses.binary_cross_entropy(probs, targets, with_logits=False, reduction="sum")

New behavior:

TINY_PROB = 1e-59
probs = mx.array([0, TINY_PROB, 1-TINY_PROB, 1])
targets = mx.array([0, 0, 1, 1])

# array([-0, -0, -0, -0], dtype=float32)
nn.losses.binary_cross_entropy(probs, targets, with_logits=False, reduction="none")

# array(0, dtype=float32)
nn.losses.binary_cross_entropy(probs, targets, with_logits=False, reduction="mean")

# array(0, dtype=float32)
nn.losses.binary_cross_entropy(probs, targets, with_logits=False, reduction="sum")

New behavior matches pytorch:

import numpy as np
import torch
import torch.nn.functional as F


# tensor([0., 0., -0., -0.])
F.binary_cross_entropy(
    torch.tensor(np.array(probs), dtype=torch.float32),
    torch.tensor(np.array(targets), dtype=torch.float32),
    reduction="none")

# tensor(0.)
F.binary_cross_entropy(
    torch.tensor(np.array(probs), dtype=torch.float32),
    torch.tensor(np.array(targets), dtype=torch.float32),
    reduction="mean")

# tensor(0.)
F.binary_cross_entropy(
    torch.tensor(np.array(probs), dtype=torch.float32),
    torch.tensor(np.array(targets), dtype=torch.float32),
    reduction="sum")

References

PR that added support for probabilities: #492

Checklist

Put an x in the boxes that apply.

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed)

[awni]

Cool, thanks. Maybe we should in addition add a note about using the logits directly instead to the docstring? If you have values that extreme the best thing to do is avoid explicitly computing the probs (use with_logits=True). It will also be faster..

[plpxsk]

Sounds good, updated docstring as suggested

[awni]

Thanks! Will merge when tests clear.