Loading topic...

PyTorch Transformer Decoder Tutorial with Code Examples

Transformer Decoder

Learn Transformer Decoder, Masked Self Attention, Cross Attention, Output Projection, and implement a complete decoder in PyTorch.

PyTorch Transformer Decoder: Architecture, Masked Self-Attention, Cross-Attention, and Implementation

A Transformer decoder is a neural network component designed to generate or transform sequences by processing previously available tokens and, in encoder-decoder architectures, attending to representations produced by an encoder.

The Transformer decoder is a fundamental component of sequence generation systems. It is used directly in decoder-only language models such as GPT and LLaMA, while encoder-decoder architectures such as T5 and BART use a decoder together with an encoder.

In this module, you will learn how a Transformer decoder works internally, why causal masking is required, how cross-attention connects a decoder to an encoder, and how to implement a decoder using PyTorch.

What You Will Learn

By the end of this module, you will understand:

What a Transformer decoder is
How a Transformer decoder block is structured
Masked self-attention and causal attention
Why future-token masking is necessary
Cross-attention in encoder-decoder Transformers
Feed-forward networks
Residual connections
Layer normalization
Output projection and vocabulary logits
How to implement a Transformer decoder in PyTorch
How to use nn.TransformerDecoder
How autoregressive decoding works
Important decoder implementation and training practices

What Is a Transformer Decoder?

A Transformer decoder processes a target sequence and produces contextual representations that can be used to predict the next token.

During autoregressive generation, the model generates tokens progressively:

text
1<SOS>
2  ↓
3Token 1
4  ↓
5Token 2
6  ↓
7Token 3
8  ↓
9...
10  ↓
11<EOS>

At each generation step, the decoder can use information from tokens that have already been generated.

In an encoder-decoder Transformer, the decoder also receives information from the encoder.

The high-level architecture is:

text
1Source Sequence
2      │
3      ▼
4Transformer Encoder
5      │
6      ▼
7Encoder Representations
8      │
9      │
10      ▼
11Target Tokens ──► Transformer Decoder
12                        │
13                        ▼
14                 Vocabulary Logits
15                        │
16                        ▼
17                  Next Token

Decoder-Only vs Encoder-Decoder Transformers

It is important to distinguish two common Transformer designs.

Decoder-Only Architecture

Models such as GPT and LLaMA use decoder-style Transformer blocks without a separate encoder.

text
1Input Tokens
2     │
3     ▼
4Masked Self-Attention
5     │
6     ▼
7Feed Forward Network
8     │
9     ▼
10Output Logits
11     │
12     ▼
13Next Token

These models are primarily used for autoregressive generation.

Encoder-Decoder Architecture

Models such as T5 and BART contain both an encoder and a decoder.

text
1Source
2  │
3  ▼
4Encoder
5  │
6  ▼
7Memory
8  │
9  ├───────────────┐
10  │               ▼
11  │          Cross-Attention
12  │               ▲
13  │               │
14  ▼               │
15Decoder ◄─────────┘
16  │
17  ▼
18Output

The decoder uses cross-attention to access information from the encoder.

Why Do We Need a Transformer Decoder?

Consider machine translation.

Input:

text
1English:
2I love AI

Target:

text
1French:
2J'aime l'IA

The encoder processes the source sentence and creates contextual representations.

The decoder then generates the target sequence autoregressively:

text
1<SOS>
2   ↓
3J'aime
4   ↓
5J'aime l'IA
6   ↓
7J'aime l'IA <EOS>

During generation, the decoder uses the tokens already generated to determine what should come next.

Transformer Decoder Architecture

A typical Transformer decoder block contains three major sub-layers:

text
1                 Decoder Input
2                       │
3                       ▼
4              Masked Self-Attention
5                       │
6                       ▼
7              Residual + LayerNorm
8                       │
9                       ▼
10                 Cross-Attention
11                       │
12                       ▼
13              Residual + LayerNorm
14                       │
15                       ▼
16              Feed-Forward Network
17                       │
18                       ▼
19              Residual + LayerNorm
20                       │
21                       ▼
22                 Decoder Output

The cross-attention layer is present in encoder-decoder Transformers.

Decoder-only models omit cross-attention because there is no separate encoder representation to attend to.

Decoder Block Components

A decoder block typically contains:

Masked self-attention
Cross-attention in encoder-decoder models
Feed-forward network
Residual connections
Layer normalization

These components are repeated across multiple decoder layers.

For example:

text
1Decoder Layer 1
2      ↓
3Decoder Layer 2
4      ↓
5Decoder Layer 3
6      ↓
7...
8      ↓
9Decoder Layer N

Masked Self-Attention

Self-attention allows tokens to interact with other tokens.

However, autoregressive generation requires an important restriction:

A token must not use information from future target tokens.

Suppose the target sequence is:

text
1I love deep learning

When predicting:

text
1love

the model should be able to use:

text
1I

but it should not use:

text
1deep learning

Otherwise, the model would receive information that would not be available during real generation.

Why Causal Masking Is Required

Without causal masking:

text
1I     love     deep     learning
2│       │        │         │
3└───────┴────────┴─────────┘
4Every token can potentially attend to every token

With causal masking:

text
1I
2↓
3I love
4↓
5I love deep
6↓
7I love deep learning

Each position can attend only to itself and earlier positions.

This is also called:

Causal attention
Autoregressive attention
Look-ahead masking
Future-token masking

Causal Attention Matrix

For a sequence of four tokens:

text
1        Token 0  Token 1  Token 2  Token 3
2
3Token 0    ✓        ✗        ✗        ✗
4
5Token 1    ✓        ✓        ✗        ✗
6
7Token 2    ✓        ✓        ✓        ✗
8
9Token 3    ✓        ✓        ✓        ✓

The diagonal and lower-triangular portion are visible.

The upper-triangular portion represents future positions and must be masked.

Creating a Causal Mask in PyTorch

A causal mask can be created using torch.triu().

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

Output:

text
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])

Here:

text
1False → attention is allowed
2True  → attention is blocked

When using PyTorch attention APIs, the exact interpretation depends on the mask argument being used, so always check the specific API semantics.

Masked Self-Attention in PyTorch

PyTorch provides nn.MultiheadAttention.

