README.md

---
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.

![Training and validation loss curves for all 19 models](images/loss_curves_all.png)

<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.