Add degradation experiments: one-parameter-at-a-time ablations

Controlled study of how the model degrades. Four variants, each = v1 with a
single field replaced (dataclasses.replace), so differences are attributable:
- small: capacity down (2L/64d, ~0.38M params)
- short: training time down (max_iters 20k->2k)
- ctx:   context window down (256->64)
- data:  data quantity down (max_train_tokens->1M; new TrainConfig knob + train.py slice)

- scripts/compare.py: sample same prompt across all trained configs with a
  shared seed, print side by side, write reports/compare.md
- tests/test_configs.py: enforces one-parameter-at-a-time (only intended fields
  differ from v1) + small param count (3 tests). Full suite: 29 passing.

README.md

@@ -166,6 +166,38 @@ experiment-llm/
       (v1: grammatical, story-arc coherent; plot logic wobbles as expected)
 - [x] Every phase re-runnable from a clean clone with documented commands
 
+## Degradation experiments (iteration 1.5)
+
+A controlled study of *how* a model gets worse: start from the v1 baseline and
+change exactly **one** variable per run, so any difference in the output is
+attributable to that variable. Each variant config is literally `v1` with one
+field replaced (see `configs/*.py`); `tests/test_configs.py` enforces that only
+the intended field changes.
+
+| config | one change from v1 | isolates | expected failure mode |
+|---|---|---|---|
+| `v1` | — (baseline) | — | grammatical, coherent story arcs |
+| `small` | 2L/64d, ~0.38M params | capacity | generic, repetitive, forgets rare words |
+| `short` | `max_iters` 20k→2k | training time | half-learned grammar, weaker structure |
+| `ctx` | `context_length` 256→64 | long-range coherence | fine locally, plot threads break sooner |
+| `data` | `max_train_tokens`→1M | data quantity | overfits: memorizes, novel prompts collapse |
+
+Run each (they all train faster than v1; `short` is ~minutes):
+
+```sh
+python src/train.py --config small     # then short, ctx, data
+python scripts/plot_loss.py --config small
+```
+
+Then compare generated text across every trained variant, side by side (same
+prompt, same RNG seed, so differences reflect the model):
+
+```sh
+python scripts/compare.py                       # default prompt, all configs
+python scripts/compare.py --suite               # every suite prompt -> reports/compare.md
+python scripts/compare.py --configs v1,small    # a subset
+```
+
 ## Iteration 2 candidates (out of scope for now)
 
 - Scale toward the coherence sweet spot (~15–30M params, longer training)

configs/ctx.py

@@ -0,0 +1,17 @@
+"""Degradation experiment: same as v1, but a smaller context window.
+
+Isolates *long-range coherence*. Only context_length changes from v1
+(256 -> 64). Local fluency should survive; callbacks and plot threads that span
+more than a few sentences should break sooner.
+
+Honest caveat: shrinking context also lowers tokens/step (batch * context), so
+this run sees fewer total tokens than v1 over the same number of steps. The
+qualitative coherence effect is still the dominant, visible change.
+"""
+
+from dataclasses import replace
+
+import v1
+
+model = replace(v1.model, context_length=64)
+train = v1.train

configs/data.py

@@ -0,0 +1,17 @@
+"""Degradation experiment: same as v1, but trained on a tiny slice of data.
+
+Isolates *data quantity*. Only max_train_tokens changes from v1: the model sees
+~1M training tokens instead of 555M. Expect overfitting — train loss falls while
+val loss (measured on the full held-out set) climbs, and novel prompts collapse
+as the model regurgitates memorized stories.
+
+Tip: this overfits within a few thousand steps. Watch the val loss turn upward
+in the log and stop early (Ctrl-C) rather than running all 20,000 steps.
+"""
+
+from dataclasses import replace
+
+import v1
+
+model = v1.model
+train = replace(v1.train, max_train_tokens=1_000_000)

configs/short.py

