Build a Transformer

Build a Transformer From Scratch With PyTorch: Encoder, Decoder, Masking, Training, and Translation

A Transformer is a neural network architecture designed to process sequences using attention mechanisms instead of relying primarily on recurrent computation.

In this module, you will build a complete encoder-decoder Transformer using PyTorch. Rather than treating nn.Transformer as a black box, you will construct the important surrounding components yourself, including token embeddings, positional embeddings, padding masks, causal masks, the output projection layer, training pipeline, and greedy decoding.

The final project demonstrates how these components can be combined to create the foundation of an English-to-French translation model.

What You Will Learn

By completing this tutorial, you will learn how to:

• Build a Transformer model with PyTorch
• Create token embeddings
• Add positional information to token representations
• Generate padding masks
• Generate causal attention masks
• Connect an encoder and decoder
• Project decoder representations into vocabulary logits
• Calculate cross-entropy loss
• Train a Transformer with backpropagation
• Perform autoregressive greedy decoding
• Understand the complete sequence-to-sequence pipeline
• Build the foundation of an English-to-French translation system

Transformer Architecture Overview

A sequence-to-sequence Transformer can be represented as:

textCopy
1Source Sentence
2      │
3      ▼
4Tokenization
5      │
6      ▼
7Token IDs
8      │
9      ▼
10Token Embedding
11      │
12      +
13      │
14Position Embedding
15      │
16      ▼
17Transformer Encoder
18      │
19      ▼
20Contextual Representations
21      │
22      ▼
23Transformer Decoder
24      ▲
25      │
26Target Token Embeddings
27      │
28      ▼
29Linear Output Layer
30      │
31      ▼
32Target Vocabulary Logits
33      │
34      ▼
35Generated Sentence

The encoder processes the source sequence and produces contextual representations. The decoder then uses those representations while generating the target sequence.

Step 1: Import PyTorch Libraries

Start by importing the libraries required to construct the model.

pythonCopy
1import math
2
3import torch
4import torch.nn as nn
5import torch.optim as optim

The main components are:

• torch for tensor operations
• torch.nn for neural network modules
• torch.optim for optimization algorithms
• math for scaling token embeddings

Step 2: Build the Token Embedding Layer

A Transformer cannot directly process words represented as integer token IDs.

For example:

textCopy
1"hello" → 42
2"world" → 781

The integers are indices into an embedding table.

nn.Embedding converts each token ID into a dense vector.

Token Embedding Implementation

pythonCopy
1class TokenEmbedding(nn.Module):
2
3    def __init__(self, vocab_size, d_model):
4        super().__init__()
5
6        self.embedding = nn.Embedding(
7            vocab_size,
8            d_model
9        )
10
11    def forward(self, x):
12        return self.embedding(x) * math.sqrt(
13            self.embedding.embedding_dim
14        )

If:

textCopy
1vocab_size = 10,000
2d_model = 512

the embedding table contains 10,000 vectors, each with 512 features.

Why Scale the Embeddings?

The implementation multiplies the embedding output by:

textCopy
1√d_model

This follows the scaling used in the original Transformer architecture and helps establish an appropriate magnitude for combining token representations with positional information.

Step 3: Add Positional Embeddings

Self-attention does not inherently encode the order of tokens.

For example:

textCopy
1"dog bites man"

and:

textCopy
1"man bites dog"

contain the same words but have different meanings because their positions differ.

The model therefore needs positional information.

Position Embedding Implementation

pythonCopy
1class PositionEmbedding(nn.Module):
2
3    def __init__(
4        self,
5        max_len,
6        d_model
7    ):
8        super().__init__()
9
10        self.embedding = nn.Embedding(
11            max_len,
12            d_model
13        )
14
15    def forward(self, x):
16
17        batch_size, seq_len = x.shape
18
19        positions = torch.arange(
20            seq_len,
21            device=x.device
22        )
23
24        positions = positions.unsqueeze(0).expand(
25            batch_size,
26            seq_len
27        )
28
29        return self.embedding(positions)

For a sequence of length 10, the position IDs are:

textCopy
10 1 2 3 4 5 6 7 8 9

Each position receives a learned vector of size d_model.

Step 4: Combine Token and Position Embeddings

The Transformer input representation can be created by adding token and position embeddings.

