Kahibaro
Discord Login Register

10.4 Teacher Forcing and Sequence Losses

What Teacher Forcing Is and When You Need It

In sequence to sequence models, you predict an output sequence one token at a time. At step $t$, the model typically consumes some representation of the input and a version of the previously generated output tokens, then produces a distribution over the next token. The key practical question during training is what to feed in as the “previous token”. Teacher forcing answers this by feeding the ground truth previous token from the dataset, not the model’s own previous prediction.

With teacher forcing, the model learns the next token conditional on the correct history, which makes optimization easier and faster early on. Without teacher forcing, the model conditions on its own mistakes, which can cause errors to compound and training to become unstable, especially for longer sequences.

Important rule: During training with teacher forcing, the input token at step $t$ is the ground truth token from step $t-1$. During inference, you do not have ground truth, so you must feed the model’s own predicted token from step $t-1$.

This mismatch between training and inference is often called exposure bias. You do not need to solve it fully to get a working system, but you should understand that training loss can look good while generated sequences at inference time are still poor if mistakes snowball.

The Shape of Sequence Outputs and Targets

Most sequence decoders output logits for each time step. If your vocabulary size is $V$ and your sequence length is $T$, then for a batch size $B$ your model commonly produces logits shaped $[B, T, V]$. Your target tokens are integer class indices shaped $[B, T]$, where each entry is in $[0, V-1]$.

In many implementations you will also have padding, so not every example uses all $T$ time steps. That is handled through masking or an ignore index in the loss.

Next Token Prediction and the Shift Operation

Training a decoder is usually framed as next token prediction. That means you feed a prefix and ask the model to predict the next token at each position.

Concretely, if a target sequence (including special tokens) is

$y = [\langle bos \rangle, y_1, y_2, \dots, y_{T-1}, \langle eos \rangle]$

then the input to the decoder is often the same sequence shifted right

$y_{\text{in}} = [\langle bos \rangle, y_1, y_2, \dots, y_{T-1}]$

and the training targets are the sequence shifted left

$y_{\text{out}} = [y_1, y_2, \dots, y_{T-1}, \langle eos \rangle]$

The model predicts distributions for each position in $y_{\text{out}}$ given the corresponding prefix in $y_{\text{in}}$.

Important rule: The logits at time step $t$ should be compared against the target token at time step $t$, after you have done the right shift for inputs and left shift for targets.

Sequence Cross Entropy Loss

For token classification, the standard loss is cross entropy applied at every time step and then aggregated. If $z_{b,t} \in \mathbb{R}^V$ are logits for batch item $b$ at time $t$, and $y_{b,t}$ is the true token index, then the per token loss is

$$
\ell_{b,t} = -\log \frac{\exp(z_{b,t,y_{b,t}})}{\sum_{k=1}^{V} \exp(z_{b,t,k})}
$$

The sequence loss is typically the mean over all non padded tokens:

$$
\mathcal{L} = \frac{1}{N_{\text{tokens}}} \sum_{b=1}^{B} \sum_{t=1}^{T} m_{b,t} \, \ell_{b,t}
$$

where $m_{b,t} \in \{0,1\}$ is a mask that is 1 for real tokens and 0 for padding, and $N_{\text{tokens}} = \sum_{b,t} m_{b,t}$.

In PyTorch, this is commonly implemented with torch.nn.CrossEntropyLoss, which expects logits of shape $[N, V]$ and targets of shape $[N]`, so you usually flatten the batch and time dimensions.

Important rule: CrossEntropyLoss expects raw logits, not softmax probabilities. Do not apply softmax before the loss.

A typical pattern looks like this in code, assuming logits is [B, T, V] and targets is [B, T]:

python
import torch
import torch.nn as nn
criterion = nn.CrossEntropyLoss(ignore_index=pad_id)
B, T, V = logits.shape
loss = criterion(logits.reshape(B*T, V), targets.reshape(B*T))

If you do not have padding you can omit ignore_index. If you do have padding, using ignore_index=pad_id is often the simplest option, as long as your padding token id is consistent.

Masking, Padding, and Ignore Index

In real datasets, sequences are padded to a common length within a batch. Those padded positions should not contribute to the loss. There are two standard approaches.

One approach uses ignore_index in CrossEntropyLoss. If your pad token id is pad_id, then any target equal to pad_id is ignored automatically, which effectively implements the mask.

The other approach is to compute per token losses with reduction="none" and then apply an explicit mask, which is useful when you need custom weighting:

python
criterion = nn.CrossEntropyLoss(reduction="none")
per_token = criterion(logits.reshape(B*T, V), targets.reshape(B*T))  # [B*T]
per_token = per_token.view(B, T)
mask = (targets != pad_id).float()
loss = (per_token * mask).sum() / mask.sum().clamp_min(1.0)

Important rule: If you average losses over time steps, average over non padded tokens only, otherwise longer padding will artificially reduce or distort the loss.

Teacher Forcing Ratio and Scheduled Sampling (Conceptually)

Sometimes you use teacher forcing only some of the time, controlled by a teacher forcing ratio. With probability $p$, you feed the ground truth previous token. With probability $1-p$, you feed the model’s previous prediction. Reducing $p$ over training is one way to reduce exposure bias. This family of ideas is often called scheduled sampling.

For absolute beginners, it is enough to know that full teacher forcing, meaning $p = 1$, is the common baseline because it is simple and trains reliably. More advanced strategies can improve generation quality but add complexity and can make training noisier.

Important rule: Even if you use a teacher forcing ratio less than 1 during training, evaluation and inference should still use the same autoregressive procedure you will use in production.

Putting It Together in a Minimal Training Step

A typical training step for a decoder model with teacher forcing uses shifted sequences and a token loss.

Assume you have a batch of tokenized target sequences called y of shape [B, T] that already includes a beginning of sequence token at position 0, and uses pad_id for padding. Then:

python
y_in = y[:, :-1]     # [B, T-1]
y_out = y[:, 1:]     # [B, T-1]
logits = model(y_in) # [B, T-1, V], model uses y_in as teacher forced inputs
loss = nn.CrossEntropyLoss(ignore_index=pad_id)(
    logits.reshape(-1, logits.size(-1)),
    y_out.reshape(-1)
)

The only job of the shifting is to align the model’s prediction at each step with the correct next token label. This is the core mechanics behind teacher forcing and sequence losses in most token based sequence models.

Common Alignment Mistakes to Watch For

One common mistake is comparing logits produced for y_in against y_in itself, which trains the model to copy the current token rather than predict the next token. Another common mistake is forgetting to drop the last time step in the inputs or forgetting to drop the first time step in the targets, which causes an off by one error.

Important rule: Inputs are shifted right, targets are shifted left. If you get this wrong, the model may still train, but it will learn the wrong task.

A final practical issue is confusing token ids with one hot vectors. For CrossEntropyLoss, your targets should be integer token ids, not one hot encodings.

Views: 67

Comments

Please login to add a comment.

Don't have an account? Register now!