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, Jonathon LY, User110},
  year         = {2026},
  url          = {https://huggingface.co/fromziro/32D-Vs-64D]
}