python
1import torch
2import torch.nn as nn
3
4attention = nn.MultiheadAttention(
5    embed_dim=512,
6    num_heads=8,
7    batch_first=True
8)
9
10x = torch.randn(
11    2,
12    5,
13    512
14)
15
16mask = torch.triu(
17    torch.ones(
18        5,
19        5
20    ),
21    diagonal=1
22).bool()
23
24output, weights = attention(
25    x,
26    x,
27    x,
28    attn_mask=mask
29)
30
31print(output.shape)

Output:

text
1torch.Size([2, 5, 512])

The three inputs represent:

text
1Query
2Key
3Value

When all three are the same tensor:

python
1attention(x, x, x)

the operation performs self-attention.

Cross-Attention

Cross-attention connects the decoder to the encoder.

Unlike self-attention, the query and key-value representations come from different sequences.

The relationship is:

text
1Query
2↓
3Decoder
4
5Key
6↓
7Encoder
8
9Value
10↓
11Encoder

Therefore:

text
1Q = Decoder Representation
2
3K = Encoder Representation
4
5V = Encoder Representation

This allows the decoder to determine which parts of the source sequence are relevant when generating the target sequence.

Cross-Attention Example

Suppose the encoder processes:

text
1I love AI

The encoder creates contextual representations:

text
1Encoder Output
2      │
3      ├── representation of "I"
4      ├── representation of "love"
5      └── representation of "AI"

When generating the target language, the decoder can use cross-attention to determine which source representations are relevant for the current output token.

Cross-Attention in PyTorch

python
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
10encoder_output = torch.randn(
11    2,
12    10,
13    512
14)
15
16decoder_hidden = torch.randn(
17    2,
18    6,
19    512
20)
21
22output, weights = cross_attention(
23    decoder_hidden,
24    encoder_output,
25    encoder_output
26)
27
28print(output.shape)

Output:

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

The output sequence length is determined by the decoder query sequence.

Feed-Forward Network

After attention operations, the decoder applies a position-wise feed-forward network.

A simplified version is:

text
1Input
2  │
3  ▼
4Linear
5  │
6  ▼
7Activation
8  │
9  ▼
10Dropout
11  │
12  ▼
13Linear
14  │
15  ▼
16Output

A PyTorch implementation is:

python
1import torch.nn as nn
2
3ffn = nn.Sequential(
4    nn.Linear(
5        512,
6        2048
7    ),
8
9    nn.GELU(),
10
11    nn.Dropout(0.1),
12
13    nn.Linear(
14        2048,
15        512
16    )
17)

The feed-forward network is applied independently to each sequence position while using shared parameters across positions.

Residual Connections

Residual connections add the input of a sub-layer to its output.

Conceptually:

text
1Input
2  │
3  ├──────────────┐
4  │              │
5  ▼              │
6Attention        │
7  │              │
8  └──────► Add ◄─┘
9             │
10             ▼
11          Output

Mathematically:

text
1Output = Input + SubLayer(Input)

Residual connections help information and gradients propagate through deep Transformer networks.

Layer Normalization

Layer normalization stabilizes the representations flowing through the network.

A simplified Transformer operation can be represented as:

text
1x
2│
3▼
4Attention
5│
6▼
7Dropout
8│
9▼
10Add x
11│
12▼
13LayerNorm

PyTorch provides:

python
1nn.LayerNorm(d_model)

For example:

python
1norm = nn.LayerNorm(512)
2
3x = torch.randn(
4    2,
5    6,
6    512
7)
8
9output = norm(x)
10
11print(output.shape)

Output:

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

Pre-Norm and Post-Norm

Transformer implementations can organize normalization differently.

Post-Norm

The original Transformer architecture is commonly described using:

text
1x → Sublayer → Add → LayerNorm

Pre-Norm

Many modern Transformer implementations use:

text
1x → LayerNorm → Sublayer → Add

Pre-normalization can provide more stable optimization for deep Transformer networks.

When implementing a decoder from scratch, it is important to understand which normalization arrangement your architecture is using.

Output Projection

The decoder produces hidden representations, not directly token IDs.

A linear layer maps each hidden representation to vocabulary logits.

text
1Decoder Hidden State
2        │
3        ▼
4    Linear Layer
5        │
6        ▼
7Vocabulary Logits
8        │
9        ▼
10Next Token

Suppose:

text
1d_model = 512
2vocabulary size = 10,000

Then the projection is:

python
1projection = nn.Linear(
2    512,
3    10000
4)

Example:

python
1x = torch.randn(
2    2,
3    8,
4    512
5)
6
7logits = projection(x)
8
9print(logits.shape)

Output:

text
1torch.Size([2, 8, 10000])

The dimensions mean:

text
12      → batch size
28      → target sequence length
310000  → vocabulary size

The values are logits, not token IDs.

From Logits to Tokens

The model produces vocabulary scores:

text
1Vocabulary
2
3hello     → 2.1
4world     → 0.4
5AI        → 4.7
6model     → 1.8
7...

A decoding strategy converts these scores into the next token.

The simplest strategy is greedy decoding:

python
1next_token = logits[:, -1].argmax(dim=-1)

More advanced generation methods include:

Temperature sampling
Top-k sampling
Top-p sampling
Beam search

Building a Transformer Decoder Block

Now combine masked self-attention, cross-attention, feed-forward layers, residual connections, and normalization.

python
1import torch
2import torch.nn as nn
3
4
5class DecoderBlock(nn.Module):
6
7    def __init__(
8        self,
9        d_model=512,
10        n_heads=8,
11        ff_dim=2048,
12        dropout=0.1
13    ):
14        super().__init__()
15
16        self.self_attention = nn.MultiheadAttention(
17            embed_dim=d_model,
18            num_heads=n_heads,
19            dropout=dropout,
20            batch_first=True
21        )
22
23        self.cross_attention = nn.MultiheadAttention(
24            embed_dim=d_model,
25            num_heads=n_heads,
26            dropout=dropout,
27            batch_first=True
28        )
29
30        self.ffn = nn.Sequential(
31            nn.Linear(
32                d_model,
33                ff_dim
34            ),
35
36            nn.GELU(),
37
38            nn.Dropout(dropout),
39
40            nn.Linear(
41                ff_dim,
42                d_model
43            )
44        )
45
46        self.norm1 = nn.LayerNorm(d_model)
47        self.norm2 = nn.LayerNorm(d_model)
48        self.norm3 = nn.LayerNorm(d_model)
49
50        self.dropout = nn.Dropout(dropout)
51
52    def forward(
53        self,
54        x,
55        encoder_output,
56        tgt_mask=None,
57        memory_key_padding_mask=None
58    ):
59
60        self_attn, _ = self.self_attention(
61            x,
62            x,
63            x,
64            attn_mask=tgt_mask
65        )
66
67        x = self.norm1(
68            x + self.dropout(self_attn)
69        )
70
71        cross_attn, _ = self.cross_attention(
72            x,
73            encoder_output,
74            encoder_output,
75            key_padding_mask=memory_key_padding_mask
76        )
77
78        x = self.norm2(
79            x + self.dropout(cross_attn)
80        )
81
82        ffn_output = self.ffn(x)
83
84        x = self.norm3(
85            x + self.dropout(ffn_output)
86        )
87
88        return x

