17. Token-Level Cross-Entropy With Label Smoothing and Padding Mask

Problem

A language model's training loss is token-level cross-entropy, but two production details matter beyond the textbook formula. First, batches pad shorter sequences to a common length, and those pad positions must be excluded from both the loss sum and the averaging denominator, or the loss (and its gradient) gets diluted by however much padding happened to be in the batch. Second, label smoothing softens the one-hot target distribution (1-smoothing probability on the true token, smoothing spread over the rest of the vocabulary) to keep the model from becoming overconfident, which is standard practice in most modern LLM pretraining and fine-tuning recipes.

Given logits (shape (batch, seq_len, vocab_size)) and targets (shape (batch, seq_len)), compute the mean label-smoothed loss over only the non-padded positions.

Source: src/17_label_smoothed_cross_entropy.py

def label_smoothed_nll_loss(logits, targets, pad_id, smoothing=0.1) -> torch.Tensor

>>> import torch
>>> logits = torch.tensor([[[2.0, 0.1, 0.1], [0.1, 2.0, 0.1]]])
>>> targets = torch.tensor([[0, 1]])
>>> label_smoothed_nll_loss(logits, targets, pad_id=99, smoothing=0.0)
tensor(0.2617)

Step-by-Step Approach

1. Compute log_probs = F.log_softmax(logits, dim=-1) over the vocabulary dimension — using the fused log-softmax rather than torch.log(torch.softmax(...)) avoids a separate numerically unstable log of a (possibly tiny) probability.
2. Build the smoothed target distribution inside a torch.no_grad() block: start with a tensor filled with smoothing / (vocab_size - 1) everywhere, then scatter_ 1 - smoothing onto each position's true target index.
3. Compute the per-token negative log-likelihood against this soft target as -(true_dist * log_probs).sum(-1) — this is the general cross-entropy-against-a-distribution formula, which reduces to standard NLL when true_dist is one-hot (i.e. when smoothing=0).
4. Build a padding mask with (targets != pad_id).float() and multiply it elementwise into the per-token loss, zeroing out every padded position's contribution.
5. Sum the masked per-token losses and divide by the number of non-padded tokens (mask.sum()), not by batch * seq_len — dividing by the wrong denominator is the single most common bug in padded-sequence loss code.
6. Clamp the denominator to at least 1 with .clamp(min=1.0) so an all-padding batch returns 0.0 instead of NaN from a division by zero.

The key insight is that masking must happen at the per-token loss level, before both the sum and the count in the final division — masking only the numerator (summing masked losses but dividing by batch * seq_len) silently produces a loss that shrinks as padding grows, which looks like training is improving as sequences get shorter relative to a fixed max length, purely as an artifact of the wrong denominator.

Reference solution

import torch
import torch.nn.functional as F


def label_smoothed_nll_loss(logits, targets, pad_id, smoothing=0.1):
    vocab_size = logits.size(-1)
    log_probs = F.log_softmax(logits, dim=-1)  # (batch, seq_len, vocab_size)

    # smoothed target distribution: 1-smoothing on the true token, the
    # remaining smoothing mass spread uniformly over the other vocab_size-1
    # tokens. Built with no_grad since it's a fixed target, not a learned
    # quantity.
    with torch.no_grad():
        true_dist = torch.full_like(log_probs, smoothing / (vocab_size - 1))
        true_dist.scatter_(-1, targets.unsqueeze(-1), 1.0 - smoothing)

    # cross-entropy against a soft target is -sum(true_dist * log_probs),
    # not gather+nll, since the target is no longer one-hot
    loss_per_token = -(true_dist * log_probs).sum(-1)  # (batch, seq_len)

    mask = (targets != pad_id).float()  # (batch, seq_len), 0 at pad positions
    loss_per_token = loss_per_token * mask
    # clamp the denominator so an all-padding batch returns 0.0 instead of NaN
    loss = loss_per_token.sum() / mask.sum().clamp(min=1.0)
    return loss

Key Functions & Tricks

• F.log_softmax(logits, dim=-1) — the fused, numerically stable log-probability computation over the vocabulary axis.
• Tensor.scatter_(-1, index, value) — writes 1-smoothing at each position's true-token index without a Python loop or one-hot matmul.
• torch.no_grad() — marks the smoothed target as a fixed constant, since it should never receive gradients itself.
• (targets != pad_id).float() — the padding mask, built directly from the target ids rather than requiring a separately-passed attention mask.
• mask.sum().clamp(min=1.0) — the division-by-zero guard for an entirely padded batch.
• Cross-checking against F.cross_entropy with smoothing=0.0 and no padding — the fastest way to confirm the from-scratch formula reduces correctly to the standard case.

How to Recognize This Pattern