Module 9 — Complete Transformer

Introduction

In the previous modules, you learned every major component of the Transformer:

• Input Embeddings
• Positional Encoding
• Multi-Head Attention
• Feed Forward Networks
• Transformer Encoder
• Transformer Decoder
• Masked Self-Attention
• Cross-Attention

In this module, you'll combine all these components to build the complete Transformer architecture introduced in the paper "Attention Is All You Need" (2017).

The complete Transformer consists of two main parts:

1. Encoder – Understands the input sequence.
2. Decoder – Generates the output sequence.

This architecture powers many sequence-to-sequence tasks, including:

• Machine Translation
• Text Summarization
• Question Answering
• Speech Recognition

By the end of this module, you'll understand how the entire Transformer works from input text to generated output.

1. Encoder–Decoder Architecture

The original Transformer consists of an encoder stack and a decoder stack.

Architecture

textCopy
1                Input Sentence
2                      │
3              Input Embedding
4                      │
5            Positional Encoding
6                      │
7             ┌────────────────┐
8             │ Encoder Stack  │
9             └────────────────┘
10                      │
11          Contextual Representations
12                      │
13                      ▼
14             ┌────────────────┐
15             │ Decoder Stack  │
16             └────────────────┘
17                      │
18            Linear Projection
19                      │
20                  Softmax
21                      │
22               Generated Tokens

The encoder processes the entire input sentence simultaneously, while the decoder generates one token at a time.

2. Input Embeddings

Neural networks cannot process raw words directly.

Each token is converted into a dense vector using an embedding layer.

Example

textCopy
1"I love AI"
2
3↓
4
5[15, 204, 88]
6
7↓
8
9Embedding Layer
10
11↓
12
13(3 × 512)

PyTorch Example

pythonCopy
1import torch
2import torch.nn as nn
3
4vocab_size = 30000
5embed_dim = 512
6
7embedding = nn.Embedding(
8    vocab_size,
9    embed_dim
10)
11
12tokens = torch.tensor([[10, 25, 40, 90]])
13
14vectors = embedding(tokens)
15
16print(vectors.shape)

Output

textCopy
1torch.Size([1, 4, 512])

3. Output Embeddings

The decoder also requires embeddings for previously generated target tokens.

Example

textCopy
1Target Tokens
2
3↓
4
5Embedding Layer
6
7↓
8
9Decoder Input

Usually, the input and output embedding dimensions are the same.

Example

pythonCopy
1target_tokens = torch.tensor([[1, 35, 18]])
2
3target_embeddings = embedding(target_tokens)
4
5print(target_embeddings.shape)

4. Attention Masks

Attention masks control which tokens are allowed to attend to each other.

Types

• Padding Mask
• Causal Mask

Masks are essential for correct Transformer training.

5. Padding Mask

Different sentences have different lengths.

Example

textCopy
1Sentence A
2
3I love AI
4
5Sentence B
6
7Hello

After padding

textCopy
1I love AI <PAD>
2
3Hello <PAD> <PAD> <PAD>

The model should ignore padding tokens.

Padding mask

textCopy
11 1 1 0
2
31 0 0 0

PyTorch Example

pythonCopy
1import torch
2
3tokens = torch.tensor([
4    [10,20,30,0],
5    [15,0,0,0]
6])
7
8padding_mask = tokens == 0
9
10print(padding_mask)

Output

textCopy
1tensor([
2 [False,False,False, True],
3 [False, True, True, True]
4])

6. Causal Mask

The decoder must not look at future words.

Example

textCopy
1I love Transformers

When predicting

textCopy
1love

Future token

textCopy
1Transformers

must remain hidden.

Causal mask

textCopy
1✓ ✗ ✗
2
3✓ ✓ ✗
4
5✓ ✓ ✓

PyTorch Example

pythonCopy
1sequence_length = 6
2
3mask = torch.triu(
4    torch.ones(sequence_length, sequence_length),
5    diagonal=1
6).bool()
7
8print(mask)

7. Training Pipeline

During training, the model sees both the source sentence and the correct target sentence.

Pipeline

textCopy
1Source Sentence
2        │
3Encoder
4        │
5Memory
6        │
7Target Tokens
8        │
9Decoder
10        │
11Linear Projection
12        │
13Softmax
14        │
15Cross Entropy Loss
16        │
17Backpropagation

Steps

1. Tokenize source and target text.
2. Create input and target embeddings.
3. Apply positional encoding.
4. Pass the source through the encoder.
5. Pass shifted target tokens through the decoder.
6. Compute vocabulary logits.
7. Calculate cross-entropy loss.
8. Update model parameters using backpropagation.

Simplified Training Loop

pythonCopy
1optimizer.zero_grad()
2
3logits = model(
4    source,
5    target_input
6)
7
8loss = criterion(
9    logits.reshape(-1, vocab_size),
10    target_output.reshape(-1)
11)
12
13loss.backward()
14
15optimizer.step()

8. Inference Pipeline

During inference, the correct target sentence is not available.

The decoder generates tokens one at a time.

Pipeline

