Transformer APIs

PyTorch Transformer APIs: Complete Guide to Encoder, Decoder, Attention, and Translation

PyTorch provides a collection of high-level Transformer APIs that make it possible to build encoder-decoder architectures without implementing every Transformer component from scratch.

In this module, you will learn how PyTorch's built-in Transformer components work, how they fit together, and how to use them to construct a basic text translation model.

The main APIs covered are:

• nn.Transformer
• nn.TransformerEncoder
• nn.TransformerEncoderLayer
• nn.TransformerDecoder
• nn.TransformerDecoderLayer
• nn.MultiheadAttention
• nn.Embedding
• nn.LayerNorm
• nn.Dropout

You will also build a complete Transformer-based translation model and train it using a basic PyTorch training loop.

What You Will Learn

By the end of this module, you should be able to:

• Understand the architecture of PyTorch's Transformer APIs
• Create an encoder-decoder Transformer
• Build individual encoder and decoder layers
• Use multi-head attention
• Convert token IDs into embeddings
• Apply layer normalization and dropout
• Understand Transformer tensor shapes
• Build a basic sequence-to-sequence translation model
• Calculate a language-model-style cross-entropy loss
• Train the model using an optimizer and backpropagation
• Understand how the individual components combine into a complete Transformer pipeline

Why Use PyTorch Transformer APIs?

Implementing a complete Transformer manually requires attention mechanisms, feed-forward networks, normalization, residual connections, masking, parameter initialization, and other components.

PyTorch provides reusable Transformer modules so developers can focus on model architecture and experimentation instead of rebuilding standard components from scratch.

The main advantages include:

• Less implementation code
• Faster model development
• Reusable components
• Easier experimentation
• Integrated PyTorch autograd support
• CPU and GPU execution
• Straightforward integration with other PyTorch modules

These APIs are especially useful for learning Transformer architecture and building custom encoder-decoder models.

PyTorch Transformer API Architecture

The main components can be understood as a hierarchy:

textCopy
1                    nn.Transformer
2                          │
3             ┌────────────┴────────────┐
4             │                         │
5             ▼                         ▼
6    TransformerEncoder       TransformerDecoder
7             │                         │
8             ▼                         ▼
9 TransformerEncoderLayer      TransformerDecoderLayer
10             │                         │
11             └────────────┬────────────┘
12                          ▼
13                 MultiheadAttention
14                          │
15                          ▼
16                  Feed-Forward Network
17                          │
18                          ▼
19                     LayerNorm
20                          │
21                          ▼
22                       Dropout

At a high level:

• nn.Transformer combines encoder and decoder stacks.
• nn.TransformerEncoder contains multiple encoder layers.
• nn.TransformerDecoder contains multiple decoder layers.
• nn.MultiheadAttention implements the attention mechanism.
• nn.Embedding converts token IDs into dense vectors.
• nn.LayerNorm normalizes features.
• nn.Dropout provides regularization.

Understanding Tensor Shapes

Transformer models process sequences as tensors.

When using:

pythonCopy
1batch_first=True

the common tensor format becomes:

textCopy
1(batch_size, sequence_length, embedding_dimension)

For example:

textCopy
1(32, 128, 512)

means:

• 32 examples in the batch
• 128 tokens per sequence
• 512 features per token

Using batch_first=True makes Transformer code easier to read because the batch dimension comes first.

1. nn.Transformer

nn.Transformer provides a complete encoder-decoder Transformer architecture.

It combines:

textCopy
1Source Sequence
2      ↓
3Embedding
4      ↓
5Transformer Encoder
6      ↓
7Encoder Memory
8      ↓
9Transformer Decoder
10      ↓
11Output Representations

Creating a Transformer

pythonCopy
1import torch
2import torch.nn as nn
3
4transformer = nn.Transformer(
5    d_model=512,
6    nhead=8,
7    num_encoder_layers=6,
8    num_decoder_layers=6,
9    dim_feedforward=2048,
10    dropout=0.1,
11    batch_first=True
12)
13
14print(transformer)

Important parameters include:

For example, with d_model=512, every token entering the Transformer has a 512-dimensional representation.

Running a Forward Pass

Create sample source and target representations:

pythonCopy
1src = torch.randn(
2    2,
3    15,
4    512
5)
6
7tgt = torch.randn(
8    2,
9    10,
10    512
11)
12
13output = transformer(
14    src,
15    tgt
16)
17
18print(output.shape)