pythonCopy
1class TransformerEmbedding(nn.Module):
2
3    def __init__(
4        self,
5        vocab_size,
6        d_model,
7        max_len
8    ):
9        super().__init__()
10
11        self.token = TokenEmbedding(
12            vocab_size,
13            d_model
14        )
15
16        self.position = PositionEmbedding(
17            max_len,
18            d_model
19        )
20
21        self.dropout = nn.Dropout(0.1)
22
23    def forward(self, x):
24
25        embedding = self.token(x)
26
27        embedding = embedding + self.position(x)
28
29        return self.dropout(embedding)

The resulting representation is:

textCopy
1Token Information
2       +
3Position Information
4       ↓
5Transformer Input

This allows the model to process both the identity of each token and its location in the sequence.

Step 5: Generate Attention Masks

Masks control which tokens the Transformer is allowed to attend to.

Two important masks are:

• Padding mask
• Causal mask

These masks serve different purposes.

Padding Mask

Padding is commonly required when sequences have different lengths.

For example:

textCopy
1Sequence A:
2[12, 45, 91, 0, 0]
3
4Sequence B:
5[21, 34, 78, 15, 62]

Here 0 represents the padding token.

The padding mask identifies those positions:

pythonCopy
1def create_padding_mask(
2    tokens,
3    pad_idx=0
4):
5    return tokens == pad_idx

For the first sequence:

textCopy
1[12, 45, 91, 0, 0]

the mask conceptually becomes:

textCopy
1[False, False, False, True, True]

The Transformer can then avoid attending to those padding positions.

Causal Mask

During autoregressive decoding, the model should not see future target tokens.

For example, when predicting:

textCopy
1"I"

the decoder should not have access to:

textCopy
1"love AI"

A causal mask creates this restriction.

pythonCopy
1def create_causal_mask(
2    seq_len,
3    device
4):
5
6    return torch.triu(
7        torch.ones(
8            seq_len,
9            seq_len,
10            device=device
11        ),
12        diagonal=1
13    ).bool()

Conceptually, the mask has the form:

textCopy
1False  True   True   True
2False  False  True   True
3False  False  False  True
4False  False  False  False

The upper-triangular portion prevents each position from accessing future positions.

Test the Masks

pythonCopy
1tokens = torch.tensor([
2    [1, 2, 3, 0, 0],
3    [5, 6, 7, 8, 9]
4])
5
6padding = create_padding_mask(tokens)
7
8causal = create_causal_mask(
9    5,
10    tokens.device
11)
12
13print(padding)
14print(causal)

Testing masks independently is useful because incorrect masking can silently produce incorrect training behavior.

Step 6: Build the Transformer Model

Now combine the embedding layers, Transformer encoder-decoder architecture, and output projection.

pythonCopy
1class Transformer(nn.Module):
2
3    def __init__(
4        self,
5        src_vocab,
6        tgt_vocab,
7        d_model=512,
8        nhead=8,
9        num_layers=6,
10        ff_dim=2048,
11        max_len=100
12    ):
13        super().__init__()
14
15        self.src_embedding = TransformerEmbedding(
16            src_vocab,
17            d_model,
18            max_len
19        )
20
21        self.tgt_embedding = TransformerEmbedding(
22            tgt_vocab,
23            d_model,
24            max_len
25        )
26
27        self.transformer = nn.Transformer(
28            d_model=d_model,
29            nhead=nhead,
30            num_encoder_layers=num_layers,
31            num_decoder_layers=num_layers,
32            dim_feedforward=ff_dim,
33            batch_first=True
34        )
35
36        self.output = nn.Linear(
37            d_model,
38            tgt_vocab
39        )
40
41    def forward(
42        self,
43        src,
44        tgt,
45        src_padding_mask=None,
46        tgt_padding_mask=None,
47        tgt_mask=None
48    ):
49
50        src = self.src_embedding(src)
51
52        tgt = self.tgt_embedding(tgt)
53
54        output = self.transformer(
55            src,
56            tgt,
57            tgt_mask=tgt_mask,
58            src_key_padding_mask=src_padding_mask,
59            tgt_key_padding_mask=tgt_padding_mask
60        )
61
62        return self.output(output)

Understanding the Model

The model contains three major stages:

textCopy
1Source Embedding
2       ↓
3Transformer Encoder
4       ↓
5Encoder Representation
6       ↓
7Transformer Decoder
8       ↓
9Linear Output Projection