This implementation demonstrates the core structure of an encoder-decoder Transformer decoder block.

Testing the Decoder Block

Create example encoder and decoder tensors:

python
1decoder = DecoderBlock()
2
3encoder_output = torch.randn(
4    2,
5    10,
6    512
7)
8
9decoder_input = torch.randn(
10    2,
11    6,
12    512
13)
14
15mask = torch.triu(
16    torch.ones(
17        6,
18        6
19    ),
20    diagonal=1
21).bool()
22
23output = decoder(
24    decoder_input,
25    encoder_output,
26    tgt_mask=mask
27)
28
29print(output.shape)

Output:

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

The decoder preserves the target sequence length and model dimension.

Building Multiple Decoder Layers

A Transformer decoder usually contains several decoder blocks.

python
1class TransformerDecoder(nn.Module):
2
3    def __init__(
4        self,
5        num_layers=6,
6        d_model=512,
7        n_heads=8,
8        ff_dim=2048,
9        dropout=0.1
10    ):
11        super().__init__()
12
13        self.layers = nn.ModuleList([
14            DecoderBlock(
15                d_model=d_model,
16                n_heads=n_heads,
17                ff_dim=ff_dim,
18                dropout=dropout
19            )
20            for _ in range(num_layers)
21        ])
22
23    def forward(
24        self,
25        x,
26        encoder_output,
27        tgt_mask=None,
28        memory_key_padding_mask=None
29    ):
30
31        for layer in self.layers:
32
33            x = layer(
34                x,
35                encoder_output,
36                tgt_mask=tgt_mask,
37                memory_key_padding_mask=memory_key_padding_mask
38            )
39
40        return x

The input passes through each decoder block sequentially:

text
1Decoder Input
2      ↓
3Layer 1
4      ↓
5Layer 2
6      ↓
7Layer 3
8      ↓
9...
10      ↓
11Layer 6
12      ↓
13Decoder Output

Testing the Transformer Decoder

python
1decoder = TransformerDecoder()
2
3encoder_output = torch.randn(
4    2,
5    10,
6    512
7)
8
9decoder_input = torch.randn(
10    2,
11    6,
12    512
13)
14
15mask = torch.triu(
16    torch.ones(
17        6,
18        6
19    ),
20    diagonal=1
21).bool()
22
23output = decoder(
24    decoder_input,
25    encoder_output,
26    tgt_mask=mask
27)
28
29print(output.shape)

Output:

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

Using PyTorch's TransformerDecoder

PyTorch already provides a Transformer decoder implementation.

python
1import torch
2import torch.nn as nn
3
4decoder_layer = nn.TransformerDecoderLayer(
5    d_model=512,
6    nhead=8,
7    dim_feedforward=2048,
8    dropout=0.1,
9    batch_first=True
10)
11
12decoder = nn.TransformerDecoder(
13    decoder_layer,
14    num_layers=6
15)

Create example inputs:

python
1memory = torch.randn(
2    2,
3    10,
4    512
5)
6
7tgt = torch.randn(
8    2,
9    6,
10    512
11)

Create a causal mask:

python
1mask = torch.triu(
2    torch.ones(
3        6,
4        6
5    ),
6    diagonal=1
7).bool()

Run the decoder:

python
1output = decoder(
2    tgt,
3    memory,
4    tgt_mask=mask
5)
6
7print(output.shape)

Output:

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

Here:

text
1tgt
2↓
3Decoder input
4
5memory
6↓
7Encoder output

The decoder uses masked self-attention over tgt and cross-attention over memory.

Decoder Padding Masks

Causal masking and padding masking solve different problems.

Causal Mask

Prevents a position from seeing future target positions.

text
1Future information
2       ↓
3     BLOCKED

Padding Mask

Prevents attention to padding tokens.

text
1<PAD>
2<PAD>
3  ↓
4Ignored by attention

A real sequence-to-sequence implementation may need both.

For example:

python
1tgt_padding_mask = (
2    tgt_tokens == pad_token_id
3)
4
5memory_padding_mask = (
6    src_tokens == pad_token_id
7)

The exact argument names depend on the PyTorch module being used.

Decoder Training With Teacher Forcing

During training, sequence-to-sequence models commonly use teacher forcing.

Suppose the target sentence is:

text
1<SOS> I love AI <EOS>

The decoder input becomes:

text
1<SOS> I love AI

The expected labels become:

text
1I love AI <EOS>

The model therefore learns:

text
1<SOS>
2   ↓
3I
4
5<SOS> I
6   ↓
7love
8
9<SOS> I love
10   ↓
11AI
12
13<SOS> I love AI
14   ↓
15<EOS>

This shifted-target setup is essential for standard autoregressive training.

Decoder Loss

Suppose the decoder produces:

text
1(batch, target_length, vocabulary_size)

Cross-entropy loss can be calculated after flattening the batch and sequence dimensions.

python
1criterion = nn.CrossEntropyLoss(
2    ignore_index=pad_token_id
3)
4
5loss = criterion(
6    logits.reshape(
7        -1,
8        vocabulary_size
9    ),
10    target.reshape(-1)
11)

The padding index is ignored so padded positions do not contribute to the loss.

Autoregressive Generation

During inference, the target sequence is not known in advance.

The decoder generates one token at a time.

text
1<SOS>
2  │
3  ▼
4Decoder
5  │
6  ▼
7Token 1
8  │
9  ▼
10Decoder
11  │
12  ▼
13Token 2
14  │
15  ▼
16Decoder
17  │
18  ▼
19Token 3

Generation continues until:

text
1<EOS>

or until the maximum generation length is reached.

Greedy Decoding

A simple decoding algorithm selects the token with the highest logit.

python
1next_token = logits[:, -1].argmax(
2    dim=-1
3)

This is easy to implement but does not always produce the best sequence.