Output:

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

The output preserves the target sequence length and Transformer feature dimension.

2. TransformerEncoderLayer

nn.TransformerEncoderLayer represents one encoder block.

A simplified encoder block can be viewed as:

textCopy
1Input
2  ↓
3Self-Attention
4  ↓
5Residual Connection + Normalization
6  ↓
7Feed-Forward Network
8  ↓
9Residual Connection + Normalization
10  ↓
11Output

The self-attention mechanism allows each token to interact with other tokens in the sequence.

Creating an Encoder Layer

pythonCopy
1encoder_layer = nn.TransformerEncoderLayer(
2    d_model=512,
3    nhead=8,
4    dim_feedforward=2048,
5    dropout=0.1,
6    batch_first=True
7)
8
9print(encoder_layer)

Testing the Encoder Layer

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

Output:

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

The encoder layer does not change the sequence length or feature dimension.

3. TransformerEncoder

nn.TransformerEncoder stacks multiple encoder layers.

Conceptually:

textCopy
1Input
2  ↓
3Encoder Layer 1
4  ↓
5Encoder Layer 2
6  ↓
7Encoder Layer 3
8  ↓
9...
10  ↓
11Encoder Layer N
12  ↓
13Output

Creating an Encoder

pythonCopy
1encoder = nn.TransformerEncoder(
2    encoder_layer,
3    num_layers=6
4)
5
6x = torch.randn(
7    2,
8    30,
9    512
10)
11
12output = encoder(x)
13
14print(output.shape)

Output:

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

Increasing num_layers increases the depth of the encoder and therefore the capacity of the network.

4. TransformerDecoderLayer

nn.TransformerDecoderLayer represents one decoder block.

A simplified decoder contains:

textCopy
1Target Sequence
2      ↓
3Masked Self-Attention
4      ↓
5Cross-Attention
6      ↓
7Feed-Forward Network
8      ↓
9Output

The important difference from the encoder is that the decoder performs attention over the target sequence and also attends to the encoder's output.

Creating a Decoder Layer

pythonCopy
1decoder_layer = nn.TransformerDecoderLayer(
2    d_model=512,
3    nhead=8,
4    dim_feedforward=2048,
5    dropout=0.1,
6    batch_first=True
7)
8
9print(decoder_layer)

Testing the Decoder Layer

The decoder receives:

• tgt: target sequence representations
• memory: encoder output

pythonCopy
1memory = torch.randn(
2    2,
3    15,
4    512
5)
6
7tgt = torch.randn(
8    2,
9    8,
10    512
11)
12
13output = decoder_layer(
14    tgt,
15    memory
16)
17
18print(output.shape)

Output:

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

The output follows the target sequence length.

5. TransformerDecoder

nn.TransformerDecoder stacks multiple decoder layers.

pythonCopy
1decoder = nn.TransformerDecoder(
2    decoder_layer,
3    num_layers=6
4)
5
6output = decoder(
7    tgt,
8    memory
9)
10
11print(output.shape)

Output:

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

The decoder uses the encoder output as its memory representation.

6. MultiheadAttention

nn.MultiheadAttention implements the core attention mechanism used by Transformer architectures.

Attention allows a token to assign different importance to other tokens in the sequence.

Creating Multi-Head Attention

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

Output:

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

The three inputs represent:

textCopy
1Query
2Key
3Value

When the same tensor is passed three times:

pythonCopy
1attention(x, x, x)

the operation represents self-attention.

Understanding Attention Weights

Attention weights describe how strongly tokens attend to other tokens.

Conceptually:

textCopy
1              Token 1   Token 2   Token 3
2Token 1          0.72      0.18      0.10
3Token 2          0.20      0.65      0.15
4Token 3          0.11      0.24      0.65

The values depend on the input and learned parameters. They should be interpreted as attention scores rather than simple statements of semantic importance.

7. Embedding Layer

Transformer models normally receive token IDs rather than raw words.

nn.Embedding maps each token ID to a learned dense vector.

Conceptually:

textCopy
1Token ID
2   ↓
3Embedding Table
4   ↓
5Dense Vector

Creating an Embedding Layer

