Full Transformer Architecture and Training: From Input to Output

calendar_today May 31, 2025 schedule 7 min read

Learning Objectives: By the end of this lesson, you will be able to:

Diagram and explain the full Transformer encoder-decoder architecture

Describe the flow of data from input to output, including embeddings and final layers

Understand and implement the training objectives for sequence-to-sequence tasks

1. Welcome and Introduction (5 minutes)

Icebreaker Activity

Imagine you are building a machine translation system. How do you think the model knows when to stop generating words? What happens “under the hood” from the moment you input a sentence to the moment you get a translation?

Lesson Overview

Duration: Approximately 2.5 hours including all exercises (or 60–90 minutes when selected activities are assigned)
Prerequisites: Understanding of Transformer blocks, attention, and positional encoding
Materials Needed: Python (Jupyter Notebook or IDE), PyTorch or TensorFlow, pen and paper for diagramming

2. Learning Objectives and Session Outline (5 minutes)

What We’ll Cover Today:

The overall encoder-decoder architecture of the Transformer
The flow of data: input embeddings, encoder stack, decoder stack, output layer
Training objectives: loss functions and optimization for sequence-to-sequence tasks
Hands-on practice and exercises
Common pitfalls and how to avoid them

Success Criteria:

By the end of this session, you should be able to sketch and explain the full Transformer architecture and describe the training process for tasks like machine translation.

3. Foundation Knowledge (15-20 minutes)

Core Concepts

Encoder-Decoder Architecture

The Transformer consists of an encoder stack and a decoder stack. The encoder processes the input sequence and produces a set of representations. The decoder generates the output sequence, attending to both its own previous outputs and the encoder’s representations.

Input/Output Embeddings

Both input and output tokens are converted to embeddings and have positional encodings added. The output of the decoder is passed through a linear layer and softmax to produce probabilities over the vocabulary.

Training Objectives

For tasks like machine translation, the model is trained to maximize the likelihood of the correct output sequence given the input. The most common loss function is cross-entropy loss.

Real-World Analogy: Think of the encoder as a reader who understands the source text, and the decoder as a writer who produces the translation, consulting the reader as needed.

Why This Matters: Understanding the full architecture and training process is essential for building, customizing, or troubleshooting Transformer-based models.

Key Terminology

Encoder Stack: Series of encoder blocks processing the input
Decoder Stack: Series of decoder blocks generating the output
Embedding Layer: Converts tokens to vectors
Linear + Softmax Layer: Produces output probabilities
Cross-Entropy Loss: Measures the difference between predicted and true output sequences

4. Detailed Explanation and Examples (25-30 minutes)

Understanding the Full Transformer Architecture

The input sequence is embedded and passed through the encoder stack. The decoder stack generates the output sequence one token at a time, using both its own previous outputs and the encoder’s representations. The final output is produced by a linear layer followed by softmax.

Practical Applications

Machine translation
Text summarization
Sequence-to-sequence tasks (e.g., question answering)

Common Misconceptions

Misconception: The decoder generates the entire output at once. Reality: The decoder generates one token at a time, using previous outputs as context.
Misconception: The encoder and decoder stacks are trained separately. Reality: The entire model is trained end-to-end.

5. Live Coding Demonstration (20-25 minutes)

Example 1: Full Transformer Skeleton (PyTorch)

import torch
import torch.nn as nn

class TransformerModel(nn.Module):
    def __init__(self, vocab_size, embed_dim, num_heads, ff_dim, num_layers):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim)
        # This is a learned position embedding (not a fixed sinusoidal encoding).
        self.pos_encoding = nn.Parameter(torch.zeros(1, 100, embed_dim))
        encoder_layer = nn.TransformerEncoderLayer(embed_dim, num_heads, ff_dim, batch_first=True)
        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers)
        decoder_layer = nn.TransformerDecoderLayer(embed_dim, num_heads, ff_dim, batch_first=True)
        self.decoder = nn.TransformerDecoder(decoder_layer, num_layers)
        self.fc_out = nn.Linear(embed_dim, vocab_size)
    def forward(self, src, tgt):
        src_emb = self.embedding(src) + self.pos_encoding[:, :src.size(1), :]
        tgt_emb = self.embedding(tgt) + self.pos_encoding[:, :tgt.size(1), :]
        memory = self.encoder(src_emb)
        # True entries prevent each target position from seeing future tokens.
        tgt_mask = torch.triu(
            torch.ones(tgt.size(1), tgt.size(1), device=tgt.device, dtype=torch.bool),
            diagonal=1,
        )
        out = self.decoder(tgt_emb, memory, tgt_mask=tgt_mask)
        logits = self.fc_out(out)
        return logits

