64D-vs-32D Heads at Small Scale
We wanted to know if more heads with less dimension per head beats fewer heads with more dimension per head. So we deciced to train two slightly different models on about 500M tokens of MiniFineWeb.
Model Architectures
Config A: 8H-32D
- Base Architecture:
LlamaForCausalLM - Tokenizer:
Harley-ml/Dillionv2-1.3M - Transformers Version:
5.13.1 - Hidden Size:
256 - Vocab Size: 2564
- Number of Layers:
6 - Number of Heads:
8 - Number of KV Heads:
2 - Intermediate Size:
512 - Head Dim:
64 - Max Position Embeddings:
256 - RoPe Theta:
2500.0 - Tie Word Embeddings:
true - Hidden Activation:
silu - MLP Bias:
false - Initializer Range:
0.2 - RMS Norm Eps:
1e-06 - Pretraining Tp:
1 - Use Cache:
false - Total Parameters:
4,002,048
Config B: 4H-64D
- Base Architecture:
LlamaForCausalLM - Tokenizer:
Harley-ml/Dillionv2-1.3M - Transformers Version:
5.13.1 - Hidden Size:
256 - Vocab Size: 2564
- Number of Layers:
6 - Number of Heads:
4 - Number of KV Heads:
1 - Intermediate Size:
512 - Head Dim:
64 - Max Position Embeddings:
256 - RoPe Theta:
2500.0 - Tie Word Embeddings:
true - Hidden Activation:
silu - MLP Bias:
false - Initializer Range:
0.2 - RMS Norm Eps:
1e-06 - Pretraining Tp:
1 - Use Cache:
false - Total Parameters:
4,002,048
Training Setup
- Epochs:
1 - Max Steps:
-1.0 - Batch Size:
400 - Sequence Length:
256 - Gradient Accumulation:
2 - Gradient Clipping:
1.0 - Gradient Checkpointing:
true - Learning Rate:
3e-3 - Eval Split:
0.00165 - Weight Decay:
0.01 - Optimizer:
AdamW - AdamW Betas:
(0.9, 0.95) - AdamW Eps:
1e-8 - Scheduler:
WSD - WSD Warmup Ratio:
0.015 - WSD Stable Ratio:
0.78 - WSD Decay Ratio:
0.20 - WSD Minium LR Ratio:
0.0 - WSD Number of Cycles:
0.5 - DType:
float16 - Torch.Compile:
true - DataLoader Workers:
4 - Seed:
311
Results
| Model | Final Val Loss | Arc Easy | HellaSwag | PiQA | Avg |
|---|---|---|---|---|---|
| Config A | 2.9821 | 29.71% | 27.39% | 52.94% | 36.68% |
| Config B | 3.0067 | 29.92% | 27.39% | 52.07% | 36.46% |
Config A barely edges out Config B on Avg - the two are close enough that at this scale, the choice mostly comes down to preference.
That said, we'd expect head dimension to matter more as models get bigger. For a >=200M-parameter model, we'd lean toward 64–128 dims per head rather than 32.
Model Checkpoints
The two models are stored separately in different folders in this repository. To load them, use:
from transformers import AutoModelForCausalLM, AutoTokenizer
config_a = AutoModelForCausalLM.from_pretrained(
"fromziro/Width-Vs-Depth",
subfolder="config_a",
)
# load config b instead:
# config_b = AutoModelForCausalLM.from_pretrained(
# "fromziro/32D-Vs-64D",
# subfolder="config_b",
# )
tokenizer = AutoTokenizer.from_pretrained("Harley-ml/Dillionv2-1.3M")
License
Apache 2.0.
Citation
@misc{32d-vs-64d,
title = {32D-vs-64D Heads at Small Scale},
organization = {FromZero},
authors = {Paul Courneya},
year = {2026},
url = {https://huggingface.co/fromziro/32D-Vs-64D]
}