pythonCopy
1embedding = nn.Embedding(
2    num_embeddings=10000,
3    embedding_dim=512
4)
5
6tokens = torch.randint(
7    0,
8    10000,
9    (4, 20)
10)
11
12vectors = embedding(tokens)
13
14print(vectors.shape)

Output:

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

Here:

• Vocabulary size = 10000
• Embedding dimension = 512
• Batch size = 4
• Sequence length = 20

Each token is converted from an integer ID into a 512-dimensional vector.

8. LayerNorm

nn.LayerNorm performs normalization over the specified feature dimensions.

Transformer architectures commonly use layer normalization around their attention and feed-forward components.

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

Output:

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

The shape remains unchanged because normalization operates on the feature dimension.

9. Dropout

Dropout is a regularization technique used during training.

pythonCopy
1dropout = nn.Dropout(0.2)
2
3x = torch.ones(
4    2,
5    5
6)
7
8output = dropout(x)
9
10print(output)

During training, dropout randomly disables elements according to the specified probability.

During evaluation, dropout is disabled.

For example:

pythonCopy
1model.train()

enables training behavior, while:

pythonCopy
1model.eval()

switches modules such as dropout into evaluation behavior.

Building a Transformer Translation Model

Now combine the components into a simple sequence-to-sequence translation model.

The model will contain:

textCopy
1Source Token IDs
2       ↓
3Source Embedding
4       ↓
5Transformer Encoder
6       ↓
7Transformer Decoder
8       ↑
9Target Embedding
10       ↓
11Linear Projection
12       ↓
13Target Vocabulary Scores

Model Definition

pythonCopy
1import torch
2import torch.nn as nn
3
4
5class Translator(nn.Module):
6
7    def __init__(
8        self,
9        src_vocab,
10        tgt_vocab,
11        d_model=512
12    ):
13        super().__init__()
14
15        self.src_embedding = nn.Embedding(
16            src_vocab,
17            d_model
18        )
19
20        self.tgt_embedding = nn.Embedding(
21            tgt_vocab,
22            d_model
23        )
24
25        self.transformer = nn.Transformer(
26            d_model=d_model,
27            nhead=8,
28            num_encoder_layers=6,
29            num_decoder_layers=6,
30            batch_first=True
31        )
32
33        self.fc = nn.Linear(
34            d_model,
35            tgt_vocab
36        )
37
38    def forward(
39        self,
40        src,
41        tgt
42    ):
43        src = self.src_embedding(src)
44        tgt = self.tgt_embedding(tgt)
45
46        output = self.transformer(
47            src,
48            tgt
49        )
50
51        return self.fc(output)

Understanding the Forward Pass

The source token IDs first pass through the source embedding:

pythonCopy
1src = self.src_embedding(src)

The target token IDs are also converted into vectors:

pythonCopy
1tgt = self.tgt_embedding(tgt)

The resulting representations are passed into the Transformer:

pythonCopy
1output = self.transformer(
2    src,
3    tgt
4)

Finally, a linear layer projects each Transformer representation into the target vocabulary:

pythonCopy
1return self.fc(output)

The result is a set of logits for every target position.

Creating the Model

pythonCopy
1model = Translator(
2    src_vocab=10000,
3    tgt_vocab=12000
4)
5
6print(model)

The model has:

textCopy
1Source vocabulary = 10,000
2Target vocabulary = 12,000
3Model dimension = 512
4Attention heads = 8
5Encoder layers = 6
6Decoder layers = 6

Creating Example Input Data

For demonstration, generate random token IDs:

pythonCopy
1src = torch.randint(
2    0,
3    10000,
4    (8, 20)
5)
6
7tgt = torch.randint(
8    0,
9    12000,
10    (8, 15)
11)

The shapes are:

textCopy
1src → (8, 20)
2tgt → (8, 15)

where:

• 8 is the batch size
• 20 is the source sequence length
• 15 is the target sequence length

This is synthetic data and is useful only for demonstrating tensor shapes and model execution. It is not sufficient to train a meaningful translation system.

Running the Forward Pass

pythonCopy
1output = model(
2    src,
3    tgt
4)
5
6print(output.shape)

Output:

textCopy
1torch.Size([8, 15, 12000])

The output dimensions represent:

textCopy
18      → batch size
215     → target sequence length
312000  → target vocabulary size

Therefore, every target position receives a vector of 12,000 vocabulary logits.

