Welcome back to Bringing Transformers to Life: Training & Inference! You've made excellent progress in this course. In your first lesson, you assembled a complete Transformer architecture, integrating all the components we've built throughout this learning journey. In the second lesson, you created a robust data preparation pipeline that transforms raw text into training-ready tensors, complete with vocabularies, special tokens, and dynamic batching.

Now we are at a pivotal point, as in today's lesson we'll be discussing training the Transformer. This is where everything comes together as we implement the actual training process: you'll learn how Transformers learn through autoregressive modeling, where the model predicts the next token given all previous tokens. We'll explore teacher forcing, a key training technique that accelerates learning, and implement sophisticated optimization strategies, including learning rate scheduling with warmup. By the end of this lesson, you'll have a complete training pipeline that can effectively train your Transformer model on sequence-to-sequence tasks.

The autoregressive training objective forms the foundation of how Transformers learn to generate sequences. Unlike traditional machine learning, where we might predict a single output, autoregressive models learn to predict each token in a sequence conditioned on all previous tokens. This creates a natural decomposition of the sequence generation probability.

Mathematically, for a target sequence $y = (y_1, y_2, ..., y_T)$, the autoregressive objective maximizes the likelihood: where is the source sequence and represents all tokens before position . This factorization allows the model to learn complex dependencies while maintaining computational tractability. During training, we use to measure how well the model's predicted probability distribution matches the true next token at each position.

Teacher forcing is the standard way to train autoregressive Transformers (both encoder-decoder and decoder-only models).
During training we feed the ground-truth tokens into the decoder (or into the masked language-model input), while asking the model to predict the next token for every position. Because the model always sees the correct previous context, gradients are better behaved and convergence is much faster than if it had to consume its own, still-noisy predictions.

The downside is the resulting train–inference mismatch, often called exposure bias:

• Training: context consists of perfect tokens from the data set.
• Inference: context consists of the model's own predictions, which may contain mistakes that propagate.

Although exposure bias is an unwanted side-effect, in practice teacher forcing is still preferred because (1) it makes optimization tractable, (2) it yields state-of-the-art performance when combined with techniques such as scheduled sampling, label smoothing, or beam search, and (3) large, diverse data sets help the model learn to recover from occasional errors.

Let's begin implementing our TransformerTrainer class, which encapsulates all the training logic, including the sophisticated learning rate scheduling:

This initialization sets up the essential training components. We use Adam optimizer with specific beta values (0.9, 0.98) that work well for Transformers, following established best practices. The CrossEntropyLoss with ignore_index=0 ensures that padding tokens don't contribute to the loss calculation, which is crucial for variable-length sequences. The warmup_steps parameter controls how long the learning rate increases before beginning to decay.

Effective learning rate scheduling is crucial for stable Transformer training. The original Transformer paper introduced a specific warmup schedule that gradually increases the learning rate during early training steps, then decreases it proportionally. This approach prevents the large parameter updates that can destabilize training in the early stages.

The warmup schedule follows the formula:

The core training logic handles teacher forcing and proper masking for both padding and causal attention:

This section demonstrates the dual masking strategy essential for proper Transformer training. The source mask prevents attention to padding tokens, while the target mask combines causal masking (preventing future token access) with padding masking. The & operator ensures both conditions must be satisfied for attention to occur, maintaining the autoregressive property while handling variable-length sequences.

The optimization phase of the training loop implements the complete forward-backward pass with loss computation:

This loop showcases teacher forcing in action: tgt_input contains the ground truth tokens (with <SOS> prefix), while tgt_output contains the targets (with <EOS> suffix). The model learns to predict each token in tgt_output given the corresponding prefix in tgt_input. The loss reshaping flattens the sequence dimension, treating each position as an independent classification problem across the vocabulary.

Now let's examine the complete training pipeline that brings everything together:

This pipeline demonstrates the complete integration of all components we've built. We create synthetic data, build vocabularies, initialize the dataset and dataloader using our custom collate_fn function, and instantiate a Transformer model with appropriate hyperparameters. The model size (237,718 parameters) is reasonable for our synthetic task while being large enough to demonstrate meaningful learning dynamics.

The final training execution brings everything together:

When we execute this training pipeline, we observe the following output:

You have successfully implemented a complete Transformer training pipeline that demonstrates the fundamental principles of autoregressive learning. Your implementation incorporates teacher forcing for efficient training, sophisticated learning rate scheduling with warmup, proper masking for both padding and causal attention, and robust optimization strategies. The training results show clear evidence of learning, with consistent loss reduction and proper learning rate dynamics.

This comprehensive training framework provides the foundation for tackling real-world sequence-to-sequence tasks, from machine translation to text summarization. In the upcoming practice exercises, you'll have the opportunity to experiment with different hyperparameters, explore various training strategies, and gain hands-on experience with the nuances of Transformer training that will make you proficient in bringing these powerful models to life. Then, in the next and final of this course we'll be discussing inference strategies such as greedy decoding and beam search. Keep learning!