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.
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:
During training, the decoder input becomes:
The training labels remain:
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:
Notice the two key details:
- '
decoder_inputis the ground-truth sequence shifted right' - '
labelsare 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:
- Start with
bos - Predict the next token
- Append the prediction
- Feed the growing sequence back into the decoder
- 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
- How is the categorical_crossentropy implemented in keras?
- How is the input tensor for TensorFlow's tf.nn.dynamic_rnn operator structured?
- How is the smooth dice loss differentiable?
- how is total loss calculated over multiple classes in Keras?
- How to accurately classify text with a lot of potential values using scikit?
- How to add new embeddings for unknown words in Tensorflow training pre-set for testing
- How is tf.summary.tensor_summary meant to be used?
- How is the Keras Conv1D input specified? I seem to be lacking a dimension
.png&w=3840&q=75)
Tackling System Design Interview Problems
A short course that equips you with the skills to approach system design interviews methodically.
Start the free courseTrack 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.