Deep Dive into a Custom GPT Architecture

Now that I have covered the fundamentals of neural networks and have a grasp of how a transformer works, I decided to build a small GPT from scratch, train it on different datasets, and see what happens.

The first step, as usual, was to implement it without relying on libraries that already provide the complete architecture. Rebuilding it from scratch helps me remember that there is nothing magical about it: it is math and scale working together to produce something that can appear surprisingly complex.

I wanted to have a solid base that I could fork whenever I wanted to play around with a model, so I created a repository that I can use as a starting point for different experiments. The details need to change depending on the experiment, but the core architecture remains the same.

You can play with this model even on a laptop, but if you want to experiment with a bit more scale, you will need a GPU. For these experiments, I am using an old PC that I built in 2020. It is not particularly powerful, but it is a good starting point. I turned it into a server running Ubuntu LTS and connected it through Tailscale, so I can access it over SSH. It has 32 GB of RAM, a Ryzen 5 3600, and, most importantly, an RX 5600 XT. To use the GPU as a backend, I had to compile a custom PyTorch build and configure it so that the code can access the GPU through the CUDA device interface exposed by PyTorch.

Before looking at the architecture, I want to derive the computational cost of training and inference for a model like this one.

Deriving the computational cost of training and inference from scratch

Computational cost is usually measured in FLOPs (floating-point operations). In this approximation, one multiplication and one addition each count as one FLOP. The goal is to derive a rough mental model of the computational cost of training and inference for an LLM.

We can start from the basics. The fundamental building block is matrix multiplication. Suppose we want to multiply the following matrices:

where $A$ has shape $M \times K$, $B$ has shape $K \times N$, and $C$ has shape $M \times N$. Each element of $C$ is obtained from the dot product between one row of $A$ and one column of $B$, so:

For each dot product, we perform $K$ multiplications and $K-1$ additions, which we approximate as $2K$ operations. We repeat this for every row of $A$ and every column of $B$.

Now we can extend that reasoning to a classic linear layer:

Nothing that we have not seen yet, it's still a matrix multiplication. Suppose that $x$ has shape $(B,T,C)$, where $B$ is the batch size, $T$ is the sequence length, and $C$ is the embedding dimension. The weight matrix $W$ has shape $(C,D)$.

We can treat the first two dimensions of $x$ as $B \cdot T$ rows, each multiplied by $W$. The cost is therefore:

When $C=D$, this becomes:

Now we can try to estimate the cost of an attention block. In an attention block we have the following matrices:

• $x$: $(B, T, C)$ input
• $W_q$: $(C, head\_size)$
• $W_k$: $(C, head\_size)$
• $W_v$: $(C, head\_size)$

Then we have to perform the following operation, for a single head:

1. q = x @ W_q → (B, T, head_size)
2. k = x @ W_k → (B, T, head_size)
3. v = x @ W_v → (B, T, head_size)
4. scores = q @ k.transpose(-2, -1) → (B, T, T)
5. attention = softmax(scores) → (B, T, T)
6. y = attention @ v → (B, T, head_size)

Now we can compute the FLOPs required by each step for a single head:

1. $flops = 2 \cdot B \cdot C \cdot T \cdot head\_size$
2. $flops = 2 \cdot B \cdot C \cdot T \cdot head\_size$
3. $flops = 2 \cdot B \cdot C \cdot T \cdot head\_size$
4. $flops = 2 \cdot B \cdot head\_size \cdot T^2$
5. $flops \approx 4 \cdot B \cdot T^2$ (derived below)
6. $flops = 2 \cdot B \cdot T^2 \cdot head\_size$

Since every attention head performs the same computation independently, and there are $h$ heads, the total cost is $h$ times the cost derived above. Using $h \cdot head\_size = C$:

After the heads are concatenated back into a $(B,T,C)$ tensor, there's one more matmul — the output projection — which is not per-head, so it isn't scaled by $h$:

Adding it in:

On the softmax cost: Softmax on a row of length $T$ requires: a max-reduction for numerical stability ($T-1$ comparisons), a subtract-and-exponentiate pass ($2T$ flops — one subtraction and one exp per element, counting exp as 1 flop by convention), a sum-reduction ($T-1$ additions), and a final divide ($T$ divisions). That's $\approx 5T$ flops per row, though $4T$ is the more common rounding in the literature and the two are interchangeable at this level of approximation. There are $B \cdot T$ rows per head (one per query position, per batch element), so one head's softmax costs $\approx 4 \cdot B \cdot T^2$ — independent of $head\_size$. This is why it has to be added on top of the per-head total above rather than folded into it via the $h \cdot head\_size = C$ substitution: it's the one term whose cost doesn't factor through that product, since it scales with the number of heads (independent softmax calls), not with how wide each head is.