For advanced generation, consider:

Beam search
Top-k sampling
Top-p sampling
Temperature scaling

KV Caching During Generation

Autoregressive generation can become expensive because the model repeatedly processes previous tokens.

Modern decoder-only Transformer implementations commonly use key-value caching.

Conceptually:

text
1Previous K/V
2     │
3     ▼
4Cached
5     │
6     +
7New Token
8     │
9     ▼
10New Attention Computation

Instead of recomputing all previous key and value representations at every generation step, cached values can be reused.

KV caching is an important optimization for large language model inference.

Decoder-Only Transformer Architecture

For a GPT- or LLaMA-style decoder-only model, the architecture is simplified because there is no encoder or cross-attention.

text
1Token IDs
2    │
3    ▼
4Token Embeddings
5    │
6    +
7Position Information
8    │
9    ▼
10Masked Self-Attention
11    │
12    ▼
13Feed-Forward Network
14    │
15    ▼
16Repeated Decoder Blocks
17    │
18    ▼
19Final Normalization
20    │
21    ▼
22Language Modeling Head
23    │
24    ▼
25Vocabulary Logits

This distinction is important when studying modern large language models.

Encoder-Decoder Transformer Architecture

For translation and other sequence-to-sequence tasks:

text
1                Source Tokens
2                     │
3                     ▼
4              Transformer Encoder
5                     │
6                     ▼
7                Encoder Memory
8                     │
9                     ▼
10Target Tokens → Transformer Decoder
11                     │
12                     ▼
13               Vocabulary Logits
14                     │
15                     ▼
16                Target Tokens

The decoder uses:

text
1Masked Self-Attention
2        +
3Cross-Attention
4        +
5Feed-Forward Network

Decoder vs Encoder

Feature	Transformer Encoder	Transformer Decoder
Self-attention	Yes	Yes
Causal masking	Usually no	Yes for autoregressive decoding
Cross-attention	No	Yes in encoder-decoder architectures
Uses encoder memory	No	Yes when applicable
Generates tokens	No	Yes
Typical role	Input representation	Sequence generation

The exact behavior depends on the model architecture and task.

Common Transformer Decoder Mistakes

Forgetting the Causal Mask

Without causal masking during autoregressive training, the decoder can access future target positions.

Confusing Causal and Padding Masks

They perform different functions.

text
1Causal mask
2→ blocks future positions
3
4Padding mask
5→ blocks padding tokens

Using Incorrect Target Shifting

The decoder input and expected target should normally be shifted during teacher-forced training.

Ignoring Padding in the Loss

Padding tokens should normally be excluded from the training loss.

Forgetting Cross-Attention in Encoder-Decoder Models

An encoder-decoder Transformer requires the decoder to attend to encoder representations.

Confusing Logits With Probabilities

The output projection normally produces logits.

Do not apply softmax before passing logits to CrossEntropyLoss, because CrossEntropyLoss internally combines the required operations.

Ignoring Tensor Shapes

Common shapes include:

text
1Input tokens:
2(batch, sequence_length)
3
4Embeddings:
5(batch, sequence_length, d_model)
6
7Decoder output:
8(batch, target_length, d_model)
9
10Vocabulary logits:
11(batch, target_length, vocabulary_size)

Incorrect dimensions are one of the most common sources of Transformer implementation errors.

Best Practices for Transformer Decoders

Use causal masking for autoregressive target generation.
Keep causal masking separate from padding masking.
Use shifted target sequences during teacher-forced training.
Ignore padding tokens in the loss.
Use residual connections around decoder sub-layers.
Use LayerNorm consistently with the chosen architecture.
Use cross-attention when implementing an encoder-decoder Transformer.
Verify tensor shapes at every major stage.
Use model.train() during training.
Use model.eval() during evaluation and inference.
Use torch.no_grad() when gradients are unnecessary during inference.
Consider KV caching for efficient autoregressive generation.
Start with a small model when debugging custom implementations.
Use PyTorch's optimized Transformer components when a custom implementation is not required.

Practice Project: Implement a Transformer Decoder

Build a small decoder with:

text
1d_model = 512
2attention_heads = 8
3feed_forward_dimension = 2048
4decoder_layers = 6

Step 1: Create Encoder Output

python
1encoder_output = torch.randn(
2    4,
3    15,
4    512
5)

Step 2: Create Decoder Input

python
1decoder_input = torch.randn(
2    4,
3    10,
4    512
5)

Step 3: Create Causal Mask

python
1mask = torch.triu(
2    torch.ones(
3        10,
4        10
5    ),
6    diagonal=1
7).bool()

Step 4: Run the Decoder

python
1decoder = TransformerDecoder()
2
3output = decoder(
4    decoder_input,
5    encoder_output,
6    tgt_mask=mask
7)
8
9print(output.shape)

Expected output:

text
1torch.Size([4, 10, 512])

Step 5: Project to the Vocabulary

python
1projection = nn.Linear(
2    512,
3    30000
4)
5
6logits = projection(output)
7
8print(logits.shape)

Expected output:

text
1torch.Size([4, 10, 30000])

This means the model produces 30,000 vocabulary scores for every one of the 10 target positions in each of the 4 examples.

Transformer Decoder Learning Checklist

Before moving to advanced Transformer architectures, make sure you can explain:

What a decoder does
Why autoregressive models need causal masking
How masked self-attention works
How query, key, and value differ
How cross-attention connects an encoder and decoder
Why residual connections are used
What LayerNorm does
How the feed-forward network works
How hidden states become vocabulary logits
Why decoder targets are shifted during training
Why padding should be ignored in the loss
How greedy decoding works
Why KV caching improves autoregressive inference
The difference between decoder-only and encoder-decoder Transformers

Module Summary

A Transformer decoder is responsible for generating target representations and, ultimately, target tokens.

A decoder block can be summarized as:

text
1Decoder Input
2      │
3      ▼
4Masked Self-Attention
5      │
6      ▼
7Residual + LayerNorm
8      │
9      ▼
10Cross-Attention
11      │
12      ▼
13Residual + LayerNorm
14      │
15      ▼
16Feed-Forward Network
17      │
18      ▼
19Residual + LayerNorm
20      │
21      ▼
22Decoder Output

For decoder-only models, the cross-attention stage is removed:

text
1Token Embedding
2      │
3      ▼
4Causal Self-Attention
5      │
6      ▼
7Feed-Forward Network
8      │
9      ▼
10Repeated Decoder Blocks
11      │
12      ▼
13Language Modeling Head
14      │
15      ▼
16Next-Token Prediction

