# Transformer Models From Scratch in Python

#52 May 26, 2025 Architecture

https://www.aibutsimple.com/p/transformer-models-from-scratch-in-python

Transformer neural networks (Vaswani et al. 2017) are a type of neural network specialized for modern natural language processing (NLP) tasks, such as translation, text generation, and chatbots.

They consist of an encoder and a decoder that work as an all-around NLP solution. In terms of the architecture, the main innovation behind the transformer is the self-attention mechanism—originally created for machine translation with RNNs.

The attention mechanism allows transformers to weigh the importance of different tokens (words or symbols) in an input sequence when making predictions.

In this issue, we’ll look at creating a transformer encoder from scratch using Python and Pytorch. We’ll dive into the full code and conceptual base behind the transformer, the architecture of choice for the most popular large language models (such as GPT or BERT).

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/40fe98b8-bbcf-4ea5-96c9-bed1c61ea060/Copy_of_Copy_of_Copy_of_Copy_of_Copy_of_CV_illust__5_.png)

This week, we have our third tutorial with code. Feel free to follow along, play around with the code, and learn about transformers. If you want to review some information on how transformers work conceptually, check out the issue below:

[Introduction to Transformers and the Attention Mechanism (Intermediate)](https://www.aibutsimple.com/p/transformers-and-the-attention-mechanism?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=transformer-models-from-scratch-in-python)

Note that basic Python knowledge along with some knowledge of libraries such as Pytorch is expected, so if you’re still new to it, we suggest going through some beginner tutorials to learn the syntax.

The code in this issue can be executed on a normal home PC and does not require a large amount of VRAM. It can also be executed on hosted notebooks such as Google Colab.

---

Big thanks to our partners for keeping this newsletter free. If you have a second, clicking the ad below helps us a ton—and who knows, you might find something you’ll enjoy.

### Start learning AI in 2025

Keeping up with AI is hard – we get it!

That’s why over 1M professionals read [Superhuman AI](https://www.superhuman.ai/c/confirmation?magiclink_subscription) to stay ahead.

- Get daily AI news, tools, and tutorials
- Learn new AI skills you can use at work in 3 mins a day
- Become 10X more productive

[Sign up and start mastering AI](https://www.superhuman.ai/c/confirmation?magiclink_subscription)

---

The first thing we’ll do to build this architecture is to inspect how it works—how the data flows through the structure. We notice that inputs are embedded and positionally encoded before being passed into the encoder.

The encoder is composed of two main components: self-attention (orange) and a feedforward layer (blue).

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/a4b26c25-bb42-43fc-8f37-3ebf32f54697/image.png)

Let’s say our input sentence is a simple one, such as “Deep learning is simple.”

The first thing we need to do is to start with all of our imports. We’ll import all libraries at once to have them ready for later—we don’t need to import them one by one:

```
from transformers import AutoTokenizer
import torch
from torch import nn
import torch.nn.functional as F
import math
```

We need to process the data first. Let’s tokenize the sentence, then embed the tokens:

```
tokenizer = AutoTokenizer.from_pretrained('bert-base-uncased')
text = 'Deep learning is simple.'
print(tokenizer(text, add_special_tokens=False, return_tensors='pt'))
input = tokenizer(text, add_special_tokens=False, return_tensors='pt')
```

The above code should produce the following output: {'input_ids': tensor([[2784, 4083, 2003, 3722, 1012]]), 'token_type_ids': tensor([[0, 0, 0, 0, 0]]), 'attention_mask': tensor([[1, 1, 1, 1, 1]])}

In the above code, we basically took a portion of the BERT model to tokenize our sentence. In the output, if we look at the input_ids tensor, we see that the tokenizer has tokenized the sentence “Deep learning is simple.” into tokens.

We’ll model most of the transformer encoder with some important values, such as the vocabulary size or the hidden size. Let’s define all of our values to use below:

```
vocab_size = 30522
hidden_size = 768
intermediate_size = 3072
num_heads = 12
```

Next, we need to create some dense embeddings—dense meaning that most entries in the embedding list are nonzero. In PyTorch, we can do this by using a torch.nn.Embedding layer that acts as a lookup table for each input ID. Then, we can actually embed our tokens by feeding in the input IDs:

```
token_embeds = nn.Embedding(vocab_size, hidden_size)
input_embeds = token_embeds(input.input_ids)
print(token_embeds) # output: Embedding(30522, 768)
print(input_embeds.size()) # output: torch.Size([1, 5, 768])
```

The output (the input embeddings) is a tensor of shape [batch_size, seq_len, hidden_size], where our hidden size is the size of the vector representing each token.

Currently, the input embeddings are not positionally embedded. For simplicity purposes, we’ll skip the positional embeddings, although we have an expanded version for [supporters only](https://www.aibutsimple.com/upgrade?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=transformer-models-from-scratch-in-python).

Now that we have our embedding matrix, we move on to the multi-head attention portion of our transformer. Having several heads allows the model to focus on several aspects at once, enhancing its NLP performance. We’ll compute it by using multiple linear layers and scaled dot-product attentions as shown on the diagram:

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/1da71d31-078c-4772-9e5f-5f752cc44cbb/Copy_of_Copy_of_Copy_of_Copy_of_CV_illust__3_.png)

The scaled dot-product attention that is in purple is simply the self-attention mechanism. The self-attention mechanism is calculated using four main steps:

- First, we project the embedding matrix into three matrices called the query (Q), key (K), and value (V) matrices. For self-attention, Q, K, and V are the same since no linear layers are used.
- We then take the dot product between the query (Q) matrix and the key matrix (K), which is a matrix multiplication between Q and K transpose. At this point, we obtained our attention scores, which is an n-by-n matrix, where n is the input size in tokens.
- We scale the attention scores by the square root of the embedding dimension (d), then pass it through a softmax function to normalize the values. As masking is optional, we’ll skip it for simplicity purposes.
- Finally, we perform a matrix multiplication of the softmax output with the value (V) matrix to obtain our final output.

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/96686868-ee5e-461a-8448-0338b791776d/image.png)

Diagram of scaled dot-product attention

If you want to learn or review exactly how data passes through transformers mathematically, check out our issue diving into the math step-by-step [here](https://www.aibutsimple.com/p/transformers-and-multi-head-attention-mathematically-explained?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=transformer-models-from-scratch-in-python).

For multi-head attention, we stack multiple self-attention mechanisms together in layers and obtain different Q, K, and V values for each layer by projecting the token embedding into three copies and multiplying each by its own specific weight matrix.

Thus, each head of attention will have different Q, K, and V matrices, along with the fact that each Q, K, and V are different from each other (not like self attention)—their weights are learned during training.

We’ll implement multi-head step-by-step by using smaller classes, then building up.

```
query = key = value = input_embeds

def scaled_dot_product_attention(query, key, value):
    dim_k = query.size(-1)
    # bmm is batch matrix multiplication
    scores = torch.bmm(query, key.transpose(1, 2)) / math.sqrt(dim_k)

    # Skip masking
    weights = F.softmax(scores, dim=-1)
    return torch.bmm(weights, value)

class AttentionHead(nn.Module):
    def __init__(self, embed_dim, head_dim):
        super().__init__()
        self.q = nn.Linear(embed_dim, head_dim)
        self.k = nn.Linear(embed_dim, head_dim)
        self.v = nn.Linear(embed_dim, head_dim)
    
    def forward(self, hidden_state):
        # Use scaled dot product attention
        attn_outputs = scaled_dot_product_attention(self.q(hidden_state), self.k(hidden_state), self.v(hidden_state))
        return attn_outputs

class MultiHeadAttention(nn.Module):
    def __init__(self):
        super().__init__()
        head_dim = hidden_size // num_heads

        # Multiple heads
        self.heads = nn.ModuleList([AttentionHead(hidden_size, head_dim) for _ in range(num_heads)])
        self.output_linear = nn.Linear(hidden_size, hidden_size)
        
    def forward(self, hidden_state):
        # Concatenate the scaled dot product outputs 
        x = torch.cat([h(hidden_state) for h in self.heads], dim=-1)
        x = self.output_linear(x)
        return x

multihead_attn = MultiHeadAttention()
attn_output = multihead_attn(input_embeds)
print(attn_output.size()) # output: torch.Size([1, 5, 768])
```

You’ll notice that the multi-head attention class uses “attention head” objects, which create “scaled dot product attention” objects, tying the pieces of code together to produce an output tensor.

In the attention head class, we initialize three independent linear layers that apply matrix multiplication to the embedding vectors to produce tensors of shape [batch_size, seq_len, head_dim], where head_dim represents the number of dimensions we are projecting into.

self.q = nn.Linear(embed_dim, head_dim) self.k = nn.Linear(embed_dim, head_dim) self.v = nn.Linear(embed_dim, head_dim)

In practice, head_dim is chosen to be a multiple of embed_dim (simply the hidden_size, the size of the full embedding vector) so that the computation across each head is constant. For instance, the BERT model has 12 heads of attention, so the dimension of each head is 768/12 = 64. We make this calculation in the code.

Our multi-head output is obtained by concatenating all the matrices from the individual heads together, and it ends up still being a tensor of size [batch_size, seq_len, hidden_size], which is constant throughout the transformer, excluding the input and output.

If we print the output, it looks a little something like this, although the numbers will vary a little bit due to the randomness in the weight initialization of nn.Embedding:

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/5637a71f-d9c2-483c-9e77-4ada1f353bdf/image.png)

Moving on, we’ll implement the feed-forward layer that comes right after our multi-head attention to complete an entire “encoder” block. Keep in mind that in the original paper, many encoder blocks are stacked on top of each other to create effective architectures.

The feed-forward block in transformers is just a fully connected neural network that has two layers—with a small twist. For transformers specifically, feed-forward layers are applied independently and identically to each token in each sequence in our batch.

This means that every token gets transformed by the same parameters (weights and biases) when it passes through. We refer to this type of feed-forward layer as position-wise feed-forward networks (FFNs).

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/e72d5d58-a2e7-41de-9ea5-977e07cf8411/Copy_of_Copy_of_Copy_of_Copy_of_Copy_of_CV_illust__2_.png)

