update card: fix HF paths, prune-stale text, add hybrid ablation + 61.2 winner

README.md

@@ -39,27 +39,56 @@ novel architecture.
 
 ## What's in this repo
 
-Six training runs, each in its own folder. Within a folder you'll find periodic
-snapshots `step_<N>.ckpt` (every 5000 steps; every 3000 for the 50M run) and
-`last.ckpt` (the final-step weights). These are **PyTorch Lightning
-checkpoints** from the MDLM codebase — they bundle model weights *and* EMA
-shadow parameters (EMA decay 0.9999), optimizer state, and config. They are
-**not** `transformers`-loadable via `from_pretrained`; load them with the
-training repo (see *How to use* below).
-
-| Folder | Backbone | Params | LR | Steps | Tokens | Val PPL ↓ |
+Each training run lives under `runs/<name>/checkpoints/` and holds two files:
+`best.ckpt` (lowest validation loss — the weights the PPL table below reports)
+and `last.ckpt` (final-step weights, for resuming). Periodic `step=<N>.ckpt`
+snapshots are written during training but pruned from this repo to save space,
+so only `best`/`last` are hosted. These are **PyTorch Lightning checkpoints**
+from the MDLM codebase — they bundle model weights *and* EMA shadow parameters
+(EMA decay 0.9999), optimizer state, and config. They are **not**
+`transformers`-loadable via `from_pretrained`; load them with the training repo
+(see *How to use* below).
+
+| HF path `runs/…` | Backbone | Params | LR | Steps | Tokens | Val PPL ↓ |
 |--------|----------|:------:|:--:|:-----:|:------:|:---------:|
-| `runB_transformer_130m`  | Transformer (DiT) | \~130M | 3e-4 | 76k | \~5B | **70.5** |
-| `runD_130m_seed1`        | BiMamba-2 (SSM)   | \~130M | 3e-4 | 76k | \~5B | 85.9 |
-| `runD_130m_seed2`        | BiMamba-2 (SSM)   | \~130M | 3e-4 | 76k | \~5B | 83.5 |
-| `runD_130m_lr1e3_seed1`  | BiMamba-2 (SSM)   | \~130M | **1e-3** | 76k | \~5B | **79.3** |
-| `scaling_100m`           | BiMamba-2 (SSM)   | \~100M | 3e-4 | 60k* | \~4B | 97.5 |
-| `scaling_50m`            | BiMamba-2 (SSM)   | \~50M  | 3e-4 | 30k | \~2B | 136.3 |
+| `runB`        | Transformer (DiT) | \~130M | 3e-4 | 76k | \~5B | **70.5** |
+| `runD1`       | BiMamba-2 (SSM)   | \~130M | 3e-4 | 76k | \~5B | 85.9 |
+| `runD2`       | BiMamba-2 (SSM)   | \~130M | 3e-4 | 76k | \~5B | 83.5 |
+| `runD_lr1e3`  | BiMamba-2 (SSM)   | \~130M | **1e-3** | 76k | \~5B | **79.3** |
+| `s100`        | BiMamba-2 (SSM)   | \~100M | 3e-4 | 60k* | \~4B | 97.5 |
+| `s50`         | BiMamba-2 (SSM)   | \~50M  | 3e-4 | 30k | \~2B | 136.3 |
 
 Val PPL = MDLM ELBO-bound validation perplexity on the OpenWebText validation
-split, measured on each run's **best EMA checkpoint** (lower is better).
-`*` the 100M run's valid final checkpoint is `step_60000`/`last.ckpt`
-(see the GitHub report for why 61k looped).
+split, measured on each run's **`best.ckpt`** (EMA weights; lower is better).
+`*` the 100M run's valid final checkpoint is `last.ckpt` at step 60000
+(see the GitHub report for why 61k looped). Also hosted (report §6.5 LR sweep,
+not shown above): `s50_lr5e4`, `s50_lr1e3`, `s50_lr2e3` — the 50M learning-rate
+sweep.
+
+### Hybrid Mamba+attention checkpoints (attention-layout ablation)
+
+The hybrid backbone inserts full bidirectional attention among the Mamba blocks.
+These runs ablate **how many** attention layers and **where**, then over-train
+the winner. All share the 130M / OpenWebText / lr-3e-4 recipe and differ only in
+the attention layout. Val PPL is on each run's final-step `last.ckpt`, matching
+report §11.5 (grid at 76k; the winner over-trained to 150k).
+
+| HF path `runs/…` | Attention layers | # | Steps | Val PPL ↓ | `+experiment=` |
+|--------|--------|:--:|:--:|:--:|--------|
+| `hybrid_130m` | `[3,7,11]`   | 3 | 76k  | 69.5 | `hybrid_130m` |
+| **`hyb_e3`**  | `[2,5,8,11]` | 4 | **150k** | **60.9** ← best | `hyb_e3` |
+| `hyb_e3_s2`   | `[2,5,8,11]` | 4 | 150k | 61.5 (seed 2) | `hyb_e3` |
+| `hyb_e6`      | `[5,11]`     | 2 | 76k  | 71.4 | `hyb_e6` |
+| `hyb_e12`     | `[11]`       | 1 | 76k  | 75.7 | `hyb_e12` |
+| `hyb_early`   | `[0,1,2]`    | 3 | 76k  | 80.8 | `hyb_early` |
+| `hyb_mid`     | `[4,5,6]`    | 3 | 76k  | 73.9 | `hyb_mid` |
+| `hyb_late`    | `[9,10,11]`  | 3 | 76k  | 73.0 | `hyb_late` |
+
+At matched 76k, more attention lowers PPL (4 → 68.1, 3 → 69.5, 2 → 71.4, 1 → 75.7)
+and **placement matters more than count** — distributed `[3,7,11]` (69.5) beats
+every clustered layout, and clustering early is catastrophic (80.8). The winner
+`hyb_e3` (`[2,5,8,11]`) over-trained to 150k reaches **61.2 mean ±0.3** (seeds
+60.9 / 61.5) — the best quality in this study, at ~2× the 76k compute.
 
 ### Results at a glance
 
