Safely Normalized Cross-Entropy Loss for Corrupted Non-Padding Tokens

Welcome back. Last lesson, your denoiser produced vocabulary logits at every sequence position—even visible and padded ones. The next question is which of those predictions should count as mistakes.

In this 15–20 minute lesson, you’ll compute cross-entropy loss only where a valid character was hidden. You’ll also handle a small-batch edge case: sometimes no characters are masked at all.

What should the denoiser be graded on?

Let $x_0$ be the clean token IDs, and let corruption_mask mark the positions replaced by [MASK] in the input. At a selected position, the target is the original token in $x_0$, not the [MASK] token now present in the corrupted input.

For example, if the clean characters are c a t [PAD] and the input is c [MASK] t [PAD], the denoiser should be graded on predicting a at the second position. Predicting c or t at their already-visible positions earns no credit; neither does predicting anything at [PAD].

The positions that contribute to loss are therefore those where both conditions hold: the position was corrupted, and its clean token is not padding. The padding attention mask from the encoder lesson controls which positions can provide context. This loss mask controls which predictions are graded. They have different jobs.

For a short PyTorch explanation, watch this portion of PyTorch Tutorial 11 - Softmax and Cross Entropy by Patrick Loeber:

Normalize by predictions that actually count

If two valid positions were corrupted, the loss should be the average of two token-prediction penalties—not an average over every position in the padded batch. In notation, if $S$ is the set of corrupted, non-padding positions, our loss is

Here $z_{b,\ell }$ is the vector of vocabulary logits at one position. The denominator is the number of predictions being graded. The $\max (1,|S|)$ handles an empty $S$: its sum is zero, so the loss is zero rather than an undefined average. This can occur when random masking selects no tokens in a small batch.

PyTorch’s ignore_index gives us a convenient way to exclude positions. We’ll place -100 in the target tensor wherever loss should not be computed; it is not a vocabulary token.

Implement and check it in Colab

The vocabulary dimension of last lesson’s logits is last, with shape [batch, length, vocab_size]. For the full-batch form of F.cross_entropy, the class dimension must be second, so transpose(1, 2) changes the shape to [batch, vocab_size, length].

import torch
import torch.nn.functional as F

def denoising_loss(logits, x0, corruption_mask, pad_id):
    # x0 and corruption_mask: [batch, length]
    # logits: [batch, length, vocab_size]
    selected = corruption_mask.bool() & x0.ne(pad_id)

    targets = x0.clone()
    targets = targets.masked_fill(~selected, -100)

    loss_sum = F.cross_entropy(
        logits.transpose(1, 2),  # [batch, vocab_size, length]
        targets,                 # [batch, length], integer token IDs
        ignore_index=-100,
        reduction="sum",
    )
    return loss_sum / selected.sum().clamp_min(1)

Notice that targets starts from clean IDs. Setting unselected targets to -100 excludes visible tokens as well as padding. Merely setting ignore_index=pad_id would still grade the visible, non-padding positions.

Here is a self-contained check; it does not require trained model weights:

pad_id = 0
x0 = torch.tensor([[2, 3, 4, 0],
                   [4, 2, 0, 0]], dtype=torch.long)
corruption_mask = torch.tensor([[ True, False, False, False],
                                [False,  True, False, False]])

# Eight equally scored vocabulary tokens at each position.
logits = torch.zeros(2, 4, 8, requires_grad=True)
loss = denoising_loss(logits, x0, corruption_mask, pad_id)

print("loss:", loss.item())
assert torch.allclose(loss, torch.log(torch.tensor(8.0)))

loss.backward()
selected = corruption_mask & x0.ne(pad_id)
assert torch.all(logits.grad[~selected] == 0)

# Even if this batch happens to mask nothing, the loss remains usable.
logits.grad.zero_()
empty_mask = torch.zeros_like(corruption_mask)
empty_loss = denoising_loss(logits, x0, empty_mask, pad_id)
print("empty-mask loss:", empty_loss.item())
empty_loss.backward()
assert torch.count_nonzero(logits.grad) == 0

Each selected position has eight equally scored candidates, so its cross-entropy is $\log 8$. Averaging the two selected positions gives the same value. The gradient check confirms that visible and padded output positions receive no direct loss.

When you replace these test logits with your denoiser’s output, run a fresh forward pass without torch.no_grad(). Last lesson used torch.no_grad() for a shape check; the upcoming training step needs gradients.

What to carry forward

A denoising loss compares raw logits with clean token IDs only at corrupted, non-padding positions. Ignored targets remove the other positions, while summing and dividing by a count clamped to at least one keeps an empty-mask batch finite and differentiable.

Next, you’ll put this loss into an optimizer step so the denoiser’s predictions can actually improve.