Now we can derive the cost of the other fundamental block: the feedforward network. It usually contains two linear layers with a non-linearity between them. The first linear layer expands the input dimension by a factor of four, and the second projects it back to the original dimension:

1. h = x @ W1 (B, T, C) @ (C, 4C) -> (B, T, 4C)
2. h = activation(h) (B, T, 4C) -> (B, T, 4C)
3. y = h @ W2 (B, T, 4C) @ (4C, C) -> (B, T, C)

We already have the formula for a linear layer, so we can reuse it directly for step 1 and step 3.

1. $flops = 2 \cdot B \cdot T \cdot C \cdot 4C = 8 \cdot B \cdot T \cdot C^2$
2. activation is elementwise, cost linear in the number of elements ($B \cdot T \cdot 4C$), negligible compared to the matmuls, so we drop it
3. $flops = 2 \cdot B \cdot T \cdot 4C \cdot C = 8 \cdot B \cdot T \cdot C^2$

Summing:

Notice this is exactly twice the $8 \cdot B \cdot T \cdot C^2$ term from attention (the qkv + output projection part) — the FFN is the more expensive of the two sub-blocks per layer, by a factor of 2, for the standard 4x expansion ratio.

We can now combine the two components to obtain the computational cost of one transformer layer, which contains an attention block and a feedforward network. A complete transformer stacks multiple layers of this type.

So at the end we obtain

If the model contains $l$ layers, their contribution to one forward pass is $l$ times $flops_{layer}$. The complete model also contains the embedding layer and the final projection over the vocabulary.

The embedding layer is basically a lookup, so we can consider its arithmetic cost equal to zero in this approximation. At the end of the model we also have the language model head, which projects the representation of every token from $C$ to vocab_size. This is another linear layer, so its cost is:

Even when the language model head shares its weights with the token embedding, we still have to perform this matrix multiplication. Weight tying reduces the number of parameters, not the computational cost of the projection.

Putting together the $l$ transformer layers and the final projection, we obtain the cost of one complete forward pass:

Often it is more useful to reason about the cost per token. In one forward pass we process $B \cdot T$ tokens, so we can divide the previous formula by this value:

This is the cost of a complete forward pass over a sequence of length $T$. In this implementation, autoregressive generation runs another forward pass for every generated token. Implementations with a KV cache can reuse previous keys and values, changing the inference cost.

These formulas make the two main scaling behaviours more visible. The projections inside the transformer grow quadratically with the embedding dimension $C$. In the complete forward pass, attention grows quadratically with the context length $T$; after dividing by the number of tokens, the corresponding per-token cost grows linearly with $T$.

For training we also have to perform the backward pass. A useful rough approximation is that the backward pass costs around two times the forward pass:

The complete training cost is then approximately three times the forward cost. Per token we obtain:

Custom GPT Architecture

Clearly I have not developed from scratch this model architecture, is derived from the original attention paper, plus some other sources, like karpathy's videos and also deep learning books. It's pretty easy, nothing too complicated, I just want to graps the fundamentals, so it's deliberately written in clanky python and pytorch, nothing that must be taken for serious.

You may already be familiar with the architecture, so you can skip this section. To better retain what I am learning, I like to sketch the systems I am building. In this case, I used Excalidraw, which I found very helpful:

As we have already seen, the classic transformer block contains a self-attention layer followed by a feedforward network. The main difference is the placement of layer normalization: the original paper uses post-norm, while this implementation uses pre-norm. Pre-norm generally makes optimization more stable, especially as the number of layers increases.

The implementation is available on GitHub, together with the code required for training and inference. Instead of loading the entire training set into RAM, I used NumPy memory mapping to create a disk-backed array and load only the required portions. This prevented crashes and out-of-memory errors during training.

Now we can look at some examples of this custom GPT at work.

SumGPT

The first experiment is SumGPT, a small model trained to sum two three-digit numbers. The input has a format like 123+456= and the model has to generate the result followed by a newline. I decided to generate the digits of the result in reverse order, because this aligns the carry operation with the autoregressive direction of the model: it can start from the units and then move to tens, hundreds and thousands.

For this task I used a deliberately small configuration, with an embedding size of 128, 4 attention heads, 4 transformer blocks and a context length of 12. The complete operation fits inside this context, so there is no reason to use a larger model. On my GPU the training took around three minutes, and on 1000 randomly generated sums it reached around 99% accuracy.