The final linear layer converts each decoder representation into a vector whose size equals the target vocabulary.

Step 7: Create the Transformer Model

Define source and target vocabulary sizes:

pythonCopy
1SRC_VOCAB = 10000
2TGT_VOCAB = 12000
3
4model = Transformer(
5    SRC_VOCAB,
6    TGT_VOCAB
7)
8
9print(model)

The model configuration is approximately:

textCopy
1Model dimension       = 512
2Attention heads       = 8
3Encoder layers        = 6
4Decoder layers        = 6
5Feed-forward dimension = 2048

These values are reasonable for demonstrating the architecture, but they are not automatically optimal for every dataset or hardware configuration.

Step 8: Create a Dummy Dataset

For testing the model, generate random token IDs.

pythonCopy
1src = torch.randint(
2    1,
3    SRC_VOCAB,
4    (8, 15)
5)
6
7tgt = torch.randint(
8    1,
9    TGT_VOCAB,
10    (8, 12)
11)

The shapes are:

textCopy
1Source:
2(8, 15)
3
4Target:
5(8, 12)

This means:

textCopy
1Batch size = 8
2Source sequence length = 15
3Target sequence length = 12

The random dataset is useful for testing the model's tensor flow, but it cannot teach the model meaningful translation because the data contains no linguistic relationship.

Step 9: Run a Forward Pass

Create the required masks:

pythonCopy
1tgt_mask = create_causal_mask(
2    tgt.size(1),
3    tgt.device
4)
5
6src_padding_mask = create_padding_mask(src)
7
8tgt_padding_mask = create_padding_mask(tgt)

Then run the Transformer:

pythonCopy
1outputs = model(
2    src,
3    tgt,
4    src_padding_mask,
5    tgt_padding_mask,
6    tgt_mask
7)
8
9print(outputs.shape)

Output:

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

The dimensions mean:

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

Therefore, the model produces vocabulary logits for every target position.

Step 10: Define the Loss Function

Use cross-entropy loss for token prediction.

pythonCopy
1criterion = nn.CrossEntropyLoss(
2    ignore_index=0
3)

The ignore_index=0 setting prevents padding tokens from contributing to the loss when 0 represents the padding token.

Step 11: Create the Optimizer

Use Adam to update the model parameters.

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

The optimizer uses gradients produced during backpropagation to adjust the trainable parameters.

Step 12: Train the Transformer

A basic training loop is:

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

The training process is:

textCopy
1Input Batch
2    ↓
3Embedding
4    ↓
5Transformer
6    ↓
7Vocabulary Logits
8    ↓
9Cross-Entropy Loss
10    ↓
11Backward Pass
12    ↓
13Gradients
14    ↓
15Optimizer
16    ↓
17Updated Parameters

Important Training Correction

For a real encoder-decoder translation model, the decoder input and prediction target should normally be shifted.

For example:

textCopy
1Decoder input:
2<SOS> je suis
3
4Expected prediction:
5je suis étudiant <EOS>

The decoder receives the previous target tokens while the loss compares predictions against the next tokens.

This is commonly called teacher forcing.

Using the same target tensor directly as both decoder input and loss target is acceptable as a simplified demonstration of tensor flow, but it is not the correct target-shifting strategy for a production translation system.

Step 13: Implement Greedy Decoding

During inference, the target sequence is generated one token at a time.

The basic process is:

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

A simple greedy decoder can be implemented as:

pythonCopy
1def greedy_decode(
2    model,
3    src,
4    max_len,
5    start_token
6):
7
8    model.eval()
9
10    generated = torch.tensor(
11        [[start_token]],
12        device=src.device
13    )
14
15    for _ in range(max_len):
16
17        mask = create_causal_mask(
18            generated.size(1),
19            generated.device
20        )
21
22        output = model(
23            src,
24            generated,
25            create_padding_mask(src),
26            create_padding_mask(generated),
27            mask
28        )
29
30        next_token = output[:, -1].argmax(-1)
31
32        generated = torch.cat(
33            [
34                generated,
35                next_token.unsqueeze(1)
36            ],
37            dim=1
38        )
39
40    return generated

Greedy decoding selects the token with the highest predicted logit at each step.

Greedy Decoding Limitation

Greedy decoding is simple but not always optimal.

Other generation strategies include:

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

For deterministic translation, beam search is often worth investigating.