Computing the Loss

For classification over vocabulary tokens, cross-entropy loss can be used.

pythonCopy
1criterion = nn.CrossEntropyLoss()
2
3loss = criterion(
4    output.reshape(
5        -1,
6        12000
7    ),
8    tgt.reshape(-1)
9)
10
11print(loss.item())

The reshape changes:

textCopy
1(batch, sequence, vocabulary)

into:

textCopy
1(batch × sequence, vocabulary)

The target is flattened into:

textCopy
1(batch × sequence)

This allows CrossEntropyLoss to calculate a classification loss for every target position.

Important Note About Decoder Training

A real sequence-to-sequence translation system normally uses shifted target sequences.

For example:

textCopy
1Decoder Input:
2<BOS> I am learning
3
4Expected Output:
5I am learning <EOS>

The decoder should not simply receive the same target sequence that it is expected to predict. Training also requires appropriate causal masking so that a position cannot use future target tokens.

This distinction is important when moving from a demonstration model to a real translation system.

Optimizer

Use an optimizer to update the model parameters.

pythonCopy
1optimizer = torch.optim.Adam(
2    model.parameters(),
3    lr=0.0001
4)

The optimizer uses gradients calculated during backpropagation to update the model's trainable parameters.

Training Loop

A basic training loop can be written as:

pythonCopy
1epochs = 5
2
3for epoch in range(epochs):
4
5    model.train()
6
7    optimizer.zero_grad()
8
9    output = model(
10        src,
11        tgt
12    )
13
14    loss = criterion(
15        output.reshape(
16            -1,
17            12000
18        ),
19        tgt.reshape(-1)
20    )
21
22    loss.backward()
23
24    optimizer.step()
25
26    print(
27        f"Epoch {epoch + 1}: {loss.item()}"
28    )

The training process follows:

textCopy
1Forward Pass
2     ↓
3Calculate Loss
4     ↓
5Backpropagation
6     ↓
7Gradient Calculation
8     ↓
9Optimizer Step
10     ↓
11Updated Parameters

For a real translation project, you would additionally need a proper dataset, tokenizer, target shifting, masking, batching, padding handling, validation, and an appropriate evaluation metric.

Complete Transformer Translation Pipeline

The complete conceptual pipeline is:

textCopy
1Source Sentence
2      │
3      ▼
4Tokenizer
5      │
6      ▼
7Source Token IDs
8      │
9      ▼
10Source Embedding
11      │
12      ▼
13Transformer Encoder
14      │
15      ▼
16Encoder Memory
17      │
18      ▼
19Transformer Decoder
20      ▲
21      │
22Target Token IDs
23      │
24      ▼
25Target Embedding
26      │
27      ▼
28Linear Projection
29      │
30      ▼
31Vocabulary Logits
32      │
33      ▼
34Token Selection
35      │
36      ▼
37Target Sentence

This architecture is the foundation of many encoder-decoder sequence-to-sequence models.

Practice Project: Build a Text Translation Model

A good practice exercise is to extend the basic Translator class into a more complete translation system.

Step 1: Create the Model

pythonCopy
1model = Translator(
2    src_vocab=5000,
3    tgt_vocab=7000
4)

Step 2: Generate Sample Data

pythonCopy
1src = torch.randint(
2    0,
3    5000,
4    (16, 25)
5)
6
7tgt = torch.randint(
8    0,
9    7000,
10    (16, 18)
11)

Step 3: Run the Model

pythonCopy
1output = model(
2    src,
3    tgt
4)
5
6print(output.shape)

Output:

textCopy
1torch.Size([16, 18, 7000])

Step 4: Calculate Loss

pythonCopy
1criterion = nn.CrossEntropyLoss()
2
3loss = criterion(
4    output.reshape(
5        -1,
6        7000
7    ),
8    tgt.reshape(-1)
9)
10
11print(loss)

Step 5: Improve the Project

After the basic model works, add:

• A real parallel translation dataset
• A tokenizer
• Padding tokens
• Beginning-of-sequence tokens
• End-of-sequence tokens
• Padding masks
• Causal masks
• Teacher forcing
• Validation data
• Evaluation metrics
• GPU training
• Checkpoint saving
• Inference code

This transforms the exercise from a tensor demonstration into a meaningful sequence-to-sequence project.

