Skip to content

← Flax NNX step 60 of 100

Medium Primitives

NNX Label Smoothing

Jump to the task

Why this matters

Label smoothing mixes a class label with a uniform distribution. Instead of asking for probability one on the labeled class, it gives every class some target probability. This can reduce overconfidence, but better calibration or accuracy is not guaranteed; choose the smoothing strength using validation data.

Smoothing changes the training objective. It does not prevent a model from producing extreme probabilities, and it does not guarantee nonzero gradients. The loss below is a plain JAX function that can be used with Flax or NNX.

The formula

For K classes and smoothing strength alpha:

target = (1 - alpha) * one_hot(label, K) + alpha / K
loss_per_example = -sum(target * log_softmax(logits))
loss = mean(loss_per_example)

The labeled class receives 1 - alpha + alpha / K; every other class receives alpha / K. These targets sum to one. This convention distributes the uniform part over all K classes, including the labeled class. A formula that assigns alpha / (K - 1) only to other classes is a different convention.

At alpha zero, this is ordinary cross-entropy. At alpha one, the target is uniform and independent of the label. It still penalizes nonuniform predictions. For finite logits and positive target probabilities, the ideal predicted probabilities equal the target. Only logit differences matter: ordinary cross-entropy does not require every incorrect logit to tend to minus infinity.

Worked mini-example

import jax
import jax.numpy as jnp

logits = jnp.array([[0.0, 0.0, 0.0, 0.0]], dtype=jnp.float32)
labels = jnp.array([2], dtype=jnp.int32)
alpha = 0.1
target = (1 - alpha) * jax.nn.one_hot(labels, 4) + alpha / 4
# target: [[0.025, 0.025, 0.925, 0.025]]
loss = -jnp.mean(jnp.sum(target * jax.nn.log_softmax(logits), axis=-1))
print(loss)  # approximately 1.3862944, or log(4)

Uniform predictions have the same cross-entropy for every target distribution that sums to one. For nonuniform predictions, smoothing generally changes the loss and its gradient.

Problem

Implement label_smoothing_loss(logits, labels_int, num_classes, alpha):

  1. Convert the integer-valued labels to int32.
  2. Build the smoothed targets using the formula above.
  3. Compute log probabilities with jax.nn.log_softmax(logits, axis=-1).
  4. Sum the negative weighted log probabilities over classes, then average over the batch.
  5. Return the mean reshaped to (1,), keeping it as a JAX array.

Inputs are finite float32 logits of shape (N, K), valid labels of shape (N,), an integer-valued num_classes equal to K, and alpha in [0, 1]. N and K are positive. Labels and num_classes may arrive as floating-point values. Return a float32 array of shape (1,) containing the mean loss.

Common pitfalls

  • Use log-softmax directly. Taking the logarithm after softmax can underflow to negative infinity for unlikely classes, even with finite logits. Log-softmax computes log probabilities, not softmax of logged logits.
  • Keep the final mean as a JAX value. Converting it with Python float breaks differentiation and JIT tracing. Use jnp.mean(losses).reshape(1).
  • If you use int(num_classes) for the one-hot width, mark num_classes static when applying jax.jit. Keep alpha in JAX arithmetic so it can remain a dynamic scalar.
  • JAX accepts integer-valued floating indices in one_hot, but explicitly casting labels makes the expected class-index contract clear. Do not silently treat fractional or out-of-range labels as valid inputs.
  • The equivalent loss is (1 - alpha) * CE + alpha * uniform_CE. Adding the uniform term to an unweighted CE gives a different objective.
  • Class reduction is a sum; batch reduction is a mean. Returning a scalar instead of a (1,) array does not meet this exercise's output contract.