Transformer
teacher-forcing
machine learning
training techniques
NLP

How is teacher-forcing implemented for the Transformer training?

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

In Transformer training, teacher forcing is implemented by feeding the decoder the ground-truth target sequence shifted to the right, rather than feeding it its own previous predictions. That allows the model to learn all target positions in parallel while still preserving the rule that position t may only attend to earlier target tokens.

The Core Idea: Shift the Target Sequence

Suppose the true target sentence is:

text
["I", "like", "coffee", "<eos>"]

During training, the decoder input becomes:

text
["<bos>", "I", "like", "coffee"]

The training labels remain:

text
["I", "like", "coffee", "<eos>"]

That is teacher forcing in a Transformer. At each position, the model sees the correct previous token, not the token it predicted on the prior step.

Use a Causal Mask So the Decoder Cannot Peek Ahead

Because Transformers process the whole sequence in parallel, teacher forcing alone is not enough. The decoder also needs a causal mask so position t cannot attend to future target tokens.

A small PyTorch example shows the shape of the training setup:

python
1import torch
2import torch.nn as nn
3
4vocab_size = 100
5d_model = 32
6
7embedding = nn.Embedding(vocab_size, d_model)
8transformer = nn.Transformer(
9    d_model=d_model,
10    nhead=4,
11    num_encoder_layers=2,
12    num_decoder_layers=2,
13    batch_first=True,
14)
15output_layer = nn.Linear(d_model, vocab_size)
16
17src = torch.tensor([[5, 6, 7, 0]])
18target = torch.tensor([[1, 20, 21, 2]])  # 1=<bos>, 2=<eos>
19
20decoder_input = target[:, :-1]
21labels = target[:, 1:]
22
23tgt_mask = nn.Transformer.generate_square_subsequent_mask(decoder_input.size(1))
24
25src_emb = embedding(src)
26tgt_emb = embedding(decoder_input)
27
28hidden = transformer(src_emb, tgt_emb, tgt_mask=tgt_mask)
29logits = output_layer(hidden)
30
31loss = nn.CrossEntropyLoss()(logits.reshape(-1, vocab_size), labels.reshape(-1))
32print(loss.item())

Notice the two key details:

  • 'decoder_input is the ground-truth sequence shifted right'
  • 'labels are the next-token targets'

Why This Is Still Teacher Forcing Even Though the Decoder Runs in Parallel

In RNNs, teacher forcing is easy to visualize step by step because the model processes one time step at a time. In Transformers, all decoder positions are computed in parallel, but the logic is equivalent:

  • position 1 gets the true first previous token
  • position 2 gets the true second previous token
  • and so on

The causal mask makes sure the model does not use future target tokens directly, even though the entire tensor is present in memory at once.

Training vs. Inference

Teacher forcing is only for training. At inference time, the model must use its own generated tokens because the ground-truth target sequence is not available.

That means inference looks more like:

  1. Start with bos
  2. Predict the next token
  3. Append the prediction
  4. Feed the growing sequence back into the decoder
  5. Stop at eos

This mismatch between training and inference is one reason people talk about exposure bias in sequence models.

Common Pitfalls

The biggest mistake is forgetting to shift the target. If you feed the unshifted target sequence as both decoder input and label, the model is effectively being asked to predict tokens it can already see.

Another issue is forgetting the causal mask. In that case, the decoder can attend to future positions and training loss looks artificially good because the task has been leaked.

Developers also sometimes confuse teacher forcing with scheduled sampling. Standard Transformer training usually uses full teacher forcing with shifted targets; scheduled sampling is a separate experimental strategy.

Finally, ensure padding tokens are masked properly in both attention and loss computation. Otherwise the model wastes capacity learning from meaningless positions.

Summary

  • Teacher forcing in Transformers is implemented by feeding the decoder the ground-truth target sequence shifted right.
  • The labels are the original target sequence shifted left by one position.
  • A causal mask prevents each decoder position from seeing future target tokens.
  • Training uses teacher forcing; inference uses the model's own generated outputs.
  • The most common bugs are forgetting the shift, forgetting the mask, or mishandling padding.

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.