The code for the position-wise FFN is shown below:

```
class FeedForward(nn.Module):
    def __init__(self):
        super().__init__()
        self.lin1 = nn.Linear(hidden_size, intermediate_size)
        self.lin2 = nn.Linear(intermediate_size, hidden_size)
        self.gelu = nn.GELU()
        self.dropout = nn.Dropout(0.1)

    def forward(self, x):
        x = self.lin1(x)
        x = self.gelu(x)
        x = self.lin2(x)
        x = self.dropout(x)
        return x
```

In this implementation, each linear layer transforms the input using a weight matrix and a bias (Wx + b) as usual. After passing through the two layers, the output of those layers passes through the GeLU activation function, which is just a smoother version of the ReLU.

The activation output finally passes through a dropout layer, deactivating random neurons by setting them to 0 to reduce overfitting.

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/06aa86db-6c52-40c6-abe1-ad70331319a4/image.png)

The Gaussian Error Linear Unit (GeLU)

Let’s create a “feed forward” object, pass the data through, and check its output:

```
feed_forward = FeedForward()
ff_outputs = feed_forward(attn_output)
print(ff_outputs.size()) # output : torch.Size([1, 5, 768])
print(ff_outputs)
```

The output tensor should look something like this with varying values:

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/ebd0c3ec-f72a-4bbd-9a48-bdf6e90ea51a/image.png)

The full encoder can be made by repeating blocks of the multi-head attention, feedforward and norm layers multiple times until the desired size is reached. After passing through the encoder, the data will flow to the decoder to continue computations.

---

Here’s a special thanks to our [biggest supporters](https://www.aibutsimple.com/upgrade?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=transformer-models-from-scratch-in-python):

Sushant Waidande

---

If you enjoy our content, consider supporting us so we can keep doing what we do. Please share this with a friend!

Want to reach 50000+ ML engineers? Let’s work together:

[WORK WITH US 📈](https://www.passionfroot.me/aibutsimple?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=transformer-models-from-scratch-in-python)

Feedback, inquiries? Send us an email at [edwin@aibutsimple.com](mailto:edwin@aibutsimple.com).

If you like, you can also donate to our team to push out better newsletters every week!

[Support us ☕](https://ko-fi.com/aibutsimple?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=transformer-models-from-scratch-in-python)

---

That’s it for this week’s issue of AI, but simple. See you next week!

—AI, but simple team