File size: 4,320 Bytes
1d5f5de 0a997cf 1d5f5de b644a66 1d5f5de | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 | ---
license: mit
tags:
- pytorch
- gpt
- transformer
- ablation-study
- attention
- rope
- mixture-of-experts
- linear-attention
datasets:
- HuggingFaceFW/fineweb-edu
---
# codemindv2: SLM Attention / Position / Optimizer Ablation Study
19 small (~63M-param) GPT-style language models, each changing exactly one
architectural axis from a shared baseline, trained on 1B tokens of
[fineweb-edu](https://huggingface.co/datasets/HuggingFaceFW/fineweb-edu)
(`sample-10BT`, GPT-2 tokenizer). Full paper, figures, and code:
https://github.com/B4K2/codemind_study.
<table><tr>
<td><img src="images/bench_vs_ppl.png" alt="Zero-shot capability vs. validation perplexity"></td>
<td><img src="images/pareto.png" alt="Perplexity vs. KV-cached decode speed Pareto trade-off"></td>
</tr></table>
Every subfolder (`m01/` ... `m19/`) is one model: `model.safetensors` (weights,
optimizer state stripped), `config.json` (architecture + training hyperparameters),
`eval_results.json` (perplexity / FLOPs / throughput / VRAM), and
`benchmark_results.json` (zero-shot LAMBADA / HellaSwag / PIQA / SciQ / ARC-Easy /
WinoGrande / BLiMP via `lm-eval`).
## Models
| ID | Attention | Position | Optimizer | FFN | Val PPL |
|-----|-----------|----------|-----------|-------------|---------|
| m01 | MHA | learned | AdamW | dense | 34.08 |
| m02 | MHA | learned | AdamW | MoE (8/2) | 33.48 |
| m03 | MQA | learned | AdamW | dense | 35.40 |
| m04 | GQA (4:1) | learned | AdamW | dense | 34.70 |
| m05 | MLA | learned | AdamW | dense | 38.28 |
| m06 | MHA | RoPE | AdamW | dense | 31.27 |
| m07 | MQA | RoPE | AdamW | dense | 32.26 |
| m08 | GQA (4:1) | RoPE | AdamW | dense | 32.10 |
| m09 | MLA | RoPE | AdamW | dense | 33.84 |
| m10 | Diff Attn | RoPE | AdamW | dense | 30.60 |
| m11 | MHA | RoPE | Muon | dense | 30.54 |
| m12 | GQA (4:1) | RoPE | Muon | dense | 32.09 |
| m13 | MHA | RoPE | Muon | dense + AttnRes | 37.14 |
| m14 | MHA | RoPE | Muon | MoE (8/2) | 31.07 |
| m15 | MHA | RoPE | Muon | Shared MoE | 32.07 |
| m16 | MHA | RoPE | Muon | Shared MoE (fat) | 31.71 |
| m17 | MHA | none | AdamW | dense | 35.69 |
| m18 | GQA (4:1) | none | AdamW | dense | 35.97 |
| m19 | KDA (linear) | RoPE | Muon | dense | 33.58 |
All models: `d_model=512`, `num_layers=12`, `num_heads=8`, `block_size=1024`,
1B training tokens. See each `config.json` for exact hyperparameters.
**Note on perplexity:** the paper's thesis is that PPL is not a reliable proxy
for downstream capability across attention variants — see `benchmark_results.json`
per model and the paper for zero-shot accuracy, which reorders several
comparisons versus the PPL table above.
## Loading a model
Weights are architecture-specific (custom PyTorch, not a `transformers`
`AutoModel`). Clone the code repo above, then:
```python
import json
import torch
from safetensors.torch import load_file
from models import build_model # dispatches on config["model_id"]
model_id = "m11"
config = json.load(open(f"{model_id}/config.json"))
state_dict = load_file(f"{model_id}/model.safetensors")
model = build_model(config)
model.load_state_dict(state_dict, strict=not config.get("tie_embeddings"))
model.eval()
```
`strict=False` is only needed for tied-embedding models (`lm_head.weight` was
dropped before saving since it's a view of `token_emb.emb.weight`; it's
re-materialized by `tie_embeddings` in the model's `__init__`).
M19 has no KV-cache path (`supports_kv_cache = False`) — its attention operator
(KDA, a linear-attention recurrence) carries a fixed-size state rather than a
growing K/V tensor, so `inference_tok_per_sec_kvcache` is `null` in its
`eval_results.json` and its `inference_tok_per_sec` (full re-forward) is not
comparable to the other models' cached numbers.
## Citation
See the paper (`paper/main.pdf` in the code repo) for the full ablation
methodology, comparison pairs, and analysis.
|