Transformer Training Essentials

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: $P(y|x) = \prod_{t=1}^{T} P(y_t | y_{<t}, x)$ where $x$ is the source sequence and $y_{<t}$ represents all tokens before position $t$. This factorization allows the model to learn complex dependencies while maintaining computational tractability. During training, we use cross-entropy loss 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:

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.

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:

where $d_{model}$ is the model dimension. This creates a learning rate that increases linearly during warmup, then decreases with the inverse square root of the step number. This schedule helps the model converge more reliably while maintaining training stability throughout the process.

Let's implement this learning rate scheduling logic:

    def get_lr_scale(self):
        """Learning rate schedule with warmup"""
        d_model = self.model.d_model
        step_num = self.step_num + 1
        return (d_model ** -0.5) * min(step_num ** -0.5, step_num * (self.warmup_steps ** -1.5))
    
    def update_lr(self):
        """Update learning rate based on schedule"""
        lr_scale = self.get_lr_scale()
        for param_group in self.optimizer.param_groups:
            param_group['lr'] = self.start_lr * lr_scale

The get_lr_scale method implements the mathematical formula we discussed earlier, computing the scaling factor based on the current step and warmup duration. During warmup, the learning rate increases linearly, then transitions to inverse square root decay. The update_lr method applies this scaling to the optimizer's parameter groups, ensuring the learning rate updates dynamically throughout training.

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!