nll_loss

Contents

nll_loss#

class braintools.metric.nll_loss(log_probs, target)#

Compute negative log likelihood loss for classification.

The negative log likelihood (NLL) loss is a standard loss function for training classification models. It expects log-probabilities as input and computes the negative log-likelihood of the correct class.

The loss is computed as:

\[\ell(x, y) = -x_{y}\]

where \(x\) contains log-probabilities and \(y\) is the target class index.

Parameters:

  • log_probs (Array | ndarray | bool | number | bool | int | float | complex | Quantity) –

    Log-probabilities of each class. The class dimension is the last axis for the 1-D/2-D cases and axis 1 for higher-dimensional inputs (PyTorch convention). Expected shapes:

    • (num_classes,) for single sample

    • (batch_size, num_classes) for batch processing

    • (batch_size, num_classes, d1, d2, ..., dK) for higher-dimensional inputs (e.g., per-pixel classification for images)

  • target (Array | ndarray | bool | number | bool | int | float | complex | Quantity) –

    Class indices in the range [0, num_classes-1]. Expected shapes:

    • () (scalar) for single sample

    • (batch_size,) for batch processing

    • (batch_size, d1, d2, ..., dK) for higher-dimensional targets

Returns:

Negative log likelihood loss values:

  • Scalar for single sample

  • (batch_size,) for batch processing

  • (batch_size, d1, d2, ..., dK) for higher-dimensional inputs

Return type:

Array | ndarray | bool | number | bool | int | float | complex | Quantity

Examples

>>> import jax.numpy as jnp
>>> import braintools
>>> # Single sample example
>>> log_probs = jnp.log(jnp.array([0.1, 0.7, 0.2]))
>>> target = 1  # Correct class is index 1
>>> loss = braintools.metric.nll_loss(log_probs, target)
>>> print(f"NLL loss: {loss:.4f}")
NLL loss: 0.3567
>>> # Batch example
>>> log_probs_batch = jnp.log(jnp.array([[0.1, 0.7, 0.2], [0.3, 0.3, 0.4]]))
>>> targets_batch = jnp.array([1, 2])
>>> losses = braintools.metric.nll_loss(log_probs_batch, targets_batch)
>>> print(losses)
[0.35667497 0.9162907 ]

Notes

This function expects log-probabilities as input, not raw logits or probabilities. Use jax.nn.log_softmax to convert logits to log-probabilities, or jnp.log to convert probabilities.

For end-to-end training with logits, consider using softmax_cross_entropy which combines softmax and cross-entropy in a numerically stable way. In fact, nll_loss(jax.nn.log_softmax(logits), labels) is equal to softmax_cross_entropy_with_integer_labels(logits, labels).

Out-of-range integer target values are gathered via jnp.take_along_axis and are silently clamped by JAX (not validated), so ensure targets lie within [0, num_classes).

Raises:

ValueError – If the dimensions of log_probs and target are incompatible.

See also

softmax_cross_entropy

Cross entropy loss starting from logits

softmax_cross_entropy_with_integer_labels

Efficient version for integer labels

Contents