Phase 04 / Math & optimization
Probability, likelihood, and stable losses
Revision v0.7.0 · Updated since your last visit. Review the changes; your progress is saved.
Lesson · about 80 minutes · CPU practice
Derive binary cross-entropy from a probability model and compute it safely from logits.
Phase study guide: preparation, milestones and project evidence
Before you start
Before this: Write gradient descent yourself
Hardware: CPU
Setup help: Python, notebook, or Cloud TPU VM
Use the tested CPU packages (the recorded environment uses Python 3.14.3). Other Python versions and accelerator environments require separate checks. The optional Cloud TPU specialization explains the target-specific execution boundaries. Logical CPU devices practice partitioning; they do not simulate TPU performance.
Download the repository ZIP or use your existing checkout. Extract it and open a terminal in the folder containing requirements-cpu.txt. The beginner ZIP is separate; these lessons use the full repository folder.
python3 -m venv .venv
source .venv/bin/activate
python -m pip install -r requirements-cpu.txtpy -3 -m venv .venv
.\.venv\Scripts\python.exe -m pip install -r requirements-cpu.txtAfter activation, run the complete worked example from that same folder:
python phases/04-optimization/09-probability-likelihood-and-stable-losses/code/main.py.\.venv\Scripts\python.exe phases/04-optimization/09-probability-likelihood-and-stable-losses/code/main.pyFor an incremental build, create main.py in your own practice folder, append each build block in order, and run python main.py with this environment active. Keep that file separate from the downloaded reference solution.
Download this lesson’s notebook. Open it with a Python kernel using the course environment. Run cells from the top after restarting the kernel; experiments depend on earlier definitions.
For local Jupyter, first install the course requirements using the Local Python steps, then use your environment’s Python to start Jupyter:
python -m pip install jupyterlab
python -m jupyter lab.\.venv\Scripts\python.exe -m pip install jupyterlab
.\.venv\Scripts\python.exe -m jupyter labSelect the downloaded optimization-09.ipynb file in JupyterLab. The notebook contains reference solutions; make your prediction and attempt each practice task before opening them. Restart the kernel before changing logical device setup.
The worked reference for this lesson is CPU-validated. Some references explicitly configure logical CPU devices; setting JAX_PLATFORMS does not adapt them to a TPU. Follow the execution matrix and reviewed multi-host example before attempting a target run.
For accelerator preparation, study TPU setup, observed runtime receipts, checkpointed jobs and precision and retained traces. These target paths remain unqualified on physical TPU hardware until an actual receipt is recorded. Cloud provisioning is optional for CPU learning.
Check: Uncertain loss about ; confident-wrong loss .
Revisit the guided CPU setup · Open 19: TPU Setup. Package instructions are CPU instructions; Windows and accelerator environments need their own validation.
What you will be able to do
- Derive binary cross-entropy from a probability model and compute it safely from logits.
- Read likelihood as support for the observed data
- Move from probabilities to logits
- Avoid overflowing intermediate values
- Extend the idea without overstating it
The problem
A classifier should tell us how strongly it believes a label, not only which label it chooses. How do we turn that belief into a loss? We will start with a yes-or-no event, derive the penalty for a prediction, and discover why an algebraically correct formula can fail on a computer.
The idea
A classification loss rewards probability assigned to the observed outcome. A logit is an unconstrained score that becomes a probability through a link such as the sigmoid. Stable formulas let us compute the same objective without overflowing or taking the logarithm of a rounded zero.
Connect a logit to the observed label
For a binary label of one, predicting probability gives loss . Predicting gives about . If the observed label is instead zero, that same prediction gives . Confidence helps only when it supports the observed outcome.
At logit zero, the sigmoid gives , so either binary label has the same loss. Moving the logit right helps the positive label and hurts the negative one. Follow the correct label's curve rather than treating a lower curve as universally better.
Extreme logits expose numerical differences between mathematically equivalent formulas. Prefer a stable logit-based loss and verify it against moderate-value hand calculations before exploring extreme cases.
Read likelihood as support for the observed data
The Bernoulli model writes the probability of an observed label as . If the label is , predictions and incur losses about and , respectively. Both can be compared even before applying a classification threshold. For independent observations conditional on their inputs and model parameters, likelihoods multiply. Taking negative logarithms turns the product into a sum; dividing by the sample count gives a mean.
Move from probabilities to logits
A logit is an unrestricted real score. The sigmoid maps it into a probability: . At , the probability is and either binary label has loss . For a positive label, differentiating the loss gives , so gradient descent increases the logit. More generally the derivative is . A very wrong, confident prediction still supplies a useful gradient.
Avoid overflowing intermediate values
Computing log(sigmoid(z)) directly can take a logarithm of a rounded zero. For and , the correct loss is close to , not infinity. Use jax.nn.softplus or an established logits-based loss. The equivalent expression below explains the stability: its exponential argument is never positive. We use the library softplus for differentiation rather than hand-assembling maximum and abs, whose derivative conventions at a tie need care.
Extend the idea without overstating it
For mutually exclusive classes, the probability of class is the softmax of its score . Use jax.nn.log_softmax and select the observed class. Adding the same constant to every score must not change the probabilities or loss. A low training cross-entropy is not proof of calibration or generalization; check held-out predictions. Squared error also has a likelihood interpretation under independent Gaussian errors with fixed variance, but that assumption differs from the Bernoulli model.
Confidence is rewarded only when it matches the label
Predict before looking: For a positive label, which direction should the logit move?
Read the figure
The horizontal axis is a logit, which becomes the predicted probability of label after a sigmoid. The vertical axis is binary cross-entropy loss. The solid curve assumes the observed label is ; the dashed curve assumes it is .
At logit , the predicted probability is , and both losses are . Moving right lowers the loss for observed label but raises it for observed label . At logit , the latter loss is about , while the former is close to zero.
Connect it to the computation
A large positive logit strongly favors label . That confidence is rewarded only if label is actually observed. For label , the same confidence is a large mistake. The mirror-image curves make the dependence on the target explicit.
Loss is not a probability or classification accuracy, so values above are valid. The near-flat parts approach zero without implying an exact zero probability for the other class. Computing this loss from logits with a stable formula avoids unnecessary overflow or taking the logarithm of a rounded zero.
Open full-size figureDownload data and provenance
The notebook contains the data experiment, plotting code, and recorded figure output. Rerun it to test your changes.
Where to put & run this code: Work in phases/04-optimization/09-probability-likelihood-and-stable-losses/code/main.py (or downloaded exercises/optimization-09.py / exercises/optimization-09.ipynb) and execute from the repository root with python3 phases/04-optimization/09-probability-likelihood-and-stable-losses/code/main.py.
1 Write a loss from logits
Create main.py. Start at an uncertain prediction so you can verify the answer by hand.
# Step 1 — Write a loss from logits: A balanced probability assigns loss \log 2 to either outcome.
import jax
import jax.numpy as jnp
import optax
def binary_loss(z, y):
return jax.nn.softplus(z) - y * z
z = jnp.array(0.0)
y = jnp.array(1.0)
assert jnp.allclose(binary_loss(z, y), jnp.log(2.0))Code building blocks in this step
jnp.array(values, dtype=...)— Constructs an immutable device-backed JAX array from Python/NumPy values.jnp.allclose(actual, expected, rtol=..., atol=...)— Checks that two arrays match elementwise within floating-point tolerance.
A balanced probability assigns loss to either outcome.
2 Check the derivative independently
Append the analytic gradient and a library comparison at moderate scores.
# Step 2 — Check the derivative independently: The derivative points toward assigning more probability to the...
assert jnp.allclose(jax.grad(binary_loss)(z, y), -0.5)
logits = jnp.array([-2.0, 0.0, 2.0])
labels = jnp.array([0.0, 1.0, 1.0])
analytic = jax.nn.sigmoid(logits) - labels
# Differentiate the objective to obtain `actual` via automatic differentiation.
actual = jax.vmap(jax.grad(binary_loss))(logits, labels)
assert jnp.allclose(actual, analytic, atol=1e-06)
assert jnp.allclose(binary_loss(logits, labels), optax.sigmoid_binary_cross_entropy(logits, labels), atol=1e-06)Code building blocks in this step
jnp.array(values, dtype=...)— Constructs an immutable device-backed JAX array from Python/NumPy values.jnp.allclose(actual, expected, rtol=..., atol=...)— Checks that two arrays match elementwise within floating-point tolerance.jax.grad(loss_fn)(params, ...)— Transforms a scalar-output function into a function returning the gradient PyTree with the same structure asparams.jax.vmap(fn, in_axes=..., out_axes=...)— Vectorizes a single-example function across a batch axis without writing a Python loop.
The derivative points toward assigning more probability to the observed label.
3 Check a confident mistake
Append the extreme case and run python main.py.
# Step 3 — Check a confident mistake: The confident mistake has a large finite loss and derivative near -1.
assert jnp.allclose(binary_loss(jnp.array(-100.0), jnp.array(1.0)), 100.0)
assert jnp.allclose(jax.grad(binary_loss)(jnp.array(-100.0), jnp.array(1.0)), -1.0)
print('uncertain loss / confident-wrong loss:', binary_loss(z, y), binary_loss(-100.0, 1.0))Code building blocks in this step
jnp.array(values, dtype=...)— Constructs an immutable device-backed JAX array from Python/NumPy values.jnp.allclose(actual, expected, rtol=..., atol=...)— Checks that two arrays match elementwise within floating-point tolerance.jax.grad(loss_fn)(params, ...)— Transforms a scalar-output function into a function returning the gradient PyTree with the same structure asparams.
The confident mistake has a large finite loss and derivative near .
View the cumulative build in main.py
These blocks append in order to the same file.
# Step 1 — Write a loss from logits: A balanced probability assigns loss \log 2 to either outcome.
import jax
import jax.numpy as jnp
import optax
def binary_loss(z, y):
return jax.nn.softplus(z) - y * z
z = jnp.array(0.0)
y = jnp.array(1.0)
assert jnp.allclose(binary_loss(z, y), jnp.log(2.0))
# Step 2 — Check the derivative independently: The derivative points toward assigning more probability to the...
assert jnp.allclose(jax.grad(binary_loss)(z, y), -0.5)
logits = jnp.array([-2.0, 0.0, 2.0])
labels = jnp.array([0.0, 1.0, 1.0])
analytic = jax.nn.sigmoid(logits) - labels
# Differentiate the objective to obtain `actual` via automatic differentiation.
actual = jax.vmap(jax.grad(binary_loss))(logits, labels)
assert jnp.allclose(actual, analytic, atol=1e-06)
assert jnp.allclose(binary_loss(logits, labels), optax.sigmoid_binary_cross_entropy(logits, labels), atol=1e-06)
# Step 3 — Check a confident mistake: The confident mistake has a large finite loss and derivative near -1.
assert jnp.allclose(binary_loss(jnp.array(-100.0), jnp.array(1.0)), 100.0)
assert jnp.allclose(jax.grad(binary_loss)(jnp.array(-100.0), jnp.array(1.0)), -1.0)
print('uncertain loss / confident-wrong loss:', binary_loss(z, y), binary_loss(-100.0, 1.0))Run the example
python3 phases/04-optimization/09-probability-likelihood-and-stable-losses/code/main.pyAbout this command
This shortcut uses an activated macOS/Linux course environment in the repository folder. For Windows, use the explicit interpreter command in “Setup help” above. The script includes worked examples and reference solutions.
Inspect the complete assembled example
# Step 1 — Write a loss from logits: A balanced probability assigns loss \log 2 to either outcome.
import jax
import jax.numpy as jnp
import optax
def binary_loss(z, y):
return jax.nn.softplus(z) - y * z
z = jnp.array(0.0)
y = jnp.array(1.0)
assert jnp.allclose(binary_loss(z, y), jnp.log(2.0))
# Step 2 — Check the derivative independently: The derivative points toward assigning more probability to the...
assert jnp.allclose(jax.grad(binary_loss)(z, y), -0.5)
logits = jnp.array([-2.0, 0.0, 2.0])
labels = jnp.array([0.0, 1.0, 1.0])
analytic = jax.nn.sigmoid(logits) - labels
# Differentiate the objective to obtain `actual` via automatic differentiation.
actual = jax.vmap(jax.grad(binary_loss))(logits, labels)
assert jnp.allclose(actual, analytic, atol=1e-06)
assert jnp.allclose(binary_loss(logits, labels), optax.sigmoid_binary_cross_entropy(logits, labels), atol=1e-06)
# Step 3 — Check a confident mistake: The confident mistake has a large finite loss and derivative near -1.
assert jnp.allclose(binary_loss(jnp.array(-100.0), jnp.array(1.0)), 100.0)
assert jnp.allclose(jax.grad(binary_loss)(jnp.array(-100.0), jnp.array(1.0)), -1.0)
print('uncertain loss / confident-wrong loss:', binary_loss(z, y), binary_loss(-100.0, 1.0))Expected result: Uncertain loss about ; confident-wrong loss .
Recorded reference execution
Reference output captured from executing the complete lesson script (code/main.py) and its assertions in the course CPU environment.
uncertain loss / confident-wrong loss: 0.6931472 100.0
uncertain loss / confident-wrong loss: 0.6931472 100.0
PASS: optimization-09
Download execution record · The notebook includes executed cells and figure outputs. Compare your run with these values; timings depend on the machine.
Break the probability-first implementation
# Experiment — Break the probability-first implementation: A stable final answer requires stable intermediate computations.
p = jax.nn.sigmoid(jnp.array(100.0))
naive = -jnp.log(1 - p)
stable = binary_loss(jnp.array(100.0), jnp.array(0.0))
# Confirm that all computed values remain finite (no NaN or Inf).
assert not jnp.isfinite(naive)
assert jnp.isfinite(stable) and jnp.allclose(stable, 100.0)Expected: The naive loss is infinite; the logits-based loss remains finite.
A stable final answer requires stable intermediate computations.
Shift every class score
# Experiment — Shift every class score: Only relative class scores matter.
scores = jnp.array([1.0, 2.0, 3.0])
target = 2
# Evaluate numerically stable log-space cross-entropy/likelihood (`nll`).
nll = -jax.nn.log_softmax(scores)[target]
# Evaluate numerically stable log-space cross-entropy/likelihood (`shifted`).
shifted = -jax.nn.log_softmax(scores + 1000.0)[target]
assert jnp.allclose(nll, shifted, atol=1e-06)
reference_loss = jnp.log(jnp.exp(-2.0) + jnp.exp(-1.0) + 1.0)
assert jnp.allclose(nll, reference_loss, atol=1e-06)Expected: Both losses are about .
Only relative class scores matter. Stable log-softmax respects this invariance.
Make it yours
Where to put your code: Append your solution at the bottom of phases/04-optimization/09-probability-likelihood-and-stable-losses/code/main.py (or inside the Make it yours cell of exercises/optimization-09.ipynb) and run python3 phases/04-optimization/09-probability-likelihood-and-stable-losses/code/main.py to verify your prediction before opening the reference solution.
Compute the binary losses for positive labels at probabilities and , and explain which prediction is better supported.
Show a writing guide and starter scaffold
How to write this exercise code: step-by-step recipe & starter scaffold
1. Key functions & syntax to use:
jnp.array(values, dtype=...)— Constructs an immutable device-backed JAX array from Python/NumPy values.jnp.allclose(actual, expected, rtol=..., atol=...)— Checks that two arrays match elementwise within floating-point tolerance.
2. Step-by-step implementation plan:
- Set up the changed input or configuration parameters requested in the prompt.
- Run the target function or transformation and assert that the output matches your prediction.
# Exercise solution: Compute the binary losses for positive labels at probabilities 0.8 and...
probabilities = jnp.array(...) # TODO: compute probabilities
z_from_p = jnp.log(...) # TODO: compute z_from_p
assert jnp.allclose(binary_loss(z_from_p, 1.0), -jnp.log(probabilities), atol=1e-06) # TODO: complete assertion check
assert binary_loss(z_from_p[0], 1.0) # TODO: complete assertion check
Reveal the reference solution
# Exercise solution: Compute the binary losses for positive labels at probabilities 0.8 and...
probabilities = jnp.array([0.8, 0.2])
z_from_p = jnp.log(probabilities) - jnp.log1p(-probabilities)
assert jnp.allclose(binary_loss(z_from_p, 1.0), -jnp.log(probabilities), atol=1e-06)
assert binary_loss(z_from_p[0], 1.0) < binary_loss(z_from_p[1], 1.0)Fit a constant probability
Transfer / diagnosis
A dataset contains labels . For a model with one constant logit, show that is stationary and has less mean loss than .
Show a writing guide and starter scaffold
How to write this practice solution: step-by-step recipe & starter scaffold
1. Key functions & syntax to use:
jnp.array(values, dtype=...)— Constructs an immutable device-backed JAX array from Python/NumPy values.x.sum(axis=...) / x.mean(axis=..., keepdims=...)— Reduces values along the namedaxis(the axis you name is collapsed unlesskeepdims=True).jnp.allclose(actual, expected, rtol=..., atol=...)— Checks that two arrays match elementwise within floating-point tolerance.jax.grad(loss_fn)(params, ...)— Transforms a scalar-output function into a function returning the gradient PyTree with the same structure asparams.
2. Step-by-step implementation plan:
- Aggregate array values to compute
constant_loss.
# Fit a constant probability (Transfer / diagnosis): The constant model matches the observed class frequency.
ys = jnp.array(...) # TODO: compute ys
# Aggregate array values to compute `constant_loss`.
constant_loss = ... # TODO: compute constant_loss
optimum = jnp.log(...) # TODO: compute optimum
assert jnp.allclose(jax.grad(constant_loss)(optimum), 0.0, atol=1e-06) # TODO: complete assertion check
assert constant_loss(optimum) # TODO: complete assertion check
Show a hint
The mean gradient is . Convert the desired probability to log-odds.
Reveal reference solution and reasoning
# Fit a constant probability (Transfer / diagnosis): The constant model matches the observed class frequency.
ys = jnp.array([1.0, 1.0, 1.0, 0.0])
# Aggregate array values to compute `constant_loss`.
constant_loss = lambda z: jnp.mean(binary_loss(z, ys))
optimum = jnp.log(3.0)
assert jnp.allclose(jax.grad(constant_loss)(optimum), 0.0, atol=1e-06)
assert constant_loss(optimum) < constant_loss(jnp.array(0.0))The constant model matches the observed class frequency. It has learned no feature-dependent discrimination.
Catch a reduction mistake
Transfer / diagnosis
Duplicate every observation. Verify that mean loss stays fixed while total loss doubles. What happens to gradients?
Show a writing guide and starter scaffold
How to write this practice solution: step-by-step recipe & starter scaffold
1. Key functions & syntax to use:
jnp.array(values, dtype=...)— Constructs an immutable device-backed JAX array from Python/NumPy values.x.sum(axis=...) / x.mean(axis=..., keepdims=...)— Reduces values along the namedaxis(the axis you name is collapsed unlesskeepdims=True).jnp.allclose(actual, expected, rtol=..., atol=...)— Checks that two arrays match elementwise within floating-point tolerance.jax.grad(loss_fn)(params, ...)— Transforms a scalar-output function into a function returning the gradient PyTree with the same structure asparams.
2. Step-by-step implementation plan:
- Differentiate the objective to obtain
mean_gradvia automatic differentiation. - Differentiate the objective to obtain
duplicated_gradvia automatic differentiation.
# Catch a reduction mistake (Transfer / diagnosis): Reduction conventions affect gradient scale and...
zs = jnp.array(...) # TODO: compute zs
ys2 = jnp.array(...) # TODO: compute ys2
original = binary_loss(...) # TODO: compute original
doubled = binary_loss(...) # TODO: compute doubled
assert jnp.allclose(jnp.mean(original), jnp.mean(doubled)) # TODO: complete assertion check
assert jnp.allclose(jnp.sum(doubled), 2 * jnp.sum(original)) # TODO: complete assertion check
# Differentiate the objective to obtain `mean_grad` via automatic differentiation.
mean_grad = ... # TODO: compute mean_grad
# Differentiate the objective to obtain `duplicated_grad` via automatic differentiation.
duplicated_grad = jax.grad(...) # TODO: compute duplicated_grad
assert jnp.allclose(mean_grad(zs), duplicated_grad) # TODO: complete assertion check
Show a hint
The mean divides by the new observation count; the sum does not.
Reveal reference solution and reasoning
# Catch a reduction mistake (Transfer / diagnosis): Reduction conventions affect gradient scale and...
zs = jnp.array([-1.0, 2.0])
ys2 = jnp.array([0.0, 1.0])
original = binary_loss(zs, ys2)
doubled = binary_loss(jnp.tile(zs, 2), jnp.tile(ys2, 2))
assert jnp.allclose(jnp.mean(original), jnp.mean(doubled))
assert jnp.allclose(jnp.sum(doubled), 2 * jnp.sum(original))
# Differentiate the objective to obtain `mean_grad` via automatic differentiation.
mean_grad = lambda z: jax.grad(lambda t: jnp.mean(binary_loss(t * z, ys2)))(1.0)
# Differentiate the objective to obtain `duplicated_grad` via automatic differentiation.
duplicated_grad = jax.grad(lambda t: jnp.mean(binary_loss(t * jnp.tile(zs, 2), jnp.tile(ys2, 2))))(1.0)
assert jnp.allclose(mean_grad(zs), duplicated_grad)Reduction conventions affect gradient scale and learning-rate interpretation; record them with the objective.
Diagnose the result
If loss becomes infinite, inspect logits and intermediate probabilities before clipping targets or shrinking the learning rate. Check label range, class axis and reduction. Keep stability checks separate from calibration and accuracy measurements.
Carry forward
- A stable final answer requires stable intermediate computations.
- Only relative class scores matter. Stable log-softmax respects this invariance.
Keep your evidence
Keep the hand log-loss, analytic gradient, extreme-logit failure/repair, multiclass shift check and empirical-probability calculation.
Save your modified code, its output, and an explanation of your prediction. Checkpoint answers and exercise evidence are tracked separately.
Saved in this browser. Export or back up your portfolio anytime in My learning.
Verified on local CPU. Stage to a Cloud TPU VM via Cloud TPU setup to record accelerator execution receipts (TODO: add automated Cloud TPU CI execution receipts).