Explanation:

Embedding and positional encoding for input and output
Encoder and decoder stacks
Final linear layer for output probabilities

Example 2: Training Loop Skeleton

import torch.optim as optim

model = TransformerModel(vocab_size=10000, embed_dim=512, num_heads=8, ff_dim=2048, num_layers=6)
optimizer = optim.Adam(model.parameters(), lr=1e-4)
criterion = nn.CrossEntropyLoss()

for epoch in range(num_epochs):
    for src, target in dataloader:
        # Teacher forcing: input excludes EOS; labels exclude BOS.
        tgt_in, tgt_out = target[:, :-1], target[:, 1:]
        logits = model(src, tgt_in)
        loss = criterion(logits.view(-1, logits.size(-1)), tgt_out.view(-1))
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

Key Points:

The model is trained end-to-end
Cross-entropy loss is used for sequence prediction

Example 3: Visualizing the Data Flow

# Pseudocode for data flow
input_seq -> [Embedding + Positional Encoding] -> [Encoder Stack] -> encoder_outputs
start_token + previous_outputs -> [Embedding + Positional Encoding] -> [Decoder Stack (uses encoder_outputs)] -> [Linear + Softmax] -> output_seq

Debugging Demonstration

Show common errors and how to fix them:

# Common mistake: mismatched input/output shapes
try:
    logits = model(src, tgt_in)
    loss = criterion(logits, tgt_out)  # Wrong shape
except Exception as e:
    print(e)

# Corrected version
loss = criterion(logits.view(-1, logits.size(-1)), tgt_out.view(-1))

6. Hands-On Practice (30-40 minutes)

Exercise 1: Basic Practice (10 minutes)

Objective: Implement the forward pass of a Transformer model.

Instructions:

Define a TransformerModel class as above
Pass random input and target sequences through it
Print the output shape

Starter Code:

import torch
# Your code here

Solution:

model = TransformerModel(vocab_size=10000, embed_dim=512, num_heads=8, ff_dim=2048, num_layers=6)
src = torch.randint(0, 10000, (8, 20))
tgt = torch.randint(0, 10000, (8, 20))
logits = model(src, tgt)
print(logits.shape)  # (8, 20, 10000)

Exercise 2: Intermediate Challenge (15 minutes)

Objective: Implement a simple training loop for sequence-to-sequence prediction.

Requirements:

Use cross-entropy loss
Update model parameters with an optimizer
Print loss at each step

Hints:

Use nn.CrossEntropyLoss
Flatten logits and targets for loss calculation

Exercise 3: Creative Application (15 minutes)

Objective: Diagram the full data flow through the Transformer for a translation task.

Challenge:

Draw or code a diagram showing how data moves from input to output
Label each sub-layer and connection

Extension Ideas:

Add comments or annotations explaining each step
Try implementing a mini-Transformer with fewer layers

7. Common Pitfalls and Troubleshooting (10 minutes)

Frequent Mistakes

Problem: Mismatched input/output shapes Cause: Incorrect tensor dimensions Solution: Check and match shapes at each step
Problem: Forgetting to add positional encoding Cause: Model loses sense of order Solution: Always add positional encoding before feeding to the model

Debugging Tips

Print tensor shapes at each layer
Visualize data flow with diagrams
Use small test cases to verify implementation

8. Best Practices and Professional Tips (5 minutes)

Industry Standards

Use residual connections and layer normalization after each sub-layer
Always mask future tokens in the decoder during training
Document model architecture and data flow

Next Steps in Your Learning Journey

Explore pre-training and fine-tuning strategies
Learn about BERT and GPT architectures
Try implementing a full translation pipeline

9. Summary and Review (5 minutes)

Key Takeaways

The full Transformer architecture processes input and output sequences through encoder and decoder stacks
Training is performed end-to-end with cross-entropy loss
Understanding data flow is crucial for building and debugging models

Quick Review Questions

What is the purpose of the encoder and decoder stacks?
How is the output sequence generated during training?
What loss function is commonly used for sequence-to-sequence tasks?

For Instructors

Teaching Tips

Use diagrams to illustrate the full architecture and data flow
Encourage students to step through the training process
Discuss real-world applications (translation, summarization)

Assessment Ideas

Quiz on full architecture and training process
Coding assignment: implement the forward pass and training loop
Peer review: students explain their diagrams and code

Adaptation Notes

For beginners: focus on basic implementation and data flow
For advanced students: experiment with different architectures or training objectives
For online delivery: use Colab or Jupyter for interactive coding and diagramming