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.

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:

class TransformerTrainer:
    def __init__(self, model, train_loader, lr=1e-3, warmup_steps=20):
        self.model = model
        self.train_loader = train_loader
        self.start_lr = lr
        self.optimizer = optim.Adam(model.parameters(), lr=lr, betas=(0.9, 0.98), eps=1e-9)
        self.criterion = nn.CrossEntropyLoss(ignore_index=0)  # Ignore padding tokens
        self.warmup_steps = warmup_steps
        self.step_num = 0

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.

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

    def train_epoch(self):
        """Train for one epoch"""
        self.model.train()
        total_loss = 0
        num_batches = 0
        
        for batch in self.train_loader:
            self.step_num += 1
            self.update_lr()
            
            # Get batch data
            src = batch['src']
            tgt_input = batch['tgt']
            tgt_output = batch['tgt_output']
            
            # Create masks
            src_mask = self.model.create_padding_mask(src)
            tgt_causal_mask = self.model.create_causal_mask(tgt_input.size(1))
            tgt_padding_mask = self.model.create_padding_mask(tgt_input)
            # Combine masks: both must be True for attention to be allowed
            # Broadcasting will handle the shape differences
            tgt_mask = tgt_causal_mask & tgt_padding_mask

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:

            # Forward pass with teacher forcing
            self.optimizer.zero_grad()
            output = self.model(src, tgt_input, src_mask, tgt_mask)
            
            # Compute loss
            loss = self.criterion(output.reshape(-1, output.size(-1)), tgt_output.reshape(-1))
            
            # Backward pass
            loss.backward()
            self.optimizer.step()
            
            total_loss += loss.item()
            num_batches += 1
            
            if num_batches % 10 == 0:
                avg_loss = total_loss / num_batches
                lr = self.optimizer.param_groups[0]['lr']
                print(f"Step {self.step_num}, Loss: {avg_loss:.4f}, LR: {lr:.6f}")
        
        return total_loss / num_batches

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:

def train_transformer():
    """Train Transformer model on synthetic data"""
    print("Training Transformer Model...")
    
    # Create data
    src_sentences, tgt_sentences = create_synthetic_data(num_samples=200)
    
    # Build vocabularies
    src_vocab = Vocabulary()
    tgt_vocab = Vocabulary()
    src_vocab.build_vocab(src_sentences)
    tgt_vocab.build_vocab(tgt_sentences)
    
    print(f"Source vocab size: {src_vocab.size}")
    print(f"Target vocab size: {tgt_vocab.size}")
    
    # Create dataset and dataloader
    dataset = TranslationDataset(src_sentences, tgt_sentences, src_vocab, tgt_vocab, max_len=15)
    train_loader = DataLoader(dataset, batch_size=8, shuffle=True, collate_fn=collate_fn)
    
    # Create model
    model = Transformer(
        src_vocab_size=src_vocab.size,
        tgt_vocab_size=tgt_vocab.size,
        d_model=64,
        num_heads=4,
        num_encoder_layers=2,
        num_decoder_layers=2,
        d_ff=256,
        dropout=0.1
    )
    
    print(f"Model parameters: {sum(p.numel() for p in model.parameters()):,}")

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:

    # Create trainer and train
    trainer = TransformerTrainer(model, train_loader, lr=1e-3, warmup_steps=25)
    
    # Training loop
    num_epochs = 3
    for epoch in range(num_epochs):
        print(f"\nEpoch {epoch + 1}/{num_epochs}")
        avg_loss = trainer.train_epoch()
        print(f"Average loss: {avg_loss:.4f}")
    
    return model, src_vocab, tgt_vocab

def main():
    model, src_vocab, tgt_vocab = train_transformer()
    print("Training completed successfully!")

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

Training Transformer Model...
Source vocab size: 22
Target vocab size: 22
Model parameters: 237,718

Epoch 1/3
Step 10, Loss: 3.9631, LR: 0.000011
Step 20, Loss: 3.8995, LR: 0.000021
Average loss: 3.8170

Epoch 2/3
Step 35, Loss: 3.2435, LR: 0.000021
Step 45, Loss: 3.1208, LR: 0.000018
Average loss: 3.0270

Epoch 3/3
Step 60, Loss: 2.7883, LR: 0.000016
Step 70, Loss: 2.6777, LR: 0.000015
Average loss: 2.6533
Training completed successfully!

The training output reveals several important patterns that demonstrate successful learning. The loss decreases consistently from 3.8170 in the first epoch to 2.6533 in the third epoch, indicating the model is learning the word reversal task. The learning rate schedule is working correctly, starting very low (0.000011) during warmup, reaching its peak around step 20-25, then gradually decreasing. This step-by-step monitoring shows that individual batch losses are decreasing within each epoch, and the learning rate adjustments follow the expected warmup-then-decay pattern.

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!