PyTorch Transformer API Summary

Encoder vs Decoder

Understanding the difference between encoder and decoder architectures is essential.

A Transformer encoder builds contextual representations of an input sequence, while a decoder can generate an output sequence while attending to encoder representations.

Transformer Architectures in Modern Models

Transformer components appear in many influential architectures.

The exact implementations and architectural details differ between these models, but the Transformer family provides the underlying attention-based design.

Best Practices for PyTorch Transformers

Use batch_first=True When Appropriate

Using:

pythonCopy
1batch_first=True

makes the expected shape:

textCopy
1(batch_size, sequence_length, d_model)

This is often easier to work with in application code.

Use Positional Information

Self-attention alone does not inherently provide sequence order.

A practical Transformer therefore needs positional information, such as positional embeddings or another positional encoding mechanism.

The exact approach depends on the architecture being implemented.

Use Attention Masks Correctly

For sequence-to-sequence models, distinguish between:

• Padding masks
• Causal masks

Padding masks prevent attention to padding tokens.

Causal masks prevent autoregressive decoding positions from accessing future target tokens.

Use Appropriate Normalization and Regularization

Layer normalization is a fundamental part of Transformer blocks, while dropout can help regularize models during training.

Switch Between Training and Evaluation Modes

Use:

pythonCopy
1model.train()

during training and:

pythonCopy
1model.eval()

during evaluation or inference.

When evaluating without gradient computation, also consider:

pythonCopy
1with torch.no_grad():
2    ...

Monitor Tensor Shapes

Transformer bugs frequently originate from incorrect dimensions.

Before debugging model behavior, verify:

textCopy
1batch size
2sequence length
3embedding dimension
4vocabulary size

at every major stage.

Common Beginner Mistakes

Forgetting Positional Information

Embeddings alone do not encode the position of a token in the sequence.

Using Incorrect Target Alignment

For autoregressive training, decoder inputs and expected outputs normally need to be shifted relative to each other.

Ignoring Masks

Incorrect masks can allow the decoder to access information that should not be available during autoregressive prediction.

Confusing Logits with Probabilities

The output of the final linear layer is normally logits.

CrossEntropyLoss expects logits and internally handles the appropriate normalization.

Using Random Data as a Real Dataset

Random token IDs are useful for testing tensor shapes and verifying that code executes, but they do not contain linguistic relationships and therefore cannot produce meaningful translation behavior.

Debugging Transformer Models

When a Transformer model fails, check the problem systematically.

Check the Input

Verify:

textCopy
1dtype
2device
3shape
4vocabulary range

Check the Embedding

Make sure every token ID satisfies:

textCopy
10 <= token_id < vocabulary_size

Check the Model Dimension

If:

pythonCopy
1d_model=512

the attention and Transformer components must receive representations with the expected feature dimension.

Check Attention Heads

The model dimension must be compatible with the number of attention heads.

For example:

textCopy
1d_model = 512
2nhead = 8

allows the representation to be divided across the attention heads.

Check Output Shape

For a target sequence of length T and vocabulary size V, the output is typically:

textCopy
1(batch_size, T, V)

This should match the expected loss calculation after reshaping.

Module Summary

In this module, you learned how PyTorch's built-in Transformer APIs can be combined to construct Transformer-based sequence-to-sequence models.

You learned:

• How nn.Transformer provides an encoder-decoder architecture
• How TransformerEncoder and TransformerEncoderLayer work
• How TransformerDecoder and TransformerDecoderLayer work
• How MultiheadAttention implements attention
• How nn.Embedding converts token IDs into vector representations
• How LayerNorm and Dropout are used in Transformer models
• How to construct a basic translation model
• How Transformer outputs are converted into vocabulary logits
• How cross-entropy loss can be calculated
• How to build a basic training loop
• Why masking, positional information, and target shifting matter in real Transformer training

The key idea is to understand the Transformer as a collection of reusable components rather than as a single black-box model.

textCopy
1Embedding
2    ↓
3Encoder
4    ↓
5Memory
6    ↓
7Decoder
8    ↓
9Linear Projection
10    ↓
11Vocabulary Logits

Once these components are understood, you can move from PyTorch's high-level APIs toward more advanced Transformer architectures and eventually implement custom attention mechanisms, modern language-model architectures, and efficient training pipelines.