Masked Cross-Entropy with Label Smoothing - Neural Networks

Masked Cross-Entropy with Label Smoothing

What you'll build

Token-level cross-entropy from logits, with a padding mask (ignore_index) and label smoothing, matching PyTorch's F.cross_entropy(logits, targets, ignore_index=..., label_smoothing=...) with mean reduction.

Background

Language-model training computes a loss per token, but batches are padded to a common length and padding tokens must not count. Many training recipes also smooth the labels: instead of a one-hot target, the model is trained towards (1 - eps) on the true class plus eps spread over all classes, which discourages over-confident logits and often improves calibration.

Formula

For one token with logits z (V classes) and target class y:

log_p = z - logsumexp(z)                        # log-softmax, computed stably
q_c   = (1 - eps) * [c == y] + eps / V          # smoothed target: eps spread over ALL V classes
loss  = -sum_c q_c * log_p_c
      = (1 - eps) * (-log_p_y) + eps * mean_c(-log_p_c)

The final loss is the mean over non-ignored tokens only:

L = sum_{tokens with y != ignore_index} loss / (number of such tokens)

where logsumexp(z) = m + log(sum_c exp(z_c - m)) with m = max(z).

Inputs

• logits: float array (N, V) or (B, T, V)
• targets: int array (N,) or (B, T), matching the leading dimensions of logits
• ignore_index: target value that marks padding (default -100; may also be a real class id such as a pad id of 0)
• smoothing: label smoothing eps in [0, 1] (default 0)

Expected return format

A Python float: the mean loss over non-ignored tokens.

Edge cases and conventions

• Flatten the leading dimensions to (N, V) and (N,) first.
• Select the kept tokens before indexing by target: -100 is not a valid class index.
• If every token is ignored, return nan (what PyTorch's mean reduction gives).
• With smoothing = 0 this is the ordinary cross-entropy.
• Logits of magnitude 1000 must not overflow: use the max-shift inside logsumexp.

Constraints

NumPy only, vectorised (no Python loop over tokens).

Follow-up

• PyTorch's reduction='sum' divides by nothing; why do LM trainers often normalise by the number of tokens across the whole (accumulated) batch instead of per micro-batch?
• What is the minimum achievable loss with smoothing eps > 0? (It is the entropy of q, not 0.)

Input: masked_label_smoothing_ce(np.array([[2.0, 1.0, 0.1], [0.5, 2.5, -1.0], [1.0, 1.0, 1.0], [-2.0, 0.0, 3.0]]), np.array([0, 1, -100, 2]))
Output: 0.2084 (rounded)

Input: masked_label_smoothing_ce(np.array([[2.0, 1.0, 0.1], [0.5, 2.5, -1.0], [1.0, 1.0, 1.0], [-2.0, 0.0, 3.0]]), np.array([0, 1, -100, 2]), smoothing=0.1)
Output: 0.3906 (rounded)

import numpy as np

def masked_label_smoothing_ce(logits, targets, ignore_index=-100, smoothing=0.0):
    """
    Token-level cross-entropy with label smoothing and a padding mask.

    Args:
        logits: Unnormalized scores, array (N, V) or (B, T, V)
        targets: Class ids, int array (N,) or (B, T); ignore_index marks padding
        ignore_index: Target value whose tokens are excluded from the loss
        smoothing: Label smoothing eps in [0, 1]

    Returns:
        Mean loss over non-ignored tokens (float); nan if every token is ignored
    """
    # Your code here
    pass