The most important concept is causal attention: during autoregressive generation, the decoder must not use future target tokens to predict the current token.

Once you understand masked self-attention, cross-attention, residual connections, normalization, feed-forward networks, output projections, and autoregressive generation, you have the foundation needed to study modern decoder architectures and large language models.

Loading topic...

Transformer Decoder

Learn Transformer Decoder, Masked Self Attention, Cross Attention, Output Projection, and implement a complete decoder in PyTorch.

PyTorch Transformer Decoder: Architecture, Masked Self-Attention, Cross-Attention, and Implementation

What You Will Learn

By the end of this module, you will understand:

What a Transformer decoder is
How a Transformer decoder block is structured
Masked self-attention and causal attention
Why future-token masking is necessary
Cross-attention in encoder-decoder Transformers
Feed-forward networks
Residual connections
Layer normalization
Output projection and vocabulary logits
How to implement a Transformer decoder in PyTorch
How to use nn.TransformerDecoder
How autoregressive decoding works
Important decoder implementation and training practices

What Is a Transformer Decoder?

A Transformer decoder processes a target sequence and produces contextual representations that can be used to predict the next token.

During autoregressive generation, the model generates tokens progressively:

text
1<SOS>
2  ↓
3Token 1
4  ↓
5Token 2
6  ↓
7Token 3
8  ↓
9...
10  ↓
11<EOS>

At each generation step, the decoder can use information from tokens that have already been generated.

In an encoder-decoder Transformer, the decoder also receives information from the encoder.

The high-level architecture is:

text
1Source Sequence
2      │
3      ▼
4Transformer Encoder
5      │
6      ▼
7Encoder Representations
8      │
9      │
10      ▼
11Target Tokens ──► Transformer Decoder
12                        │
13                        ▼
14                 Vocabulary Logits
15                        │
16                        ▼
17                  Next Token

Decoder-Only vs Encoder-Decoder Transformers

It is important to distinguish two common Transformer designs.

Decoder-Only Architecture

Models such as GPT and LLaMA use decoder-style Transformer blocks without a separate encoder.

text
1Input Tokens
2     │
3     ▼
4Masked Self-Attention
5     │
6     ▼
7Feed Forward Network
8     │
9     ▼
10Output Logits
11     │
12     ▼
13Next Token

These models are primarily used for autoregressive generation.

Encoder-Decoder Architecture

Models such as T5 and BART contain both an encoder and a decoder.

text
1Source
2  │
3  ▼
4Encoder
5  │
6  ▼
7Memory
8  │
9  ├───────────────┐
10  │               ▼
11  │          Cross-Attention
12  │               ▲
13  │               │
14  ▼               │
15Decoder ◄─────────┘
16  │
17  ▼
18Output

The decoder uses cross-attention to access information from the encoder.

Why Do We Need a Transformer Decoder?

Consider machine translation.

Input:

text
1English:
2I love AI

Target:

text
1French:
2J'aime l'IA

The encoder processes the source sentence and creates contextual representations.

The decoder then generates the target sequence autoregressively:

text
1<SOS>
2   ↓
3J'aime
4   ↓
5J'aime l'IA
6   ↓
7J'aime l'IA <EOS>

During generation, the decoder uses the tokens already generated to determine what should come next.

Transformer Decoder Architecture

A typical Transformer decoder block contains three major sub-layers:

text
1                 Decoder Input
2                       │
3                       ▼
4              Masked Self-Attention
5                       │
6                       ▼
7              Residual + LayerNorm
8                       │
9                       ▼
10                 Cross-Attention
11                       │
12                       ▼
13              Residual + LayerNorm
14                       │
15                       ▼
16              Feed-Forward Network
17                       │
18                       ▼
19              Residual + LayerNorm
20                       │
21                       ▼
22                 Decoder Output

The cross-attention layer is present in encoder-decoder Transformers.

Decoder-only models omit cross-attention because there is no separate encoder representation to attend to.

Decoder Block Components

A decoder block typically contains:

Masked self-attention
Cross-attention in encoder-decoder models
Feed-forward network
Residual connections
Layer normalization

These components are repeated across multiple decoder layers.

For example:

text
1Decoder Layer 1
2      ↓
3Decoder Layer 2
4      ↓
5Decoder Layer 3
6      ↓
7...
8      ↓
9Decoder Layer N

Masked Self-Attention

Self-attention allows tokens to interact with other tokens.

However, autoregressive generation requires an important restriction:

A token must not use information from future target tokens.

Suppose the target sequence is:

text
1I love deep learning

When predicting:

text
1love

the model should be able to use:

text
1I

but it should not use:

text
1deep learning

Otherwise, the model would receive information that would not be available during real generation.

Why Causal Masking Is Required

Without causal masking:

text
1I     love     deep     learning
2│       │        │         │
3└───────┴────────┴─────────┘
4Every token can potentially attend to every token

With causal masking:

text
1I
2↓
3I love
4↓
5I love deep
6↓
7I love deep learning

Each position can attend only to itself and earlier positions.

This is also called:

Causal attention
Autoregressive attention
Look-ahead masking
Future-token masking

Causal Attention Matrix

For a sequence of four tokens:

text
1        Token 0  Token 1  Token 2  Token 3
2
3Token 0    ✓        ✗        ✗        ✗
4
5Token 1    ✓        ✓        ✗        ✗
6
7Token 2    ✓        ✓        ✓        ✗
8
9Token 3    ✓        ✓        ✓        ✓

The diagonal and lower-triangular portion are visible.

The upper-triangular portion represents future positions and must be masked.

Creating a Causal Mask in PyTorch

A causal mask can be created using torch.triu().

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

Output:

text
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])

Here:

text
1False → attention is allowed
2True  → attention is blocked

When using PyTorch attention APIs, the exact interpretation depends on the mask argument being used, so always check the specific API semantics.

Masked Self-Attention in PyTorch

PyTorch provides nn.MultiheadAttention.

python
1import torch
2import torch.nn as nn
3
4attention = nn.MultiheadAttention(
5    embed_dim=512,
6    num_heads=8,
7    batch_first=True
8)
9
10x = torch.randn(
11    2,
12    5,
13    512
14)
15
16mask = torch.triu(
17    torch.ones(
18        5,
19        5
20    ),
21    diagonal=1
22).bool()
23
24output, weights = attention(
25    x,
26    x,
27    x,
28    attn_mask=mask
29)
30
31print(output.shape)

