Controlled experiments with attention-head counts and activation functions in a small GPT model.
TL;DR
I tested two architectural choices in a small GPT model while keeping the rest of the setup stable:
Attention heads: 1, 4, and 8 heads under fixed-token and fixed-time budgets.
Activation functions: GELU, ReLU, ReLU², and Tanh under a fixed-token budget.
At equal token exposure, increasing the number of attention heads produced only a small improvement in validation loss. Under the fixed-time budget, eight heads performed noticeably better—but primarily because that configuration processed more tokens.
For activation functions, ReLU² achieved the lowest validation loss, while Tanh performed the worst. Their internal behavior differed dramatically: ReLU and ReLU² were highly sparse, ReLU² produced extremely large positive activations, and Tanh bounded everything between −1 and 1.
The broader takeaway was simple:
Architectural choices can significantly change how a model computes without producing equally large differences in model quality.
The complete code for these experiments is available in my GitHub repository.
So far in this series, I have experimented with learning rates, context length, model size, and fine-tuning strategies. Most of those experiments changed how the model was trained.
This time, I wanted to look inside the model itself.
A GPT model contains several architectural decisions: how attention is divided, which activation function is used, where normalization happens, and how position is represented.
Testing all of them together would make it difficult to identify what caused each result. I therefore focused on two controlled experiments:
Changing the number of attention heads
Changing the activation function
The baseline model contained approximately 13.67 million parameters, used a context length of 128, and was trained on WikiText data.
Multi-head attention divides the model’s embedding space into smaller attention subspaces.
The embedding dimension remained fixed at 128:
| Heads | Dimension per head |
|---|---|
| 1 | 128 |
| 4 | 32 |
| 8 | 16 |
The parameter count remained identical across all three models. The only difference was how the attention representation was divided.
I evaluated the configurations under two budgets.
Under a fixed-token budget, I expected the differences to be small. Every model would see the same amount of training data and complete the same number of optimizer updates.
Under a fixed-time budget, I expected multiple heads to potentially benefit from better hardware parallelism. However, I did not expect head count alone to guarantee higher computational efficiency.
At first glance, multi-head attention can look like a computational shortcut: instead of performing one large attention operation, we split the embedding space into smaller subspaces and process them separately. Conceptually, that is true—but it does not automatically reduce the total amount of work.
Practical throughput can still change because of:
Hardware parallelism
Memory-access patterns
Tensor shapes
Framework and kernel overhead
In my fixed-time run, additional heads appeared to benefit from the execution environment. However, the timing results were not consistent across both budgets, so I would not generalize that more heads are inherently faster.
Each model processed exactly 200,192 tokens across 782 optimizer steps.
Eight heads achieved the lowest validation loss, followed by four heads and one head.
However, the difference between the best and worst result was only 0.027 loss points. Accuracy also improved only slightly, from approximately 12.07% with one head to 12.29% with eight heads.
This broadly supported the hypothesis. Multiple heads may have provided a modest representational advantage, but head count was not a major performance lever when training exposure was held constant.
Each model received approximately 300 seconds of measured training time.
The eight-head model processed the most tokens and achieved the lowest validation loss. It processed approximately:
21% more tokens than the one-head model
12% more tokens than the four-head model
The fixed-token experiment had shown only a small architectural advantage for eight heads. This suggests that much of the larger fixed-time improvement came from processing more training data—not simply from having more heads.
For this small GPT model, attention-head count appeared to be a secondary architectural choice. Performance under equal token exposure changed only modestly, while the amount of training completed within the available time had a larger effect.
I chose four functions that represent meaningfully different behaviors.
GELU is a strong modern baseline for Transformer models. It is smooth, allows small negative outputs, and does not abruptly shut values off. I would use it when I want a balanced and reliable default.
ReLU sets negative inputs to zero while allowing positive inputs to grow linearly. This creates sparse representations, but it can also discard useful information on the negative side.
ReLU² retains ReLU’s hard zeroing behavior but squares positive values. Small positive values stay modest, while larger values are amplified aggressively. I would use it to test whether emphasizing a small number of strong signals improves learning—while monitoring activation and gradient stability.
Tanh is smooth and dense but bounded between −1 and 1. This controls activation magnitude, but it can also saturate. Once inputs become large, further input changes produce increasingly small output changes. I would primarily use it when bounded outputs are important or as an extreme comparison against unbounded functions.
All four models had identical parameter counts and processed exactly 200,192 tokens.
I expected GELU to remain a stable baseline.
ReLU might create greater sparsity but lose information by discarding negative values.
ReLU² might learn efficiently by amplifying strong positive signals, but it could also create unstable activation magnitudes.
Tanh represented the opposite extreme. Its bounded outputs could control activation growth, but saturation might weaken learning.
ReLU² achieved the lowest validation loss, slightly ahead of GELU. The difference was small, so I would describe it as a promising result rather than a decisive win.
Its internal behavior was much more dramatic.
Approximately 76% of ReLU² outputs were zero, but its maximum activation exceeded 203, compared with around 12 for GELU and ReLU. It suppressed most activations while strongly amplifying a small number of positive signals.
Tanh behaved exactly like the extreme comparison I wanted. Its activations remained bounded between −1 and 1, but it produced the highest validation loss, lowest accuracy, and slowest throughput.
Bounding activations prevented large values, but it may also have compressed useful distinctions between strongly activated features. Tanh showed that controlling activation magnitude too aggressively can restrict learning.
The activation functions changed the model’s internal behavior far more dramatically than they changed its final validation loss.
For this setup, eight attention heads and ReLU² achieved the lowest validation losses.
However, neither result was decisive.
Changing the number of heads produced only a modest improvement when token exposure was controlled. Changing the activation function created radically different sparsity and activation magnitudes, but the final loss differences remained relatively small.
These results describe early learning for one small model, dataset, training budget, and random seed. Repeating the closest configurations across multiple seeds would be necessary before treating the rankings as general conclusions.
1. MoH: Multi-Head Attention as Mixture-of-Head Attention — 2024
2. Peri-LN: Revisiting Layer Normalization in the Transformer Architecture — 2025
3. 2 OLMo 2 Furious — 2024/2025
4. Round and Round We Go! What Makes Rotary Positional Encodings Useful? — 2024
5. GLU Variants Improve Transformer — 2020