An interesting thing that I noticed during training was that the loss was not measuring only the ability to perform the sum. The operands before the = sign are generated randomly, so asking the model to predict them introduces a part of the loss that cannot be reduced: there is no pattern that allows it to know which random digit comes next. I fixed this by masking all targets before =, so the cross-entropy only measures the generated result.

This is a small and constrained experiment, but it shows that the architecture can learn an algorithmic pattern and not only the statistical structure of natural language. It also shows how important it is to understand what the loss is actually measuring: a high loss does not always mean that the model is failing at the task we care about.

BookGPT

The second experiment is BookGPT, a character-level language model trained on Italian books. Here the task is more open: given a sequence of characters, the model has to predict the next one and generate prose one character at a time.

The first dataset was composed of 19 books from an Hugging Face dataset, plus I Promessi Sposi. During the experiment I discovered that most of those books were English Gutenberg texts translated automatically into Italian. The corpus contained duplicated fragments, strange sentences and even a book with the wrong title. The model learned these artifacts: the generated text had a recognizable Italian structure, but it also produced many invented words and English names like Sighbury and Mr. Wergomestane.

For this reason I rebuilt the dataset using 15 native Italian works from authors like Manzoni, Verga, Pirandello, Nievo, De Amicis, Deledda and Collodi. The final clean corpus contains around 9.48 million characters and a vocabulary of 140 characters. Every book was split independently into 90% training and 10% validation before concatenating them, to avoid having a validation set biased toward only the last books in the corpus.

I trained three different configurations and compared their validation loss and generated text at matching training steps.

The first one was the baseline trained on the old corpus. It had around 852K parameters, an embedding size of 128, 4 heads, 4 blocks and a context length of 256. It reached a validation loss of 1.4355 at step 5000 and 1.3416 at step 8250. Training and validation loss remained very close, so there was no clear sign of overfitting. The generated text started to reproduce Italian syntax, but the bad quality of the corpus was clearly visible in its vocabulary.

The second run used the clean corpus and a larger model, with around 4.09M parameters, an embedding size of 256, 8 heads, 6 blocks and the same context length of 256. This was the best result. At step 2000 its validation loss was 1.5673, compared with 1.7487 for the baseline, and at step 4750 it reached 1.3969, compared with 1.4497. The validation loss continued to decrease, while the larger gap from the training loss stabilized instead of continuing to grow.

The generated text improved too. The model started to generate names present in the corpus, such as zio Trao from I Malavoglia, and it often maintained a plausible agreement between subjects and verbs. However, the text was still not coherent for long. Subjects could disappear or change inside the same generation, and many words were still invented even if they had a plausible Italian morphology.

Since the main limitation seemed to be coherence, I tried a third run with the context length increased from 256 to 512. The hypothesis was that seeing more previous characters would help the model maintain subjects and references for longer. The result was negative: at step 4750 the validation loss was 1.5241, worse than both the baseline at 1.4497 and the larger model at 1.3969. The generated text did not show a clear improvement either. With prompts like Roma or lungo i fiumi, the model produced locally plausible Italian fragments, but the sentence still lost its direction and introduced invented words.

On this hardware the larger context was not free. Without an efficient attention kernel, its memory cost forced me to reduce the batch size from 128 to 64. The model was therefore learning a more difficult task with noisier gradient updates, without receiving more representational capacity. In this setup, investing the available compute in width and depth produced better results than investing it in a longer context.

The main observation from BookGPT is that context length alone does not create coherence. A wider window only gives the model the possibility to look further back; it does not guarantee that it will learn how to use that information. In this experiment, a clean corpus and more model capacity had a much clearer effect than simply doubling the context length.

Conclusion

These experiments made the trade-offs of a GPT architecture much more concrete. By implementing it from scratch I also got a better grasp of its fundamental building block: a transformer block composed of causal self-attention and a feedforward network, together with residual connections and layer normalization. Deriving its computational cost also made it clear which parts scale with the model dimension and which ones become expensive as the context grows.

SumGPT showed that even a very small model can learn a precise algorithmic task, but only if the training objective measures the right thing. BookGPT showed that for language modeling the architecture is only one part of the problem: cleaning the dataset and increasing useful model capacity improved the results more than simply giving the model a longer context.

The longer-context run was probably the most useful negative result. A larger context window does not automatically produce better coherence, especially when it forces a smaller batch size and the model does not have enough capacity to use the additional information. On this hardware and at this scale, a context of 256 with a wider and deeper model was a better allocation of compute than a context of 512.

The generated text is still far from coherent prose, but that is consistent with the size of the model and the character-level setup. The useful result was not producing a good language model; it was understanding which changes actually improved it, which ones did not, and why.