@@ -67,7 +96,7 @@ split, measured on each run's **best EMA checkpoint** (lower is better).
   tokens, the Transformer denoiser (70.5) is modestly but consistently stronger
   than pure BiMamba-2. BiMamba prefers a **\~3.3× higher learning rate**; a
   50M LR sweep found `1e-3` best, and retraining 130M at `1e-3` (the
-  `runD_130m_lr1e3_seed1` checkpoints) closes **\~43%** of the gap (85.9 → 79.3)
+  `runD_lr1e3` checkpoints) closes **\~43%** of the gap (85.9 → 79.3)
   but does not close it.
 - **Scaling** (BiMamba, lr 3e-4): 50M → 136.3, 100M → 97.5, 130M → 84.7 —
   clean, monotonic, seed-stable (Δ≈2.4 between seeds).
@@ -77,6 +106,19 @@ split, measured on each run's **best EMA checkpoint** (lower is better).
 - **Honest finding:** *pure* BiMamba-2 trades quality for long-context
   throughput — consistent with DiffuApriel, where a *hybrid* Mamba+attention
   model is what recovers quality.
+- **Hybrid backbone + attention ablation (best quality here).** Inserting sparse
+  bidirectional attention recovers DiT-class quality: at matched 130M / 76k /
+  lr 3e-4, **3 of 12 layers as attention** (`hybrid_130m`, `[3,7,11]`) reaches
+  **69.5** — matching the DiT (70.5). An attention-layout ablation (grid runs
+  `hyb_e3`/`hyb_e6`/`hyb_e12`/`hyb_early`, all hosted here at 76k) shows
+  **placement matters more than count**: distribute attention through depth
+  (clustering it early is catastrophic, 80.8), and **4 evenly-spread layers**
+  (`hyb_e3`, `[2,5,8,11]`) is best (68.1 at 76k). Over-training that winner to
+  150k steps reaches **61.2** val PPL (2-seed mean ±0.3) — the strongest result
+  here, though at ~2× the compute of the matched table above (not a
+  matched-compute claim vs the 70.5 DiT). All hybrid checkpoints are hosted (see
+  the table above); full grid + over-train detail:
+  [report §11.5](https://github.com/shivnarainms22/DiffMamba/blob/master/DiffMamba_VLM_Report.md).
 
 Full numbers, caveats, and the LR-fairness analysis are in the
 [technical report on GitHub](https://github.com/shivnarainms22/DiffMamba/blob/master/DiffMamba_Report.md).
@@ -113,24 +155,26 @@ git clone https://github.com/shivnarainms22/DiffMamba
 cd DiffMamba
 # set up the environment (see requirements.yaml / scripts/)
 
-# download a checkpoint, e.g. the LR-tuned BiMamba-130M
+# download a checkpoint, e.g. the LR-tuned BiMamba-130M (best.ckpt reproduces
+# the table PPL; --local-dir preserves the runs/…/checkpoints/ path)
 huggingface-cli download Shiv-22/diffmamba-checkpoints \
-  runD_130m_lr1e3_seed1/last.ckpt --local-dir ./ckpts
+  runs/runD_lr1e3/checkpoints/best.ckpt --local-dir ./ckpts
 
 # validation perplexity (EMA), matching the table above
 python main.py mode=ppl_eval +experiment=runD_130m \
-  eval.checkpoint_path=./ckpts/runD_130m_lr1e3_seed1/last.ckpt \
+  eval.checkpoint_path=./ckpts/runs/runD_lr1e3/checkpoints/best.ckpt \
   data.cache_dir=<path>/data loader.eval_batch_size=32
 
 # generate samples
 python main.py mode=sample_eval +experiment=runD_130m \
-  eval.checkpoint_path=./ckpts/runD_130m_lr1e3_seed1/last.ckpt \
+  eval.checkpoint_path=./ckpts/runs/runD_lr1e3/checkpoints/best.ckpt \
   loader.eval_batch_size=4
 ```
 
-Use the matching `+experiment=` for each folder: `runD_130m` (BiMamba-130M and
-its LR-tuned variant), `runB_transformer_130m` (DiT-130M), `scaling_100m`,
-`scaling_50m`.
+Use the matching `+experiment=` for each run: `runD_130m` (BiMamba-130M `runD1`/
+`runD2` and the LR-tuned `runD_lr1e3`), `runB_transformer_130m` (DiT `runB`),
+`scaling_100m` (`s100`), `scaling_50m` (`s50`). For the hybrid runs use
+`hybrid_130m` / `hyb_e3` / `hyb_e6` / `hyb_e12` (see below).
 
 ## Limitations