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
(sample-10BT, GPT-2 tokenizer). Full paper, figures, and code:
https://github.com/B4K2/codemind_study.
 |
 |
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:
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.