@@ -0,0 +1,13 @@
+"""Degradation experiment: same as v1, but far fewer training steps.
+
+Isolates *training time*. Only max_iters changes from v1 (20,000 -> 2,000).
+The cosine LR schedule anneals fully within these 2,000 steps, so this is the
+best model achievable in that budget — not v1 halted mid-schedule.
+"""
+
+from dataclasses import replace
+
+import v1
+
+model = v1.model
+train = replace(v1.train, max_iters=2000)

configs/small.py

@@ -0,0 +1,14 @@
+"""Degradation experiment: same as v1, but a much smaller model.
+
+Isolates *capacity*. Only the model dimensions change; context length and the
+training schedule match v1 exactly. ~0.38M params vs v1's 4.26M — and only
+~0.1M of those are non-embedding (the 4,096-token embedding table dominates a
+model this small, so the actual 'reasoning' capacity shrinks even harder).
+"""
+
+from dataclasses import replace
+
+import v1
+
+model = replace(v1.model, n_layers=2, d_model=64, n_heads=2, d_ff=256)
+train = v1.train

configs/v1.py

@@ -25,6 +25,7 @@ class ModelConfig:
 class TrainConfig:
     # Data / batching. tokens per step = batch_size * context_length = 16,384.
     batch_size: int = 64
+    max_train_tokens: int | None = None  # None = whole train set; set to starve data
 
     # AdamW optimiser.
     learning_rate: float = 6e-4  # peak LR (reached after warmup)

scripts/compare.py

@@ -0,0 +1,90 @@
+"""Compare generated text across the degradation experiments.
+
+Samples the same prompt(s) from every config that has a trained checkpoint,
+using the same RNG seed for each so differences reflect the model, not sampling
+luck. Prints the results side by side and writes reports/compare.md.
+
+Usage:
+    python scripts/compare.py                              # default prompt, all configs
+    python scripts/compare.py --prompt "The dog and the cat"
+    python scripts/compare.py --suite                      # every suite prompt
+    python scripts/compare.py --configs v1,small,ctx       # a subset
+"""
+
+import argparse
+import sys
+from pathlib import Path
+
+import torch
+
+ROOT = Path(__file__).resolve().parent.parent
+sys.path.insert(0, str(ROOT / "src"))
+
+from sample import (  # noqa: E402
+    SUITE_PROMPTS,
+    TOK_PATH,
+    load_model,
+    pick_device,
+    sample_text,
+)
+from tokenizer import BPETokenizer  # noqa: E402
+
+# Baseline first, then the four one-parameter degradations.
+DEFAULT_CONFIGS = ["v1", "small", "short", "ctx", "data"]
+
+LABELS = {
+    "v1": "v1 (baseline)",
+    "small": "small (less capacity)",
+    "short": "short (less training)",
+    "ctx": "ctx (smaller context)",
+    "data": "data (less data)",
+}
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--configs", default=",".join(DEFAULT_CONFIGS))
+    ap.add_argument("--prompt", default="Once upon a time,")
+    ap.add_argument("--suite", action="store_true")
+    ap.add_argument("--max-new-tokens", type=int, default=200)
+    ap.add_argument("--temperature", type=float, default=0.8)
+    ap.add_argument("--top-k", type=int, default=200)
+    ap.add_argument("--seed", type=int, default=1234)
+    args = ap.parse_args()
+
+    configs = [c.strip() for c in args.configs.split(",") if c.strip()]
+    prompts = SUITE_PROMPTS if args.suite else [args.prompt]
+
+    device = pick_device()
+    tok = BPETokenizer.load(TOK_PATH)
+
+    # Load each available model once; warn about any not yet trained.
+    models = {}
+    for name in configs:
+        if (ROOT / "checkpoints" / name / "ckpt.pt").exists():
+            models[name] = load_model(name, device)
+        else:
+            print(f"[skip {name}: no checkpoint — train it first]")
+    if not models:
+        sys.exit("no trained checkpoints found among: " + ", ".join(configs))
+
+    out = ["# Degradation comparison", ""]
+    for prompt in prompts:
+        print(f"\n########## prompt: {prompt!r} ##########")
+        out += [f"## Prompt: {prompt}", ""]
+        for name, model in models.items():
+            torch.manual_seed(args.seed)  # same RNG state per model = fair compare
+            text = sample_text(model, tok, prompt, device,
+                               args.max_new_tokens, args.temperature, args.top_k)
+            label = LABELS.get(name, name)
+            print(f"\n----- {label} -----\n{text}")
+            out += [f"### {label}", "", f"> {text}", ""]
+
+    out_path = ROOT / "reports" / "compare.md"
+    out_path.parent.mkdir(exist_ok=True)
+    out_path.write_text("\n".join(out))
+    print(f"\nsaved {out_path}")
+
+
+if __name__ == "__main__":
+    main()