Output:

text
1torch.Size([2, 5, 512])

The three inputs represent:

text
1Query
2Key
3Value

When all three are the same tensor:

python
1attention(x, x, x)

the operation performs self-attention.

Cross-Attention

Cross-attention connects the decoder to the encoder.

Unlike self-attention, the query and key-value representations come from different sequences.

The relationship is:

text
1Query
2↓
3Decoder
4
5Key
6↓
7Encoder
8
9Value
10↓
11Encoder

Therefore:

text
1Q = Decoder Representation
2
3K = Encoder Representation
4
5V = Encoder Representation

This allows the decoder to determine which parts of the source sequence are relevant when generating the target sequence.

Cross-Attention Example

Suppose the encoder processes:

text
1I love AI

The encoder creates contextual representations:

text
1Encoder Output
2      │
3      ├── representation of "I"
4      ├── representation of "love"
5      └── representation of "AI"

When generating the target language, the decoder can use cross-attention to determine which source representations are relevant for the current output token.

Cross-Attention in PyTorch

python
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
10encoder_output = torch.randn(
11    2,
12    10,
13    512
14)
15
16decoder_hidden = torch.randn(
17    2,
18    6,
19    512
20)
21
22output, weights = cross_attention(
23    decoder_hidden,
24    encoder_output,
25    encoder_output
26)
27
28print(output.shape)

Output:

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

The output sequence length is determined by the decoder query sequence.

Feed-Forward Network

After attention operations, the decoder applies a position-wise feed-forward network.

A simplified version is:

text
1Input
2  │
3  ▼
4Linear
5  │
6  ▼
7Activation
8  │
9  ▼
10Dropout
11  │
12  ▼
13Linear
14  │
15  ▼
16Output

A PyTorch implementation is:

python
1import torch.nn as nn
2
3ffn = nn.Sequential(
4    nn.Linear(
5        512,
6        2048
7    ),
8
9    nn.GELU(),
10
11    nn.Dropout(0.1),
12
13    nn.Linear(
14        2048,
15        512
16    )
17)

The feed-forward network is applied independently to each sequence position while using shared parameters across positions.

Residual Connections

Residual connections add the input of a sub-layer to its output.

Conceptually:

text
1Input
2  │
3  ├──────────────┐
4  │              │
5  ▼              │
6Attention        │
7  │              │
8  └──────► Add ◄─┘
9             │
10             ▼
11          Output

Mathematically:

text
1Output = Input + SubLayer(Input)

Residual connections help information and gradients propagate through deep Transformer networks.

Layer Normalization

Layer normalization stabilizes the representations flowing through the network.

A simplified Transformer operation can be represented as:

text
1x
2│
3▼
4Attention
5│
6▼
7Dropout
8│
9▼
10Add x
11│
12▼
13LayerNorm

PyTorch provides:

python
1nn.LayerNorm(d_model)

For example:

python
1norm = nn.LayerNorm(512)
2
3x = torch.randn(
4    2,
5    6,
6    512
7)
8
9output = norm(x)
10
11print(output.shape)

Output:

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

Pre-Norm and Post-Norm

Transformer implementations can organize normalization differently.

Post-Norm

The original Transformer architecture is commonly described using:

text
1x → Sublayer → Add → LayerNorm

Pre-Norm

Many modern Transformer implementations use:

text
1x → LayerNorm → Sublayer → Add

Pre-normalization can provide more stable optimization for deep Transformer networks.

When implementing a decoder from scratch, it is important to understand which normalization arrangement your architecture is using.

Output Projection

The decoder produces hidden representations, not directly token IDs.

A linear layer maps each hidden representation to vocabulary logits.

text
1Decoder Hidden State
2        │
3        ▼
4    Linear Layer
5        │
6        ▼
7Vocabulary Logits
8        │
9        ▼
10Next Token

Suppose:

text
1d_model = 512
2vocabulary size = 10,000

Then the projection is:

python
1projection = nn.Linear(
2    512,
3    10000
4)

Example:

python
1x = torch.randn(
2    2,
3    8,
4    512
5)
6
7logits = projection(x)
8
9print(logits.shape)

Output:

text
1torch.Size([2, 8, 10000])

The dimensions mean:

text
12      → batch size
28      → target sequence length
310000  → vocabulary size

The values are logits, not token IDs.

From Logits to Tokens

The model produces vocabulary scores:

text
1Vocabulary
2
3hello     → 2.1
4world     → 0.4
5AI        → 4.7
6model     → 1.8
7...

A decoding strategy converts these scores into the next token.

The simplest strategy is greedy decoding:

python
1next_token = logits[:, -1].argmax(dim=-1)

More advanced generation methods include:

Temperature sampling
Top-k sampling
Top-p sampling
Beam search

Building a Transformer Decoder Block

Now combine masked self-attention, cross-attention, feed-forward layers, residual connections, and normalization.

python
1import torch
2import torch.nn as nn
3
4
5class DecoderBlock(nn.Module):
6
7    def __init__(
8        self,
9        d_model=512,
10        n_heads=8,
11        ff_dim=2048,
12        dropout=0.1
13    ):
14        super().__init__()
15
16        self.self_attention = nn.MultiheadAttention(
17            embed_dim=d_model,
18            num_heads=n_heads,
19            dropout=dropout,
20            batch_first=True
21        )
22
23        self.cross_attention = nn.MultiheadAttention(
24            embed_dim=d_model,
25            num_heads=n_heads,
26            dropout=dropout,
27            batch_first=True
28        )
29
30        self.ffn = nn.Sequential(
31            nn.Linear(
32                d_model,
33                ff_dim
34            ),
35
36            nn.GELU(),
37
38            nn.Dropout(dropout),
39
40            nn.Linear(
41                ff_dim,
42                d_model
43            )
44        )
45
46        self.norm1 = nn.LayerNorm(d_model)
47        self.norm2 = nn.LayerNorm(d_model)
48        self.norm3 = nn.LayerNorm(d_model)
49
50        self.dropout = nn.Dropout(dropout)
51
52    def forward(
53        self,
54        x,
55        encoder_output,
56        tgt_mask=None,
57        memory_key_padding_mask=None
58    ):
59
60        self_attn, _ = self.self_attention(
61            x,
62            x,
63            x,
64            attn_mask=tgt_mask
65        )
66
67        x = self.norm1(
68            x + self.dropout(self_attn)
69        )
70
71        cross_attn, _ = self.cross_attention(
72            x,
73            encoder_output,
74            encoder_output,
75            key_padding_mask=memory_key_padding_mask
76        )
77
78        x = self.norm2(
79            x + self.dropout(cross_attn)
80        )
81
82        ffn_output = self.ffn(x)
83
84        x = self.norm3(
85            x + self.dropout(ffn_output)
86        )
87
88        return x

