---
license: apache-2.0
language:
- en
datasets:
- nampdn-ai/mini-fineweb
tags:
- llama
- ablation
- heads
- slms
- dillionv2-tokenizer
- llms
---

# 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`](https://huggingface.co/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`](https://huggingface.co/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:

```python
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]
}
```