textCopy
1Source Sentence
2
3↓
4
5Encoder
6
7↓
8
9Decoder
10
11↓
12
13Generate Next Token
14
15↓
16
17Append Token
18
19↓
20
21Repeat
22
23↓
24
25<EOS>

Unlike training, inference is autoregressive.

Simple Inference Loop

pythonCopy
1generated = [bos_token]
2
3for _ in range(max_length):
4
5    logits = model(source, generated)
6
7    next_token = logits.argmax(dim=-1)
8
9    generated.append(next_token.item())
10
11    if next_token.item() == eos_token:
12        break

9. Greedy Search

Greedy Search selects the token with the highest probability at every decoding step.

Example

textCopy
1Step 1
2
3A : 0.70
4B : 0.20
5C : 0.10
6
7↓
8
9Choose A

Advantages

• Fast
• Simple
• Low memory usage

Disadvantages

• May miss better overall sequences
• Can produce repetitive outputs

Greedy Search Example

pythonCopy
1probabilities = torch.tensor(
2    [0.2, 0.6, 0.1, 0.1]
3)
4
5token = torch.argmax(probabilities)
6
7print(token)

Output

textCopy
1tensor(1)

10. Beam Search

Beam Search keeps multiple candidate sequences instead of only one.

Example

textCopy
1Beam Width = 3
2
3Sentence 1
4
5Sentence 2
6
7Sentence 3
8
9↓
10
11Keep Top 3
12
13↓
14
15Continue Expansion

Advantages

• Better translations
• More accurate generation
• Explores multiple hypotheses

Disadvantages

• Slower than Greedy Search
• Requires more memory

Simplified Beam Search Concept

pythonCopy
1beam = [
2    ("<BOS>", 0.0)
3]
4
5# Expand top sequences
6# Keep highest scoring beams

In production systems, beam search tracks cumulative log-probabilities and prunes lower-scoring candidates after each decoding step.

Practice — Machine Translation Transformer

The following example builds a simplified encoder-decoder Transformer using PyTorch's built-in nn.Transformer.

pythonCopy
1import torch
2import torch.nn as nn
3
4class TranslationTransformer(nn.Module):
5
6    def __init__(
7        self,
8        vocab_size,
9        embed_dim=512,
10        num_heads=8,
11        num_encoder_layers=6,
12        num_decoder_layers=6
13    ):
14        super().__init__()
15
16        self.embedding = nn.Embedding(
17            vocab_size,
18            embed_dim
19        )
20
21        self.transformer = nn.Transformer(
22            d_model=embed_dim,
23            nhead=num_heads,
24            num_encoder_layers=num_encoder_layers,
25            num_decoder_layers=num_decoder_layers,
26            batch_first=True
27        )
28
29        self.output_layer = nn.Linear(
30            embed_dim,
31            vocab_size
32        )
33
34    def forward(
35        self,
36        src,
37        tgt,
38        tgt_mask=None,
39        src_padding_mask=None,
40        tgt_padding_mask=None
41    ):
42
43        src = self.embedding(src)
44        tgt = self.embedding(tgt)
45
46        output = self.transformer(
47            src=src,
48            tgt=tgt,
49            tgt_mask=tgt_mask,
50            src_key_padding_mask=src_padding_mask,
51            tgt_key_padding_mask=tgt_padding_mask
52        )
53
54        return self.output_layer(output)
55
56
57vocab_size = 30000
58
59model = TranslationTransformer(vocab_size)
60
61src = torch.randint(
62    0,
63    vocab_size,
64    (2, 12)
65)
66
67tgt = torch.randint(
68    0,
69    vocab_size,
70    (2, 10)
71)
72
73mask = nn.Transformer.generate_square_subsequent_mask(10)
74
75output = model(
76    src,
77    tgt,
78    tgt_mask=mask
79)
80
81print(output.shape)

Output

textCopy
1torch.Size([2, 10, 30000])

What You'll Learn

• Build a complete encoder-decoder Transformer
• Embed source and target token sequences
• Apply causal masking during decoding
• Produce vocabulary logits for every output position
• Understand the complete data flow from input sentence to translated output

End-to-End Transformer Workflow

textCopy
1Input Sentence
2      │
3Tokenization
4      │
5Input Embeddings
6      │
7Positional Encoding
8      │
9Encoder Stack
10      │
11Encoder Memory
12      │
13Masked Decoder Input
14      │
15Decoder Stack
16      │
17Linear Projection
18      │
19Softmax
20      │
21Next Token Prediction
22      │
23Repeat Until <EOS>

Module Summary

After completing this module, you will be able to:

• Explain the complete encoder-decoder Transformer architecture.
• Create input and output embeddings for source and target sequences.
• Apply padding masks and causal masks correctly.
• Understand the complete Transformer training pipeline.
• Explain autoregressive inference and token-by-token generation.
• Compare Greedy Search and Beam Search decoding strategies.
• Build a complete machine translation Transformer using PyTorch.
• Trace the entire workflow from tokenization to generated output.

Next Module: Module 10 – Transformer Variants, where you'll explore encoder-only (BERT), decoder-only (GPT), encoder-decoder (T5, BART), Vision Transformers (ViT), Swin Transformers, and multimodal Transformer architectures.