src/train.py

@@ -116,6 +116,10 @@ def main() -> None:
     tokens_dir = ROOT / "data" / "tokens"
     train_data = load_tokens(tokens_dir / "train.bin")
     val_data = load_tokens(tokens_dir / "val.bin")
+    # Data-starvation experiments restrict training to a slice; val stays full so
+    # the generalization gap (overfitting) is visible.
+    if tc.max_train_tokens is not None:
+        train_data = train_data[: tc.max_train_tokens]
     print(f"train tokens: {len(train_data):,} | val tokens: {len(val_data):,}")
 
     model = GPT(mc).to(device)

tests/test_configs.py

@@ -0,0 +1,72 @@
+"""Guardrail: each degradation config changes exactly one thing from v1.
+
+Enforces the controlled-experiment discipline — if a variant accidentally
+drifts from the baseline in an unintended field, this fails. Forward-only.
+Run: `.venv/bin/python tests/test_configs.py`
+"""
+
+import importlib
+import sys
+from dataclasses import asdict
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parent.parent
+sys.path.insert(0, str(ROOT / "src"))
+sys.path.insert(0, str(ROOT / "configs"))
+
+from model import GPT  # noqa: E402
+
+# For each variant: which model fields and which train fields may differ from v1.
+EXPECTED_DIFFS = {
+    "small": ({"n_layers", "d_model", "n_heads", "d_ff"}, set()),
+    "short": (set(), {"max_iters"}),
+    "ctx": ({"context_length"}, set()),
+    "data": (set(), {"max_train_tokens"}),
+}
+
+
+def _diff_keys(a, b):
+    return {k for k in a if a[k] != b[k]}
+
+
+def test_each_variant_changes_only_intended_fields():
+    v1 = importlib.import_module("v1")
+    base_model, base_train = asdict(v1.model), asdict(v1.train)
+    for name, (exp_model, exp_train) in EXPECTED_DIFFS.items():
+        cfg = importlib.import_module(name)
+        model_diffs = _diff_keys(asdict(cfg.model), base_model)
+        train_diffs = _diff_keys(asdict(cfg.train), base_train)
+        assert model_diffs == exp_model, f"{name}: model diffs {model_diffs} != {exp_model}"
+        assert train_diffs == exp_train, f"{name}: train diffs {train_diffs} != {exp_train}"
+
+
+def test_each_variant_builds_a_model():
+    for name in EXPECTED_DIFFS:
+        cfg = importlib.import_module(name)
+        model = GPT(cfg.model)
+        assert model.num_params() > 0
+
+
+def test_small_param_count():
+    cfg = importlib.import_module("small")
+    model = GPT(cfg.model)
+    # tok_emb 4096*64 + pos_emb 256*64
+    # + 2*(qkv 64*192 + proj 64*64 + 2 norms*64 + mlp_fc 64*256 + mlp_proj 256*64)
+    # + final norm 64
+    expected = 4096 * 64 + 256 * 64 + 2 * (64 * 192 + 64 * 64 + 2 * 64 + 64 * 256 + 256 * 64) + 64
+    assert expected == 377_152
+    assert model.num_params() == expected
+    # The embedding table dominates: non-embedding capacity is tiny.
+    assert model.num_params(non_embedding=True) == 98_624
+
+
+def main() -> None:
+    tests = [v for k, v in sorted(globals().items()) if k.startswith("test_")]
+    for t in tests:
+        t()
+        print(f"ok  {t.__name__}")
+    print(f"\n{len(tests)} passed")
+
+
+if __name__ == "__main__":
+    main()