Binary Cross-Entropy with Logits

The problem

Binary cross-entropy is the loss for yes/no predictions: is this answer correct, is this text safe, does the reward model prefer answer A over answer B. The network outputs one logit $x$ per question, any real number, and the probability of "yes" is its sigmoid:

where $x$ is the logit and $q$ is a number between 0 and 1: large positive logits give almost 1, large negative ones almost 0, and $\sigma(0) = 0.5$.

The target $y$ is 1 for "yes" and 0 for "no". The loss is minus the log of the probability given to the right answer: $-\log q$ when $y = 1$, and $-\log(1 - q)$ when $y = 0$. In one formula:

where one of the two terms is always 0: with $y = 1$ only the first remains, with $y = 0$ only the second.

Example. $x = 2$, so $q = \sigma(2) \approx 0.881$.

• If $y = 1$: $\ell = -\log 0.881 \approx 0.127$. The model said "probably yes", and it was yes: small loss.
• If $y = 0$: $\ell = -\log(1 - 0.881) = -\log 0.119 \approx 2.127$. It said "probably yes", and it was no: large loss.

Why the formula breaks, and the fix

For large logits the sigmoid rounds to exactly 1: $\sigma(40)$ is 1.0 in float64. With $y = 0$ the loss then computes $\log(1 - 1) = \log 0 = -\infty$. The same loss can be rewritten without ever computing $q$ (derivation on the theory page):

where $|x|$ is the absolute value of $x$. The exponent $-|x|$ is never positive, so $e^{-|x|}$ is between 0 and 1: nothing overflows, and the log never sees 0.

Example.

• $x = 2$, $y = 1$: $\max(2, 0) - 2 \cdot 1 + \log(1 + e^{-2}) = 0 + 0.127 = 0.127$, the same as above.
• $x = 40$, $y = 0$: $40 - 0 + \log(1 + e^{-40}) = 40$. The naive formula returned inf.

The gradient

As for cross-entropy, it is prediction minus target. With the mean over $n$ elements:

where $x_i$ and $y_i$ are element $i$ of the logits and targets, and $n$ is the number of elements.

Example. x = [2, -3, 0], y = [1, 1, 0.5], so $n = 3$.

• Losses: $0.127$, $3.049$, $0.693$; the mean is about $1.290$.
• Gradient: $\sigma(x) \approx (0.881, 0.047, 0.5)$, minus $y$, divided by 3: about $(-0.040, -0.318, 0)$. The second logit gets the biggest push: it said "no" with confidence, and the answer was yes.

The third target, $0.5$, is a soft target: "half yes". Any number between 0 and 1 works in the same formulas.

What to implement

For arrays x (logits) and y (targets between 0 and 1) of the same shape:

• sigmoid(x): elementwise, finite for any x.
• bce_with_logits(x, y): the mean of $\ell$ over all elements.
• bce_with_logits_backward(x, y): the gradient with respect to x.

Requirements

• bce_with_logits(np.array([-800.0]), np.array([1.0])) is $800$, not inf or nan.
• Soft targets must work.
• Use np.log1p(t) for $\log(1 + t)$: it stays accurate when $t$ is tiny. For example np.log(1 + 1e-20) returns 0.0, because 1 + 1e-20 rounds to 1, while np.log1p(1e-20) returns the correct 1e-20.

Hint

For the sigmoid, $\sigma(x) = \tfrac{1}{2}\left(1 + \tanh(x/2)\right)$ is exact and never computes a huge exponential. (1 / (1 + np.exp(-x)) also gives the right value, but on a desktop NumPy warns that np.exp(800) overflowed.)