-
Notifications
You must be signed in to change notification settings - Fork 0
Experiment Results Mar 11 2026
Hardware: Apple M1 Max (32-core GPU, 64GB unified memory)
Backend: MLX
Program: program.md autonomous experiment loop
Time budget: 300s per experiment (+ 10 uncounted warmup steps)
25 experiments, 8 kept. val_bpb improved from 2.094 → 1.621 (−22.6%).
The dominant optimization lever was reducing batch size to maximize training steps within the fixed 5-minute time budget. Secondary gains came from narrowing the MLP to trade per-step capacity for throughput.
Optimization Curve
| Exp | val_bpb | Δ | Steps | Description |
|---|---|---|---|---|
| baseline | 2.094 | — | 14 | Original config (batch 131K, depth 8, MLP 4x) |
| exp1 | 1.854 | −0.241 | 87 | Batch 131K→32K, dev_batch 32→16, warmup 10% |
| exp3 | 1.750 | −0.104 | 97 | Remove warmup, cooldown 0.5→0.7 |
| exp4 | 1.673 | −0.077 | 175 | Batch 32K→16K, dev_batch 16→8 |
| exp13 | 1.668 | −0.005 | 183 | MLP ratio 4x→3x |
| exp16 | 1.654 | −0.014 | 186 | MLP ratio 3x→2.5x |
| exp17 | 1.622 | −0.032 | 205 | MLP ratio 2.5x→2x |
| exp18 | 1.621 | −0.001 | 210 | MLP ratio 2x→1.5x |
| # | Commit | val_bpb | Mem (GB) | Status | Description |
|---|---|---|---|---|---|
| 0 | 0dc8208 | 2.094 | 94.3 | keep | Baseline (no changes) |
| 1 | b24dd14 | 1.854 | 25.9 | keep | Batch 32768, dev_batch 16, warmup 0.1 |
| 2 | dcf27b8 | 2.420 | 52.7 | discard | Depth 12 (too slow, too few counted steps) |
| 3 | a085cec | 1.750 | 25.9 | keep | No warmup, cooldown 0.7 |
| 4 | 2298d5f | 1.673 | 13.2 | keep | Batch 16384, dev_batch 8 |
| 5 | f86db47 | 1.743 | 7.0 | discard | Batch 8192 (too noisy, fewer tokens) |
| 6 | e444549 | 1.970 | 13.2 | discard | Halve LRs (too low, under-trains) |
| 7 | e58db13 | 1.961 | 19.6 | discard | Depth 10 (too slow, fewer steps) |
| 8 | 11f05ea | 1.728 | 13.2 | discard | MATRIX_LR 0.06 (slightly worse) |
| 9 | a0465e3 | 1.706 | 13.2 | discard | Cooldown 0.5 (0.7 was better) |
| 10 | a7d9238 | 1.735 | 13.2 | discard | Cooldown 0.85 (too much, 0.7 optimal) |
| 11 | 9aa3b5b | 1.714 | 13.2 | discard | Weight_decay 0.05 (0.2 was better) |
| 12 | e9c2ca8 | 1.704 | 12.9 | discard | GQA n_kv=n_head//2 (fewer params hurts) |
| 13 | c1d62fd | 1.668 | 12.4 | keep | MLP ratio 3x |
| 14 | afbe72c | 1.812 | 17.5 | discard | Wider model 768-dim (too slow) |
| 15 | b9111f3 | 1.698 | 12.4 | discard | EMBEDDING_LR 1.0 (worse) |
| 16 | 6a41e2c | 1.654 | 12.1 | keep | MLP ratio 2.5x |
| 17 | 10c7fd8 | 1.622 | 11.7 | keep | MLP ratio 2.0x |
| 18 | c3605da | 1.621 | 11.3 | keep | MLP ratio 1.5x |
| 19 | 0196fce | 1.776 | 11.4 | discard | No value embeds (too much capacity loss) |
| 20 | 5f0b644 | 1.802 | 15.8 | discard | Depth=9 MLP=2x (too slow) |
| 21 | 76bf2e6 | 1.719 | 16.3 | discard | HEAD_DIM=64 (slower, worse) |
| 22 | 3319a0b | 1.632 | 11.3 | discard | Softcap 10 (too restrictive) |
| 23 | cac38ae | 1.706 | 11.3 | discard | Softcap 20 (worse + slower) |
| 24 | 716cac2 | 1.750 | 11.3 | discard | No softcap (much worse) |
-
Smaller batch size (−0.421 bpb, 89% of total gain): The original batch size of 131K gave only 14 total steps (~4 counted). Reducing to 16K gave 175+ steps, dramatically more training in the same wall-clock budget. This was the single biggest lever.
-
Narrower MLP (−0.052 bpb): Reducing MLP expansion from 4x to 1.5x traded per-layer capacity for faster step times, netting ~30 more steps and a small quality improvement.
-
Cooldown ratio 0.7: Longer LR decay than the default 0.5 helped final convergence. Both shorter (0.5) and longer (0.85) were worse.
- Depth increases (10, 12): Always net-negative. Slower steps outweighed capacity gains.
- LR tuning (higher or lower): Original MATRIX_LR=0.04 and EMBEDDING_LR=0.6 were optimal.
- Weight decay changes: 0.2 was already well-tuned.
- GQA: Fewer KV heads saved params but hurt quality.
- HEAD_DIM=64: More heads but slower overall.
- Softcap tuning: 15 was optimal; 10 too restrictive, 20 and removal both worse.
- Removing value embeddings: Major quality regression despite freeing 33% of params.
- Warmup: Wasted counted steps at low LR with no benefit.
On M1 Max with MLX and a 5-minute time budget, throughput dominates quality. Any change that slows per-step time—even if it adds model capacity—is net-negative because the model can't train long enough to exploit that capacity. The optimal strategy is the smallest, fastest model that can still learn effectively.
# Winning hyperparameters (val_bpb = 1.621) TOTAL_BATCH_SIZE = 16384 DEVICE_BATCH_SIZE = 8 DEPTH = 8 ASPECT_RATIO = 64 # → model_dim=512, 4 heads HEAD_DIM = 128 MLP_RATIO = 1.5 EMBEDDING_LR = 0.6 UNEMBEDDING_LR = 0.004 MATRIX_LR = 0.04 SCALAR_LR = 0.5 WEIGHT_DECAY = 0.2 ADAM_BETAS = (0.8, 0.95) WARMUP_RATIO = 0.0 WARMDOWN_RATIO = 0.7 SOFTCAP = 15
Model: 39.8M params, ~210 steps in 5 min, 11.3 GB peak memory