Test the Decoder

pythonCopy
1prediction = greedy_decode(
2    model,
3    src[:1],
4    10,
5    start_token=1
6)
7
8print(prediction)

The result is a sequence of generated token IDs.

Because the example model has not been trained on a real translation dataset, the generated tokens will not represent meaningful French text.

Complete Transformer Training Pipeline

A practical translation system follows this general flow:

textCopy
1Dataset
2   │
3   ▼
4Tokenization
5   │
6   ▼
7Vocabulary
8   │
9   ▼
10Special Tokens
11   │
12   ▼
13Padding + Batching
14   │
15   ▼
16Token Embeddings
17   │
18   ▼
19Position Information
20   │
21   ▼
22Transformer Encoder
23   │
24   ▼
25Encoder Memory
26   │
27   ▼
28Transformer Decoder
29   │
30   ▼
31Output Projection
32   │
33   ▼
34Vocabulary Logits
35   │
36   ▼
37Cross-Entropy Loss
38   │
39   ▼
40Backpropagation
41   │
42   ▼
43Optimizer Update

During inference:

textCopy
1Source Sentence
2      ↓
3Tokenization
4      ↓
5Encoder
6      ↓
7Decoder
8      ↓
9Autoregressive Generation
10      ↓
11Generated Token IDs
12      ↓
13Detokenization
14      ↓
15Target Sentence

Practice Project: English-to-French Translator

Now extend the model into a small translation project.

Sample Training Data

A toy dataset could contain examples such as:

pythonCopy
1english = [
2    "i love ai",
3    "good morning",
4    "thank you",
5    "how are you"
6]
7
8french = [
9    "j aime ia",
10    "bonjour",
11    "merci",
12    "comment allez vous"
13]

This tiny dataset is useful for demonstrating preprocessing and training concepts, but it is far too small for a useful translation system.

Building a Vocabulary

A vocabulary maps tokens to integer IDs.

pythonCopy
1vocab = {
2    "<PAD>": 0,
3    "<SOS>": 1,
4    "<EOS>": 2,
5    "<UNK>": 3,
6    "i": 4,
7    "love": 5,
8    "ai": 6
9}

Important special tokens include:

A complete tokenizer must also create a target vocabulary and convert every sentence into token IDs.

Teacher Forcing in Translation

For sequence-to-sequence training, create shifted decoder inputs and labels.

Conceptually:

textCopy
1Target sentence:
2<SOS> j aime ia <EOS>
3
4Decoder input:
5<SOS> j aime ia
6
7Training target:
8j aime ia <EOS>

At each position, the model learns to predict the next token.

This prevents the model from simply receiving the token it is supposed to predict.

Training the Translation Model

The training loop should:

1. Load a batch.
2. Create shifted decoder inputs.
3. Generate padding masks.
4. Generate a causal target mask.
5. Run the encoder-decoder Transformer.
6. Calculate cross-entropy loss.
7. Ignore padding positions.
8. Backpropagate gradients.
9. Update parameters.
10. Monitor validation performance.

A simplified training structure is:

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

Important Improvements for a Real Translation System

The basic example demonstrates the architecture, but a practical translation model requires significantly more engineering.

Use a Real Dataset

Use a sufficiently large parallel corpus containing source and target language pairs.

Use Subword Tokenization

Instead of relying only on word-level tokenization, modern systems commonly use subword approaches.

Examples include:

• Byte Pair Encoding
• SentencePiece
• Unigram-based tokenization

Subword tokenization helps handle rare words and previously unseen word forms.

Handle Variable-Length Sequences

Real datasets contain sentences with different lengths.

Use:

• Padding
• Padding masks
• Batching strategies
• Length-aware batching where appropriate

Use Target Shifting

The decoder should receive previous target tokens and predict the next token.

Use Causal Masking

The decoder must not access future target tokens during autoregressive training.

Use Validation Data

Do not evaluate the model only on the training dataset.

Track validation loss and appropriate translation metrics.

Save Checkpoints

Save model and optimizer states during training so training can resume and the best checkpoint can be restored.

Use a Learning-Rate Scheduler

Transformer training can benefit from carefully designed learning-rate schedules, including warmup strategies.

Common Transformer Mistakes

Mistake 1: Forgetting Positional Information

Without positional information, the model has difficulty distinguishing token order.