This implementation demonstrates the core structure of an encoder-decoder Transformer decoder block.

Testing the Decoder Block

Create example encoder and decoder tensors:

python
1decoder = DecoderBlock()
2
3encoder_output = torch.randn(
4    2,
5    10,
6    512
7)
8
9decoder_input = torch.randn(
10    2,
11    6,
12    512
13)
14
15mask = torch.triu(
16    torch.ones(
17        6,
18        6
19    ),
20    diagonal=1
21).bool()
22
23output = decoder(
24    decoder_input,
25    encoder_output,
26    tgt_mask=mask
27)
28
29print(output.shape)

Output:

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

The decoder preserves the target sequence length and model dimension.

Building Multiple Decoder Layers

A Transformer decoder usually contains several decoder blocks.

python
1class TransformerDecoder(nn.Module):
2
3    def __init__(
4        self,
5        num_layers=6,
6        d_model=512,
7        n_heads=8,
8        ff_dim=2048,
9        dropout=0.1
10    ):
11        super().__init__()
12
13        self.layers = nn.ModuleList([
14            DecoderBlock(
15                d_model=d_model,
16                n_heads=n_heads,
17                ff_dim=ff_dim,
18                dropout=dropout
19            )
20            for _ in range(num_layers)
21        ])
22
23    def forward(
24        self,
25        x,
26        encoder_output,
27        tgt_mask=None,
28        memory_key_padding_mask=None
29    ):
30
31        for layer in self.layers:
32
33            x = layer(
34                x,
35                encoder_output,
36                tgt_mask=tgt_mask,
37                memory_key_padding_mask=memory_key_padding_mask
38            )
39
40        return x

The input passes through each decoder block sequentially:

text
1Decoder Input
2      ↓
3Layer 1
4      ↓
5Layer 2
6      ↓
7Layer 3
8      ↓
9...
10      ↓
11Layer 6
12      ↓
13Decoder Output

Testing the Transformer Decoder

python
1decoder = TransformerDecoder()
2
3encoder_output = torch.randn(
4    2,
5    10,
6    512
7)
8
9decoder_input = torch.randn(
10    2,
11    6,
12    512
13)
14
15mask = torch.triu(
16    torch.ones(
17        6,
18        6
19    ),
20    diagonal=1
21).bool()
22
23output = decoder(
24    decoder_input,
25    encoder_output,
26    tgt_mask=mask
27)
28
29print(output.shape)

Output:

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

Using PyTorch's TransformerDecoder

PyTorch already provides a Transformer decoder implementation.

python
1import torch
2import torch.nn as nn
3
4decoder_layer = nn.TransformerDecoderLayer(
5    d_model=512,
6    nhead=8,
7    dim_feedforward=2048,
8    dropout=0.1,
9    batch_first=True
10)
11
12decoder = nn.TransformerDecoder(
13    decoder_layer,
14    num_layers=6
15)

Create example inputs:

python
1memory = torch.randn(
2    2,
3    10,
4    512
5)
6
7tgt = torch.randn(
8    2,
9    6,
10    512
11)

Create a causal mask:

python
1mask = torch.triu(
2    torch.ones(
3        6,
4        6
5    ),
6    diagonal=1
7).bool()

Run the decoder:

python
1output = decoder(
2    tgt,
3    memory,
4    tgt_mask=mask
5)
6
7print(output.shape)

Output:

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

Here:

text
1tgt
2↓
3Decoder input
4
5memory
6↓
7Encoder output

The decoder uses masked self-attention over tgt and cross-attention over memory.

Decoder Padding Masks

Causal masking and padding masking solve different problems.

Causal Mask

Prevents a position from seeing future target positions.

text
1Future information
2       ↓
3     BLOCKED

Padding Mask

Prevents attention to padding tokens.

text
1<PAD>
2<PAD>
3  ↓
4Ignored by attention

A real sequence-to-sequence implementation may need both.

For example:

python
1tgt_padding_mask = (
2    tgt_tokens == pad_token_id
3)
4
5memory_padding_mask = (
6    src_tokens == pad_token_id
7)

The exact argument names depend on the PyTorch module being used.

Decoder Training With Teacher Forcing

During training, sequence-to-sequence models commonly use teacher forcing.

Suppose the target sentence is:

text
1<SOS> I love AI <EOS>

The decoder input becomes:

text
1<SOS> I love AI

The expected labels become:

text
1I love AI <EOS>

The model therefore learns:

text
1<SOS>
2   ↓
3I
4
5<SOS> I
6   ↓
7love
8
9<SOS> I love
10   ↓
11AI
12
13<SOS> I love AI
14   ↓
15<EOS>

This shifted-target setup is essential for standard autoregressive training.

Decoder Loss

Suppose the decoder produces:

text
1(batch, target_length, vocabulary_size)

Cross-entropy loss can be calculated after flattening the batch and sequence dimensions.

python
1criterion = nn.CrossEntropyLoss(
2    ignore_index=pad_token_id
3)
4
5loss = criterion(
6    logits.reshape(
7        -1,
8        vocabulary_size
9    ),
10    target.reshape(-1)
11)

The padding index is ignored so padded positions do not contribute to the loss.

Autoregressive Generation

During inference, the target sequence is not known in advance.

The decoder generates one token at a time.

text
1<SOS>
2  │
3  ▼
4Decoder
5  │
6  ▼
7Token 1
8  │
9  ▼
10Decoder
11  │
12  ▼
13Token 2
14  │
15  ▼
16Decoder
17  │
18  ▼
19Token 3

Generation continues until:

text
1<EOS>

or until the maximum generation length is reached.

Greedy Decoding

A simple decoding algorithm selects the token with the highest logit.

python
1next_token = logits[:, -1].argmax(
2    dim=-1
3)

This is easy to implement but does not always produce the best sequence.

For advanced generation, consider:

Beam search
Top-k sampling
Top-p sampling
Temperature scaling

KV Caching During Generation

Autoregressive generation can become expensive because the model repeatedly processes previous tokens.

