Implement Masked Cross-Entropy with Label Smoothing

Problem Statement

In autoregressive language model training, the loss function must handle sequence padding via boolean/binary masks, avoid overconfidence via label smoothing, and maintain numerical stability across mixed-precision regimes.

Implement the function masked_smoothed_cross_entropy in pure NumPy. The function computes the smoothed cross-entropy loss averaged across all unmasked tokens and calculates the analytical gradient with respect to the input logits.

Mathematical Formulation

For a vocabulary of size C, label smoothing parameter epsilon in [0, 1), and ground-truth index y, the target probability distribution q is:

• q(c) = (1 - epsilon) + (epsilon / C) if c == y
• q(c) = epsilon / C if c != y

Given logits z and temperature T > 0, the scaled logits are s = z / T. The model probabilities are obtained via softmax: p(c) = exp(s_c) / sum_k exp(s_k).

The cross-entropy loss for an active token is: loss = - sum_c q(c) * log p(c) = (1 - epsilon) * (-log p(y)) - (epsilon / C) * sum_c log p(c)

The total loss is the average over all valid tokens where mask == 1. If no valid tokens exist, the loss should be 0.0.

Function Signature

import numpy as np
from typing import Tuple

def masked_smoothed_cross_entropy(
    logits: np.ndarray,      # shape: (B, S, C), float32/float64
    targets: np.ndarray,     # shape: (B, S), int32/int64
    mask: np.ndarray,        # shape: (B, S), bool or {0, 1}
    epsilon: float = 0.1,
    temperature: float = 1.0
) -> Tuple[float, np.ndarray]:
    """
    Returns:
        loss: Scalar float representing the average loss over valid tokens.
        grad_logits: ndarray of shape (B, S, C), gradient w.r.t. input logits.
    """
    pass

Requirements and Edge Cases

• Numerical Stability: Use the log-sum-exp trick to compute log p(c) directly without underflowing or overflowing.
• Input Validation: Check that logits.ndim == 3, targets.ndim == 2, mask.ndim == 2, matching batch/sequence dimensions, valid target ranges 0 <= targets < C, and all inputs are finite. Raise ValueError if validation fails.
• Zero Valid Tokens: If np.sum(mask) == 0, return 0.0 for loss and an array of zeros with shape (B, S, C) for the gradient.
• Gradient: Compute the exact gradient d(loss) / d(logits). Masked token positions must have zero gradient.