updated model card with ablation results and all 4 runs

README.md

@@ -21,13 +21,15 @@ pipeline_tag: text-generation
 
 Checkpoints from a research project studying expert specialization in Mixture-of-Experts models. I fine-tuned GPT-2 small on three domains -- code, math, and prose -- to see whether experts naturally specialize by domain when given the right routing incentives.
 
-Short answer: they do. MoE beats the dense baseline by 3.6% overall and 14% on math, with zero expert collapse across 10,000 training steps.
+Short answer: they do. MoE beats the dense baseline by 3.6% overall and 14% on math, with zero expert collapse across 10,000 training steps. Two ablation runs confirmed that load balancing loss is essential (without it, one expert captures 73.6% of tokens by step 500) and that top-2 routing provides negligible improvement over top-1.
 
 ---
 
 ## 1. Results
 
-| Metric | Dense Baseline | MoE (8 experts, top-1) | Delta |
+### Main comparison
+
+| Metric | Dense Baseline | MoE (top-1) | Delta |
 |---|---|---|---|
 | eval/loss | 2.157 | 2.080 | -3.6% |
 | loss_code | 1.554 | 1.521 | -2.1% |
@@ -37,6 +39,13 @@ Short answer: they do. MoE beats the dense baseline by 3.6% overall and 14% on m
 
 Math benefits the most from expert routing. Prose is the one domain where dense wins; diverse web text doesn't lend itself to clean expert specialization. The MoE model crossed the dense baseline at step ~3,600 (36% of training).
 
+### Ablations
+
+| Run | What it tests | Result |
+|---|---|---|
+| No-LB ablation | Remove load balancing loss (`lb_coef=0.0`) | Expert collapse at step 500. Single expert handles 73.6% of tokens. Z-loss alone doesn't prevent it. |
+| Top-2 directional | Route to 2 experts instead of 1 | eval/loss=2.077 vs top-1's 2.080, which is a 0.14% difference. Not worth 2x expert compute. |
+
 ---
 
 ## 2. Files
@@ -45,14 +54,28 @@ Math benefits the most from expert routing. Prose is the one domain where dense
 dense-baseline/
 ├── final-model.safetensors    # 622 MB -- dense GPT-2, 124M params
 ├── final-model.json           # metadata sidecar (config, metrics)
-├── ckpt-step-4999.pt          # 1.4 GB -- full resume checkpoint (optimizer, scheduler, RNG)
+├── ckpt-step-4999.pt          # 1.4 GB -- full resume checkpoint
 └── metrics.jsonl              # per-step training + eval metrics
 
 moe-main/
 ├── final-model.safetensors    # 1.1 GB -- MoE GPT-2, 257M params (8 experts × 4 layers)
-├── final-model.json           # metadata sidecar (config, metrics)
+├── final-model.json           # metadata sidecar
 ├── ckpt-step-9999.pt          # 2.9 GB -- full resume checkpoint
-└── metrics.jsonl              # per-step training + eval metrics
+└── metrics.jsonl
+
+no-lb-ablation/
+├── final-model.safetensors    # 1.1 GB -- collapsed MoE model at step 500
+├── best-model.safetensors     # 1.1 GB -- best eval loss (step 400, pre-collapse)
+├── ckpt-step-500.pt           # 2.9 GB -- full resume checkpoint
+├── config.json, run_summary.json
+└── metrics.jsonl
+
+top2-main-10k/
+├── final-model.safetensors    # 1.2 GB -- top-2 MoE model at step 9999
+├── best-model.safetensors     # 1.2 GB -- best eval loss (step 8000)
+├── ckpt-step-9999.pt          # 2.9 GB -- full resume checkpoint
+├── config.json, run_summary.json
+└── metrics.jsonl
 ```
 
 The `.safetensors` files are the trained model weights. The `.pt` files contain the full training state for resuming runs (optimizer, LR scheduler, RNG states). The `.json` sidecars store architecture config and final eval metrics.
@@ -119,7 +142,7 @@ Routing uses the Straight-Through Estimator. Forward pass routes to one expert w
 
 ## 5. Training
 
-Both models trained on ~6.6M tokens across three domains, balanced to equal token counts:
+All models trained on ~6.6M tokens across three domains, balanced to equal token counts:
 
 | Domain | Source | Size |
 |---|---|---|
@@ -129,17 +152,18 @@ Both models trained on ~6.6M tokens across three domains, balanced to equal toke
 
 Training config:
 
-| Parameter | Dense | MoE |
-|---|---|---|
-| Max steps | 5,000 | 10,000 |
-| Batch size | 2 × 4 grad accum = 8 | 2 × 4 grad accum = 8 |
-| Block size | 512 | 512 |
-| Learning rate | 5e-5 | 5e-5 |
-| Warmup | 10% | 10% |
-| Schedule | Cosine | Cosine |
-| Hardware | 1× RTX 4090 24GB | 1× RTX 4090 24GB |
-| Wall time | ~30 min | ~85 min |
-| Throughput | ~25.7k tok/s | ~14.2k tok/s |
+| Parameter | Dense | MoE (top-1) | MoE (top-2) | No-LB |
+|---|---|---|---|---|
+| Max steps | 5,000 | 10,000 | 10,000 | 2,000 (early-stopped at 500) |
+| Batch size | 8 | 8 | 8 | 8 |
+| Block size | 512 | 512 | 512 | 512 |
+| Learning rate | 5e-5 | 5e-5 | 5e-5 | 5e-5 |
+| lb_coef | — | 0.01 | 0.01 | 0.0 |
+| noise_std | — | 0.1 | 0.1 | 0.0 |
+| Hardware | 1× RTX 4090 | 1× RTX 4090 | 1× RTX 4090 | 1× RTX 4090 |
+| Wall time | ~30 min | ~85 min | ~48 min | ~5 min |
+
+Total GPU cost for all 4 runs: ~$2.79 (including setup/idle overhead).
 
 ---
 
@@ -149,13 +173,15 @@ Training curves are on Weights & Biases:
 
 - [Dense baseline](https://wandb.ai/sumit-ml/moe-emergence/runs/fqhfblfv)
 - [MoE main run](https://wandb.ai/sumit-ml/moe-emergence/runs/j08s2d1m)
+- [No-LB ablation](https://wandb.ai/sumit-ml/moe-emergence/runs/06pljhrv)
+- [Top-2 directional](https://wandb.ai/sumit-ml/moe-emergence/runs/6mw6qbac)
 
 ---
 
 ## 7. Links
 
 - Code: [github.com/sumitdotml/moe-emergence](https://github.com/sumitdotml/moe-emergence)
-- Experiment docs: [run-004 (dense)](https://github.com/sumitdotml/moe-emergence/blob/main/docs/experiments/run-004-dense-baseline.md), [run-005 (MoE)](https://github.com/sumitdotml/moe-emergence/blob/main/docs/experiments/run-005-moe-main.md)
+- Experiment docs: [run-004 (dense)](https://github.com/sumitdotml/moe-emergence/blob/main/docs/experiments/run-004-dense-baseline.md), [run-005 (MoE)](https://github.com/sumitdotml/moe-emergence/blob/main/docs/experiments/run-005-moe-main.md), [run-006 (no-LB)](https://github.com/sumitdotml/moe-emergence/blob/main/docs/experiments/run-006-no-lb-ablation.md), [run-007 (top-2)](https://github.com/sumitdotml/moe-emergence/blob/main/docs/experiments/run-007-top2-directional.md)
 
 ---