Module 8 — Transformer Decoder

Introduction

The Transformer Decoder is responsible for generating output sequences one token at a time. While the encoder focuses on understanding the input sequence, the decoder generates the target sequence by combining information from:

1. Previously generated output tokens
2. The encoder's contextual representations

Unlike the encoder, the decoder contains an additional attention mechanism called Cross-Attention, which allows it to attend to the encoder output while generating each token.

Decoder-based architectures are widely used in modern Large Language Models (LLMs), including:

• GPT
• LLaMA
• Mistral
• Gemma
• Qwen
• DeepSeek

In this module, you'll learn:

• Decoder Block
• Masked Self-Attention
• Cross-Attention
• Feed Forward Layer
• Residual Connections
• Layer Normalization
• Output Projection
• Build a Transformer Decoder from Scratch

1. Decoder Block

A Transformer Decoder block is similar to the encoder but includes one additional attention layer.

Architecture

textCopy
1                Input Tokens
2                      │
3          Positional Encoding
4                      │
5                      ▼
6      ┌────────────────────────┐
7      │ Masked Self-Attention  │
8      └────────────────────────┘
9                      │
10              Add & LayerNorm
11                      │
12                      ▼
13      ┌────────────────────────┐
14      │    Cross Attention     │
15      └────────────────────────┘
16                      │
17              Add & LayerNorm
18                      │
19                      ▼
20      ┌────────────────────────┐
21      │ Feed Forward Network   │
22      └────────────────────────┘
23                      │
24              Add & LayerNorm
25                      │
26                      ▼
27                  Output

Each decoder layer contains:

• Masked Multi-Head Self-Attention
• Cross-Attention
• Feed Forward Network (FFN)
• Residual Connections
• Layer Normalization
• Dropout

2. Masked Self-Attention

Unlike the encoder, the decoder cannot see future tokens during training.

Example

Sentence:

textCopy
1I love Transformers

When predicting:

textCopy
1love

The decoder can only attend to:

textCopy
1I
2love

It cannot attend to:

textCopy
1Transformers

This is enforced using a causal mask.

Attention Matrix

textCopy
1        I   Love  Transformers
2
3I       ✓     ✗        ✗
4
5Love    ✓     ✓        ✗
6
7Trans   ✓     ✓        ✓

Upper-triangular positions are masked to prevent information leakage.

Create a Causal Mask

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

Output

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

3. Cross-Attention

Cross-Attention connects the decoder with the encoder.

Source sentence

textCopy
1English
2
3↓
4
5Encoder

Target sentence

textCopy
1French
2
3↓
4
5Decoder

During Cross-Attention:

• Query comes from the decoder.
• Key comes from the encoder.
• Value comes from the encoder.

Architecture

textCopy
1Encoder Output
2     │
3     ├── Keys
4     └── Values
5
6Decoder Hidden State
7        │
8      Query
9        │
10        ▼
11 Cross Attention
12        │
13        ▼
14 Updated Decoder State

Cross-Attention enables the decoder to focus on the most relevant source words while generating each output token.

PyTorch Example

pythonCopy
1import torch
2import torch.nn as nn
3
4cross_attention = nn.MultiheadAttention(
5    embed_dim=512,
6    num_heads=8,
7    batch_first=True
8)
9
10decoder = torch.randn(2, 8, 512)
11
12encoder = torch.randn(2, 10, 512)
13
14output, weights = cross_attention(
15    decoder,
16    encoder,
17    encoder
18)
19
20print(output.shape)

Output

textCopy
1torch.Size([2, 8, 512])

4. Feed Forward Layer

After attention, every token independently passes through a feed-forward neural network.

Architecture

textCopy
1Input
2 │
3 ▼
4Linear
5 │
6GELU
7 │
8Linear
9 │
10Output

The FFN expands the embedding dimension and then projects it back.

Build Feed Forward Layer

pythonCopy
1import torch
2import torch.nn as nn
3
4class FeedForward(nn.Module):
5
6    def __init__(self, embed_dim):
7
8        super().__init__()
9
10        self.network = nn.Sequential(
11
12            nn.Linear(embed_dim, embed_dim * 4),
13
14            nn.GELU(),
15
16            nn.Linear(embed_dim * 4, embed_dim)
17        )
18
19    def forward(self, x):
20
21        return self.network(x)
22
23
24ffn = FeedForward(512)
25
26x = torch.randn(2, 6, 512)
27
28print(ffn(x).shape)

Output

textCopy
1torch.Size([2, 6, 512])

5. Residual Connections

Residual connections allow information to bypass individual sublayers.

Instead of learning:

textCopy
1Output = F(x)

The decoder learns:

textCopy
1Output = x + F(x)

Advantages

• Stable gradients
• Easier optimization
• Enables deep Transformer models

Example

pythonCopy
1import torch
2
3x = torch.randn(2, 5, 512)
4
5attention_output = torch.randn(2, 5, 512)
6
7output = x + attention_output
8
9print(output.shape)

6. Layer Normalization

Layer Normalization stabilizes training by normalizing features across the embedding dimension.

Benefits

• Stable gradients
• Faster convergence
• Improved optimization

PyTorch Example

