Stable categorical and binary classification losses

Your task

Function signature

stable_classification_loss(
    logits: torch.Tensor,
    targets: torch.Tensor,
    mask: torch.Tensor | None = None,
    weights: torch.Tensor | None = None,
    variant: str = 'categorical',
    reduction: str = 'mean',
    dim: int = -1,
    eps: float = 1e-08,
    ignore_index: int = -100,
) -> torch.Tensor

Worked example

A correct target with probability three quarters

Parameters

Return

torch.Tensor · Exactly targets.shape for reduction="none"; scalar tensor [] for "sum" or "mean".

Required behavior

• Implement stable categorical cross-entropy and binary cross-entropy from logits, with optional masks and differentiable sample weights.
• variant is 'categorical' (default) or 'binary'.
• In categorical mode logits has rank R\ge 1 with C\ge 1 categories on dim (integer, default -1; negative axes count from the end). int64 targets has logits.shape with that axis removed and stores a class in [0,C) on valid positions. ignore_index defaults to -100 and always excludes that target.
• Optional bool mask broadcasts to targets.shape; True means valid. Validity is mask AND \mathrm{target}\ne \mathrm{ignore\_index}. Masked/ignored labels may be arbitrary integers, including out-of-range padding: replace them safely before any gather.
• Per-valid-example loss is \mathrm{ell}=\operatorname{log}(\sum_{c} \operatorname{exp}(z_{c}))-z_{\mathrm{target}}, evaluated through shifted logits for stability. All C categories participate for a valid example; there is no class mask, class weighting, label smoothing or causal shift in this primitive.
• In binary mode logits and floating targets have the SAME arbitrary shape, including scalar tensors or zero-length batch dimensions; target t lies in [0,1] at valid entries, including soft labels. dim and ignore_index do not apply. Optional bool mask broadcasts to that shape.
• Per-valid-element loss is \mathrm{ell}=\operatorname{max}(z,0)-z\cdot t+\operatorname{log}(1+\operatorname{exp}(-\left|z\right|)). Do not log a sigmoid: confident incorrect logits must have finite losses and nonvanishing gradients. The mathematical derivative at z=0 is 0.5-t.
• Masked logits/targets may contain arbitrary finite padding, and must contribute neither loss nor gradients.
• In either mode optional weights is nonnegative, floating, shares logits dtype/device, and broadcasts to targets.shape (not the categorical class axis); omitted weights means one, omitted mask means all True. Floating logits/targets/weights share dtype/device and all required outputs must be representable.
• Define v_i as final validity, q_{i}=v_{i}\cdot w_{i}\cdot \mathrm{ell\_i}, S=\sum_{i} q_{i}, D=\sum_{i}(v_{i}\cdot w_{i}). reduction is 'none', 'sum', or 'mean' (default).
• Return the targets-shaped tensor q for none, scalar S for sum, or scalar S/(D+\varepsilon) for mean when D>0 and an exact zero scalar when D=0. The scalar reductions are GLOBAL over every example axis, not an average of example means.
• eps defaults to 10^{-8}, is finite and \ge 0, and is added once after summing weights only in the mean denominator; it does not modify logits, probabilities or any unreduced loss. All-masked/ignored batches return zero with zero gradients.
• Zero-length example axes are permitted in either variant: none returns the empty target shape, and \sum/\operatorname{mean} return zero. The mean's D=0 branch is constant zero. Support noncontiguous layouts, broadcasting and singleton categories, and never modify inputs.

Why this matters

Starter code

You write the body of this function in the BioTorch workspace. Your code runs in an isolated PyTorch sandbox; Run checks the public examples and Submit grades the full test suite, including gradients.

import torch


def stable_classification_loss(
    logits: torch.Tensor,  # Categorical: any shape with class axis dim of size C >= 1. Binary: exactly targets.shape.
    targets: torch.Tensor,  # Categorical: logits.shape with the class axis removed. Binary: exactly logits.shape.
    mask: torch.Tensor | None = None,  # Any shape broadcastable to targets.shape.
    weights: torch.Tensor | None = None,  # Any shape broadcastable to targets.shape.
    variant: str = 'categorical',
    reduction: str = 'mean',
    dim: int = -1,
    eps: float = 1e-08,
    ignore_index: int = -100,
) -> torch.Tensor:
    """Return a tensor with shape targets.shape for reduction="none"; scalar tensor () for
        "sum" or "mean"."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement stable_classification_loss")

Related exercises

• Builds onCategorical Cross-Entropy
• Builds onBinary Cross-Entropy with Logits

Sources and further reading

• aqlaboratory/openfold · loss.py
• facebookresearch/esm · modules.py