Write the Training Step for a Padded Token Classifier
You are asked to train a named-entity tagger. Each example is a sentence
of 3 to 80 token ids with one tag id (0-8) per token. The vocabulary has
30,000 entries and id 0 is reserved for padding. You have an
nn.Module that maps a (B, T) batch of token ids to (B, T, 9)
logits.
- Write the
collate_fnthat turns a list of(tokens, tags)pairs into a batch. What value do you pad the tags with, and why not 0? - Write the loss computation for one batch so that padding positions contribute nothing and the loss is a per-token mean. Why does the normalization matter when batches have different amounts of padding?
- A 7-sentence toy sample has lengths 5, 12, 3, 40, 7, 9 and 22 and
you use
batch_size=3without shuffling. How many real tokens and how many padded cells does each batch contain, and what does that suggest for throughput? - How do you compute token accuracy for the whole validation set so that it is correct?
1. Collate and the padding value
from torch.nn.utils.rnn import pad_sequence
PAD_ID, IGNORE = 0, -100
def collate(batch): # batch = [(tokens, tags), ...]
xs, ys = zip(*batch)
x = pad_sequence(xs, batch_first=True, padding_value=PAD_ID) # (B, T_max)
y = pad_sequence(ys, batch_first=True, padding_value=IGNORE) # (B, T_max)
return x, y
Tags are padded with -100, the default ignore_index of
cross_entropy, not with 0: tag 0 is a real class, so padding with 0
would train the model to predict class 0 on padding and would corrupt
the metric. Inputs are padded with the reserved id 0, ideally with
nn.Embedding(..., padding_idx=0) so that row receives no gradient.
2. Per-token loss
logits = model(x) # (B, T, 9)
loss_sum = nn.functional.cross_entropy(
logits.reshape(-1, logits.size(-1)), y.reshape(-1),
ignore_index=-100, reduction="sum")
n_tok = (y != -100).sum()
(loss_sum / n_tok).backward()
(With reduction="mean" the default also averages over only the
non-ignored targets; the explicit sum makes the count available for
aggregation.) Normalization matters because a per-sequence mean
over-weights short sentences. Example: a 3-token sentence with total
loss 3.0 and a 12-token sentence with total loss 6.0 have means 1.0 and
0.5, so the per-sequence average is 0.75, but the per-token average is
9/15 = 0.60. If you used mean over all padded cells (B * T_max)
the loss would shrink as padding grows, silently changing the
effective learning rate from batch to batch.
3. Padding waste
Batches of 3 in order: lengths (5, 12, 3), (40, 7, 9) and (22). Batch 1 has 20 real tokens in 3 \times 12 = 36 cells (16 padded); batch 2 has 56 real tokens in 3 \times 40 = 120 cells (64 padded); batch 3 has 22 tokens in 22 cells. Total: 98 real tokens in 178 cells, 55% efficiency, i.e. nearly half the compute is spent on padding. The fix is length bucketing: sort or group examples of similar length (a bucketing sampler, or sort within shuffled mega-batches of, say, 50 batches) so padding falls to a few percent while batches remain random enough to train well. Bucketing changes the batch composition, so shuffle at the bucket level every epoch.
4. Correct validation accuracy
Count tokens, not batches. Accumulate
correct += ((logits.argmax(-1) == y) & (y != -100)).sum().item() and
total += (y != -100).sum().item(), and report correct / total once
at the end, under model.eval() and torch.inference_mode(), with
drop_last=False. Averaging per-batch accuracies would weight a
3-sentence batch of long sentences the same as a 1-sentence batch of
short ones. Entity-level F1 (what NER is normally scored on) is not a
per-token average at all and must be computed from the concatenated
predicted and gold spans of the whole set. With gradient accumulation
the same logic applies to the loss: divide each micro-batch's summed
loss by the total token count of the accumulated group, not by
the number of micro-batches.
Share this question