pythonCopy
1import torch
2import torch.nn as nn
3
4layer_norm = nn.LayerNorm(512)
5
6x = torch.randn(2, 6, 512)
7
8output = layer_norm(x)
9
10print(output.shape)

Output

textCopy
1torch.Size([2, 6, 512])

7. Output Projection

After the decoder finishes processing, its hidden states are converted into vocabulary logits.

Pipeline

textCopy
1Decoder Output
2
3↓
4
5Linear Layer
6
7↓
8
9Vocabulary Logits
10
11↓
12
13Softmax
14
15↓
16
17Predicted Token

Suppose:

textCopy
1Vocabulary Size = 30,000
2Embedding Size = 768

The final projection is:

textCopy
1768 → 30,000

PyTorch Example

pythonCopy
1import torch
2import torch.nn as nn
3
4vocab_size = 30000
5
6projection = nn.Linear(
7    512,
8    vocab_size
9)
10
11decoder_output = torch.randn(2, 10, 512)
12
13logits = projection(decoder_output)
14
15print(logits.shape)

Output

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

The logits are later passed through Softmax during inference to obtain token probabilities.

Practice — Build a Transformer Decoder

The following implementation combines masked self-attention, cross-attention, feed-forward layers, residual connections, LayerNorm, and dropout into a complete decoder layer.

pythonCopy
1import torch
2import torch.nn as nn
3
4class TransformerDecoderLayer(nn.Module):
5
6    def __init__(
7        self,
8        embed_dim,
9        num_heads,
10        dropout=0.1
11    ):
12        super().__init__()
13
14        self.self_attention = nn.MultiheadAttention(
15            embed_dim,
16            num_heads,
17            dropout=dropout,
18            batch_first=True
19        )
20
21        self.cross_attention = nn.MultiheadAttention(
22            embed_dim,
23            num_heads,
24            dropout=dropout,
25            batch_first=True
26        )
27
28        self.ffn = nn.Sequential(
29            nn.Linear(embed_dim, embed_dim * 4),
30            nn.GELU(),
31            nn.Linear(embed_dim * 4, embed_dim)
32        )
33
34        self.norm1 = nn.LayerNorm(embed_dim)
35        self.norm2 = nn.LayerNorm(embed_dim)
36        self.norm3 = nn.LayerNorm(embed_dim)
37
38        self.dropout = nn.Dropout(dropout)
39
40    def forward(
41        self,
42        target,
43        memory,
44        mask=None
45    ):
46
47        self_output, _ = self.self_attention(
48            target,
49            target,
50            target,
51            attn_mask=mask
52        )
53
54        target = self.norm1(
55            target + self.dropout(self_output)
56        )
57
58        cross_output, _ = self.cross_attention(
59            target,
60            memory,
61            memory
62        )
63
64        target = self.norm2(
65            target + self.dropout(cross_output)
66        )
67
68        ffn_output = self.ffn(target)
69
70        target = self.norm3(
71            target + self.dropout(ffn_output)
72        )
73
74        return target
75
76
77decoder = TransformerDecoderLayer(
78    embed_dim=512,
79    num_heads=8
80)
81
82target = torch.randn(2, 8, 512)
83
84memory = torch.randn(2, 10, 512)
85
86mask = torch.triu(
87    torch.ones(8, 8),
88    diagonal=1
89).bool()
90
91output = decoder(
92    target,
93    memory,
94    mask
95)
96
97print(output.shape)

Output

textCopy
1torch.Size([2, 8, 512])

Bonus — Stack Multiple Decoder Layers

pythonCopy
1class TransformerDecoder(nn.Module):
2
3    def __init__(
4        self,
5        embed_dim,
6        num_heads,
7        num_layers
8    ):
9        super().__init__()
10
11        self.layers = nn.ModuleList(
12            [
13                TransformerDecoderLayer(
14                    embed_dim,
15                    num_heads
16                )
17                for _ in range(num_layers)
18            ]
19        )
20
21    def forward(
22        self,
23        target,
24        memory,
25        mask=None
26    ):
27
28        for layer in self.layers:
29            target = layer(
30                target,
31                memory,
32                mask
33            )
34
35        return target
36
37
38model = TransformerDecoder(
39    embed_dim=512,
40    num_heads=8,
41    num_layers=6
42)
43
44output = model(
45    target,
46    memory,
47    mask
48)
49
50print(output.shape)

Output

textCopy
1torch.Size([2, 8, 512])

Module Summary

After completing this module, you will be able to:

• Explain the architecture of a Transformer Decoder block.
• Understand how Masked Self-Attention prevents future token leakage.
• Build and use Cross-Attention to connect encoder and decoder representations.
• Implement a Feed Forward Network (FFN) with GELU activation.
• Apply Residual Connections and Layer Normalization for stable training.
• Convert decoder hidden states into vocabulary logits using an output projection layer.
• Stack multiple decoder layers to build a deep Transformer decoder.
• Implement a complete Transformer Decoder from scratch using PyTorch.

Next Module: Module 9 – Complete Transformer, where you'll combine the encoder, decoder, embeddings, positional encoding, masking, and output projection to build the original Transformer architecture described in Attention Is All You Need.