ranking_softmax_loss

ranking_softmax_loss#

class braintools.metric.ranking_softmax_loss(logits, labels, *, where=None, weights=None, reduction=None, reduce_fn=<function mean>)#

Compute ranking softmax loss for learning-to-rank applications.

Calculates a differentiable ranking loss that measures the cost of a ranking induced by item scores compared to ground truth relevance labels. This loss is particularly effective for information retrieval, recommendation systems, and other ranking tasks where the goal is to prioritize relevant items.

The loss is the negative sum of the per-item log_softmax over the scores, each weighted by its raw relevance label (the Rax / softmax-loss convention):

\[\ell(s, y) = -\sum_{i=1}^{n} y_i \log \frac{\exp(s_i)}{\sum_{j=1}^{n} \exp(s_j)}\]

where \(s_i\) are the logit scores, \(y_i\) are the relevance labels, and \(n\) is the number of items in the list. Note that the labels \(y_i\) enter as plain multiplicative coefficients on \(\log\mathrm{softmax}(s)_i\); they are not normalized into a probability distribution (this is the softmax-loss convention, not ListNet, which would first apply a softmax over the labels).

Parameters:

  • logits (Array | ndarray | bool | number | bool | int | float | complex | Quantity) – Predicted scores for each item with shape (..., list_size). Higher scores should indicate higher relevance. The function operates on the last dimension, treating leading dimensions as batch dimensions.

  • labels (Array | ndarray | bool | number | bool | int | float | complex | Quantity) – Ground truth relevance labels with shape (..., list_size). Typically non-negative values where higher values indicate greater relevance. These are used as raw multiplicative weights on the per-item log_softmax of the logits; they are not normalized into a probability distribution.

  • where (Array | ndarray | bool | number | bool | int | float | complex | Quantity | None) – Boolean mask with shape (..., list_size) indicating valid items. Items where where is False are excluded from loss computation. This is useful for handling variable-length lists or missing data.

  • weights (Array | ndarray | bool | number | bool | int | float | complex | Quantity | None) – Per-item weights with shape (..., list_size) for emphasizing certain items in the loss calculation. Applied to labels before computing the softmax cross-entropy.

  • reduction (str | None) – How to reduce the per-list losses across the leading (batch) dimensions. One of 'none', 'mean' or 'sum'. When provided, this takes precedence over reduce_fn. Lists in which every item is masked out (empty where row) contribute 0 and, for 'mean', are excluded from the average. If None (default), the legacy reduce_fn argument is used instead.

  • reduce_fn (Callable[..., Array | ndarray | bool | number | bool | int | float | complex | Quantity] | None) –

    Legacy reduction callable kept for backward compatibility. Used only when reduction is None. Common choices:

    • jax.numpy.mean (default): Average loss across lists

    • jax.numpy.sum: Sum loss across lists

    • None: Return unreduced per-list losses

    Prefer the reduction string for new code.

Returns:

Ranking softmax loss. Shape depends on the requested reduction:

  • 'mean'/'sum' (or a non-None reduce_fn): scalar loss value

  • 'none' (or reduce_fn=None): array with shape equal to all the leading (batch) dimensions of logits, i.e. logits.shape[:-1]

Return type:

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

Notes

This loss function implements a probabilistic approach to ranking where:

  • Items with higher relevance labels should receive higher probability mass

  • The softmax operation ensures valid probability distributions

  • Masked items (where where is False) are effectively ignored

  • The loss is differentiable w.r.t. logits, enabling gradient-based optimization

The function handles edge cases gracefully:

  • Empty masks (all items invalid) return 0.0 instead of NaN

  • Numerical stability is maintained through log-softmax computation

  • Mixed data types are handled by casting labels to match logit precision

Examples

Basic ranking loss with single query:

>>> import jax.numpy as jnp
>>> import braintools as braintools
>>> # Scores for 3 items
>>> logits = jnp.array([2.0, 1.0, 3.0])
>>> # Relevance: item 3 most relevant, item 1 second, item 2 least
>>> labels = jnp.array([1.0, 0.0, 2.0])
>>> loss = braintools.metric.ranking_softmax_loss(logits, labels)
>>> print(f"Loss: {loss:.4f}")
Loss: 2.2228

Batch processing with masking:

>>> # Batch of 2 queries with 3 items each
>>> logits = jnp.array([[2.0, 1.0, 0.0], [1.0, 0.5, 1.5]])
>>> labels = jnp.array([[1.0, 0.0, 0.0], [0.0, 0.0, 1.0]])
>>> # Second query only has first 2 items valid
>>> where = jnp.array([[True, True, False], [True, True, True]])
>>> loss = braintools.metric.ranking_softmax_loss(logits, labels, where=where)
>>> print(f"Batch loss: {loss:.4f}")
Batch loss: 0.4968

Per-item weighting:

>>> weights = jnp.array([1.0, 2.0, 1.0])  # Emphasize middle item
>>> loss = braintools.metric.ranking_softmax_loss(logits[0], labels[0], weights=weights)
>>> print(f"Weighted loss: {loss:.4f}")
Weighted loss: 0.4076

Unreduced losses for analysis (legacy reduce_fn API):

>>> batch_losses = braintools.metric.ranking_softmax_loss(
...     logits, labels, where=where, reduce_fn=None
... )
>>> print(f"Individual losses: {batch_losses}")
Individual losses: [0.31326163 0.68026966]

Using the reduction string API (preferred):

>>> per_list = braintools.metric.ranking_softmax_loss(
...     logits, labels, reduction='none'
... )
>>> total = braintools.metric.ranking_softmax_loss(
...     logits, labels, reduction='sum'
... )
>>> print(jnp.allclose(jnp.sum(per_list), total))
True

See also

jax.nn.log_softmax

Underlying log-softmax computation

braintools.metric.softmax_cross_entropy

Related cross-entropy function

braintools.metric.sigmoid_binary_cross_entropy

Alternative for binary relevance

References