Modern decoder-only Transformer implementations commonly use key-value caching.

Conceptually:

text
1Previous K/V
2     │
3     ▼
4Cached
5     │
6     +
7New Token
8     │
9     ▼
10New Attention Computation

Instead of recomputing all previous key and value representations at every generation step, cached values can be reused.

KV caching is an important optimization for large language model inference.

Decoder-Only Transformer Architecture

For a GPT- or LLaMA-style decoder-only model, the architecture is simplified because there is no encoder or cross-attention.

text
1Token IDs
2    │
3    ▼
4Token Embeddings
5    │
6    +
7Position Information
8    │
9    ▼
10Masked Self-Attention
11    │
12    ▼
13Feed-Forward Network
14    │
15    ▼
16Repeated Decoder Blocks
17    │
18    ▼
19Final Normalization
20    │
21    ▼
22Language Modeling Head
23    │
24    ▼
25Vocabulary Logits

This distinction is important when studying modern large language models.

Encoder-Decoder Transformer Architecture

For translation and other sequence-to-sequence tasks:

text
1                Source Tokens
2                     │
3                     ▼
4              Transformer Encoder
5                     │
6                     ▼
7                Encoder Memory
8                     │
9                     ▼
10Target Tokens → Transformer Decoder
11                     │
12                     ▼
13               Vocabulary Logits
14                     │
15                     ▼
16                Target Tokens

The decoder uses:

text
1Masked Self-Attention
2        +
3Cross-Attention
4        +
5Feed-Forward Network

Decoder vs Encoder

Feature	Transformer Encoder	Transformer Decoder
Self-attention	Yes	Yes
Causal masking	Usually no	Yes for autoregressive decoding
Cross-attention	No	Yes in encoder-decoder architectures
Uses encoder memory	No	Yes when applicable
Generates tokens	No	Yes
Typical role	Input representation	Sequence generation

The exact behavior depends on the model architecture and task.

Common Transformer Decoder Mistakes

Forgetting the Causal Mask

Without causal masking during autoregressive training, the decoder can access future target positions.

Confusing Causal and Padding Masks

They perform different functions.

text
1Causal mask
2→ blocks future positions
3
4Padding mask
5→ blocks padding tokens

Using Incorrect Target Shifting

The decoder input and expected target should normally be shifted during teacher-forced training.

Ignoring Padding in the Loss

Padding tokens should normally be excluded from the training loss.

Forgetting Cross-Attention in Encoder-Decoder Models

An encoder-decoder Transformer requires the decoder to attend to encoder representations.

Confusing Logits With Probabilities

The output projection normally produces logits.

Do not apply softmax before passing logits to CrossEntropyLoss, because CrossEntropyLoss internally combines the required operations.

Ignoring Tensor Shapes

Common shapes include:

text
1Input tokens:
2(batch, sequence_length)
3
4Embeddings:
5(batch, sequence_length, d_model)
6
7Decoder output:
8(batch, target_length, d_model)
9
10Vocabulary logits:
11(batch, target_length, vocabulary_size)

Incorrect dimensions are one of the most common sources of Transformer implementation errors.

Best Practices for Transformer Decoders

Use causal masking for autoregressive target generation.
Keep causal masking separate from padding masking.
Use shifted target sequences during teacher-forced training.
Ignore padding tokens in the loss.
Use residual connections around decoder sub-layers.
Use LayerNorm consistently with the chosen architecture.
Use cross-attention when implementing an encoder-decoder Transformer.
Verify tensor shapes at every major stage.
Use model.train() during training.
Use model.eval() during evaluation and inference.
Use torch.no_grad() when gradients are unnecessary during inference.
Consider KV caching for efficient autoregressive generation.
Start with a small model when debugging custom implementations.
Use PyTorch's optimized Transformer components when a custom implementation is not required.

Practice Project: Implement a Transformer Decoder

Build a small decoder with:

text
1d_model = 512
2attention_heads = 8
3feed_forward_dimension = 2048
4decoder_layers = 6

Step 1: Create Encoder Output

python
1encoder_output = torch.randn(
2    4,
3    15,
4    512
5)

Step 2: Create Decoder Input

python
1decoder_input = torch.randn(
2    4,
3    10,
4    512
5)

Step 3: Create Causal Mask

python
1mask = torch.triu(
2    torch.ones(
3        10,
4        10
5    ),
6    diagonal=1
7).bool()

Step 4: Run the Decoder

python
1decoder = TransformerDecoder()
2
3output = decoder(
4    decoder_input,
5    encoder_output,
6    tgt_mask=mask
7)
8
9print(output.shape)

Expected output:

text
1torch.Size([4, 10, 512])

Step 5: Project to the Vocabulary

python
1projection = nn.Linear(
2    512,
3    30000
4)
5
6logits = projection(output)
7
8print(logits.shape)

Expected output:

text
1torch.Size([4, 10, 30000])

This means the model produces 30,000 vocabulary scores for every one of the 10 target positions in each of the 4 examples.

Transformer Decoder Learning Checklist

Before moving to advanced Transformer architectures, make sure you can explain:

What a decoder does
Why autoregressive models need causal masking
How masked self-attention works
How query, key, and value differ
How cross-attention connects an encoder and decoder
Why residual connections are used
What LayerNorm does
How the feed-forward network works
How hidden states become vocabulary logits
Why decoder targets are shifted during training
Why padding should be ignored in the loss
How greedy decoding works
Why KV caching improves autoregressive inference
The difference between decoder-only and encoder-decoder Transformers

Module Summary

A Transformer decoder is responsible for generating target representations and, ultimately, target tokens.

A decoder block can be summarized as:

text
1Decoder Input
2      │
3      ▼
4Masked Self-Attention
5      │
6      ▼
7Residual + LayerNorm
8      │
9      ▼
10Cross-Attention
11      │
12      ▼
13Residual + LayerNorm
14      │
15      ▼
16Feed-Forward Network
17      │
18      ▼
19Residual + LayerNorm
20      │
21      ▼
22Decoder Output

For decoder-only models, the cross-attention stage is removed:

text
1Token Embedding
2      │
3      ▼
4Causal Self-Attention
5      │
6      ▼
7Feed-Forward Network
8      │
9      ▼
10Repeated Decoder Blocks
11      │
12      ▼
13Language Modeling Head
14      │
15      ▼
16Next-Token Prediction

The most important concept is causal attention: during autoregressive generation, the decoder must not use future target tokens to predict the current token.