Transformer Implementation in PyTorch

The Transformer can be implemented as connected PyTorch modules. The main purpose is to trace the data through attention, encoder layers, decoder layers, masks, and the final vocabulary output.

Model configuration

A useful configuration uses these values:

Setting	Value
Batch size	64 sequences
Sequence length	100 token positions
Model width	512 features per token
Attention heads	8
Features per head	64
Encoder and decoder layers	6 each
Feed-forward width	2,048
Source and target vocabulary	5,000 tokens each

The source token IDs start with shape (64, 100). Token embedding and positional encoding produce (64, 100, 512).

Multi-head attention

Each projection keeps the model width:

self.W_q = nn.Linear(d_model, d_model)
self.W_k = nn.Linear(d_model, d_model)
self.W_v = nn.Linear(d_model, d_model)
self.W_o = nn.Linear(d_model, d_model)

For these settings, the data flow is:

Q, K, V                         (64, 100, 512)
split into 8 heads              (64, 8, 100, 64)
scaled attention scores         (64, 8, 100, 100)
weighted values                 (64, 8, 100, 64)
combine heads                   (64, 100, 512)
output projection               (64, 100, 512)

d_model must be divisible by num_heads. Each head then has width d_k = d_model // num_heads.

Positional encoding

A useful implementation creates sinusoidal positional encodings once and stores them with register_buffer. A buffer moves with the model but is not a trainable parameter.

def forward(self, x):
    return x + self.pe[:, :x.size(1)]

x.size(1) is the sequence length. The slice selects one positional vector for each token position in the current input.

Encoder layer

One encoder layer performs these operations:

self-attention
-> residual addition and layer normalisation
-> position-wise feed-forward network
-> residual addition and layer normalisation

The self-attention call uses the same tensor for query, key, and value:

attn_output = self.self_attn(x, x, x, mask)

Decoder layer

One decoder layer performs:

masked self-attention
-> residual and normalisation
-> encoder-decoder attention
-> residual and normalisation
-> feed-forward network
-> residual and normalisation

In encoder-decoder attention, the decoder supplies the query. The encoder output supplies the keys and values.

attn_output = self.cross_attn(x, enc_output, enc_output, src_mask)

Masks

We can use two types of mask:

A source padding mask blocks token ID 0
A target mask blocks padding and future target positions

The no-peek part is a lower-triangular Boolean mask. It prevents a target position from reading later target tokens during training.

Complete Transformer flow

source IDs -> source embedding and position -> encoder stack -> encoder output
target IDs -> target embedding and position -> masked decoder stack
decoder output -> Linear(d_model, target vocabulary size) -> token logits

The final output has one set of 5,000 vocabulary logits for each target position.

Training shift

The decoder input and prediction target use the same target sequence with a one-position shift:

output = transformer(src_data, tgt_data[:, :-1])
loss = criterion(
    output.contiguous().view(-1, tgt_vocab_size),
    tgt_data[:, 1:].contiguous().view(-1),
)

The decoder receives all target tokens except the final token. The loss compares each output position with the following target token.

CrossEntropyLoss(ignore_index=0) excludes padding targets. The training step uses this order:

optimizer.zero_grad()
-> forward pass
-> loss
-> loss.backward()
-> optimizer.step()

Implementation checkpoints

A basic decoder can return None as its third item so the training loop can use the same output structure as an attention decoder
BahdanauAttention uses two Linear(hidden_size, hidden_size) projections and one Linear(hidden_size, 1) scoring projection
An attention decoder joins the embedded decoder input with the attended context, so its GRU receives 2 * hidden_size input features
Multi-head attention uses four Linear(d_model, d_model) projections for query, key, value, and output
Positional encoding returns x + self.pe[:, :x.size(1)]

Code

PyTorch: configure and train the complete model

This block uses the Transformer class assembled from the attention, position, encoder, decoder, and mask components above.

import torch
from torch import nn, optim
 
src_vocab_size = 5000
tgt_vocab_size = 5000
d_model = 512
num_heads = 8
num_layers = 6
d_ff = 2048
max_seq_length = 100
dropout = 0.1
 
transformer = Transformer(
    src_vocab_size,
    tgt_vocab_size,
    d_model,
    num_heads,
    num_layers,
    d_ff,
    max_seq_length,
    dropout,
)
 
src_data = torch.randint(1, src_vocab_size, (64, max_seq_length))
tgt_data = torch.randint(1, tgt_vocab_size, (64, max_seq_length))
 
criterion = nn.CrossEntropyLoss(ignore_index=0)
optimizer = optim.Adam(
    transformer.parameters(),
    lr=0.0001,
    betas=(0.9, 0.98),
    eps=1e-9,
)
 
transformer.train()
 
for epoch in range(100):
    optimizer.zero_grad()
 
    output = transformer(src_data, tgt_data[:, :-1])
    output_for_loss = output.contiguous().view(-1, tgt_vocab_size)
    expected_tokens = tgt_data[:, 1:].contiguous().view(-1)
 
    loss = criterion(output_for_loss, expected_tokens)
    loss.backward()
    optimizer.step()
 
    print(f"Epoch: {epoch + 1}, Loss: {loss.item():.4f}")

The target slice gives the decoder 99 input positions and asks it to predict the following 99 positions.

Value	Shape	Meaning
`src_data`	`(64, 100)`	Source token IDs
`tgt_data[:, :-1]`	`(64, 99)`	Decoder input without the final target token
`output`	`(64, 99, 5000)`	Vocabulary logits for every decoder position
`output_for_loss`	`(6336, 5000)`	All position logits in one classification batch
`expected_tokens`	`(6336,)`	The next-token ID for every output row

transformer.train() enables training behaviour such as dropout. optimizer.zero_grad() clears gradients from the earlier update. loss.backward() calculates new gradients, and optimizer.step() updates the model parameters.

JJ

Explorer