Mistake 2: Incorrect Mask Shapes

Attention masks and padding masks have different roles and expected shapes. Always verify the PyTorch API requirements for the specific operation being used.

Mistake 3: Using the Same Target for Input and Label

For autoregressive training, shift the target sequence so the decoder predicts the next token.

Mistake 4: Not Ignoring Padding

Padding tokens should normally be excluded from the training loss.

Mistake 5: Training on Random Data

Random token IDs can verify that the model executes correctly, but they cannot teach meaningful language relationships.

Mistake 6: Forgetting Evaluation Mode

Use:

pythonCopy
1model.eval()

during evaluation and inference.

Mistake 7: Ignoring Device Placement

When using a GPU, ensure the model, input tensors, and masks are placed on compatible devices.

For example:

pythonCopy
1device = torch.device(
2    "cuda" if torch.cuda.is_available() else "cpu"
3)
4
5model = model.to(device)

Debugging a Transformer

When debugging a Transformer, inspect the tensor dimensions at each stage.

Check:

textCopy
1Source shape
2Target shape
3Embedding shape
4Encoder output shape
5Decoder output shape
6Vocabulary logits shape
7Mask shapes
8Target labels

For example:

textCopy
1Source:
2(batch, source_length)
3
4Target:
5(batch, target_length)
6
7Embedding:
8(batch, sequence_length, d_model)
9
10Output:
11(batch, target_length, target_vocab)

Shape debugging is one of the fastest ways to identify Transformer implementation errors.

Applications of Transformer Architecture

Transformer models are widely used for:

• Machine translation
• Text generation
• Text summarization
• Question answering
• Code generation
• Language modeling
• Information extraction
• Multimodal learning
• Document understanding

Transformer-based architectures also form the foundation of many modern NLP and generative AI systems.

Best Practices for Building Transformers in PyTorch

• Use clear tensor-shape conventions throughout the model.
• Use batch_first=True when it simplifies the implementation.
• Include positional information.
• Distinguish padding masks from causal masks.
• Use shifted decoder targets for autoregressive training.
• Ignore padding tokens when calculating loss.
• Use model.train() during training.
• Use model.eval() during evaluation and inference.
• Validate the model on data that was not used for training.
• Save checkpoints during long training runs.
• Use an appropriate tokenizer for real-world language data.
• Monitor both training and validation loss.
• Move tensors and the model to the same device.
• Start with a small model for debugging before scaling up.

Transformer From Scratch: Key Concepts

The most important concepts from this module can be summarized as:

textCopy
1Token IDs
2   ↓
3Token Embeddings
4   ↓
5Position Information
6   ↓
7Encoder
8   ↓
9Contextual Representations
10   ↓
11Decoder
12   ↓
13Linear Projection
14   ↓
15Vocabulary Logits
16   ↓
17Next Token

The encoder understands the source sequence, while the decoder generates the target sequence one step at a time.

Understanding this pipeline gives you the foundation required to study more advanced Transformer architectures.

Module Summary

In this module, you learned how to construct the major components required for an encoder-decoder Transformer using PyTorch.

You learned:

• How token embeddings represent vocabulary IDs as dense vectors.
• How positional embeddings provide sequence-order information.
• How padding masks prevent attention to padding positions.
• How causal masks restrict autoregressive attention to previous positions.
• How to assemble an encoder-decoder Transformer using nn.Transformer.
• How to project decoder representations into target vocabulary logits.
• How to calculate cross-entropy loss.
• How to train a Transformer using backpropagation and Adam.
• How greedy decoding generates tokens sequentially.
• How to structure an English-to-French translation project.
• Why target shifting and teacher forcing are important for sequence-to-sequence training.
• Why random data is useful for debugging but cannot train a meaningful translation system.

The complete architecture can be remembered as:

textCopy
1Input Tokens
2     ↓
3Token + Position Embeddings
4     ↓
5Transformer Encoder
6     ↓
7Encoder Memory
8     ↓
9Transformer Decoder
10     ↓
11Linear Output Layer
12     ↓
13Vocabulary Logits
14     ↓
15Autoregressive Decoding
16     ↓
17Generated Sequence

Once you understand these components, you are ready to move toward more advanced Transformer topics such as attention implementation, positional encoding strategies, BERT, GPT, T5, LLaMA-style architectures, efficient training, inference optimization, and large language model fine-tuning.