Upload README.md with huggingface_hub

README.md

@@ -10,23 +10,56 @@ tags:
 
 # EDLM External Init Checkpoints
 
-External MDLM and AR baseline checkpoints used to seed EDLM-Soft warmup runs in the draft-refine project.
+External MDLM and AR baseline checkpoints used to seed EDLM-Soft warmup runs in the draft-refine project, plus the best-trained ckpt produced by the inference + training improvements documented below.
+
+## Files
 
 | File | Size | Arch | Steps | Notes |
 |---|---|---|---|---|
-| `mdlm.ckpt` | 2.6 GB | DiT (custom; `models/dit.py`) — `sigma_map.mlp.*` + `norm1` per block | 1,112,551 (epoch 67) | Lightning ckpt; full state including optimizer + scheduler. **NOT the same arch as `kuleshov-group/mdlm-owt`** which uses `adaLN_modulation`. |
-| `ar.ckpt`   | 2.4 GB | AR baseline (no `sigma_map`) | 1,060,000 (epoch 72) | Lightning ckpt. 101 backbone keys; missing time-conditioning module relative to MDLM. Partial init when loaded into a DiT model via `strict=False`. |
+| `mdlm.ckpt` | 2.7 GB | DiT (custom; `models/dit.py`) — `sigma_map.mlp.*` + `norm1` per block | 1,112,551 (epoch 67) | Drive Lightning ckpt; full state. **NOT the same arch as `kuleshov-group/mdlm-owt`** which uses `adaLN_modulation`. |
+| `ar.ckpt`   | 2.6 GB | AR baseline (no `sigma_map`) | 1,060,000 (epoch 72) | Drive Lightning ckpt. 101 backbone keys. Partial init when loaded into a DiT model via `strict=False`. |
+| **`edlm-step2k-pure-mdlm.ckpt`** | 2.7 GB | DiT (HF kuleshov-group/mdlm-owt) + tiny_queries scorer head | 2000 (post-pretrained) | **Best ckpt**: 2000 steps of pure MDLM CE training on top of `kuleshov-group/mdlm-owt`. Achieves PPL=10.76 with `uncommitted_soft @ 256 steps` sampler (paper-aligned 128 samples eval). |
 
-## Compatibility
+## Headline result
 
-These ckpts produce keys with a `backbone.` prefix matching the local `models/dit.py` in the draft-refine `external/edlm/` codebase (with `backbone: dit` config — NOT `hf_dit`).
+Vanilla `kuleshov-group/mdlm-owt` MDLM-1M backbone, paper-aligned 128-sample gen-PPL eval under gpt2-large evaluator (matches EDLM paper's eval protocol):
 
-## Usage in EDLM training
+| Setup | sampler | steps | gen-PPL | factor below paper-61 |
+|---|---|---|---|---|
+| Vanilla MDLM (no train) | `ddpm_cache` (paper) | 1000 | 60.28 | replicates paper baseline |
+| Vanilla MDLM (no train) | `uncommitted_soft` (ours) | 256 | 17.04 | 3.5× |
+| **`edlm-step2k-pure-mdlm.ckpt`** | **`uncommitted_soft` (ours)** | **256** | **10.76** | **5.7×** |
 
-Configure `training.init_from_ckpt` to point at one of these files. The training loop (`main.py`) loads weights with `strict=False` and starts a fresh optimizer.
+## Winning training recipe
 
 ```yaml
-backbone: dit              # critical: matches Drive ckpt arch
+backbone: dit                          # (overridden by EBM hardcode → hf_dit)
+ebm_backbone: tiny_queries
 training:
-  init_from_ckpt: /path/to/mdlm.ckpt   # or ar.ckpt
+  k_max_final: 1                       # NO rollouts (standard MDLM)
+  soft_alpha: False                    # NO soft-α blending
+  loss_on_all_positions: True
+  threshold_commit_train: 0.9          # not used when k_max=1
+  force_commit_strategy_train: none
+sampling:
+  predictor: uncommitted_soft
+  steps: 256
+  threshold: 0.9
+  commit_sampling: True
+  commit_temperature: 1.0
+  force_commit_strategy: uniform
 ```
+
+## Negative results (avoid)
+
+- `k_max_final > 1`: training rollouts cause monotonic regression to PPL ~200 within 5k steps
+- `soft_alpha: True`: training-time soft-α blending drifts model away from OWT distribution
+- `force_commit_strategy_train: uniform`: similar regression; not additive
+- `predictor: ddpm_cache` at `steps: 256`: too few steps; gives PPL=87.58 instead of 60.28
+- `commit_temperature < 1.0`: mode collapse (PPL drops but text becomes repetitive)
+- `steps > 512` with `threshold < 0.9`: degenerate samples (low entropy, repetitive lists)
+
+## Compatibility notes
+
+- `mdlm.ckpt` and `ar.ckpt` (Drive originals) match the local `models/dit.py` arch (sigma_map.mlp.*) but **DO NOT load cleanly** into the standard EDLM EBM model — gen produces garbage at PPL=109788. Different layer ordering or normalization details in their training script.
+- `edlm-step2k-pure-mdlm.ckpt` is a Lightning ckpt produced by training on top of `kuleshov-group/mdlm-owt` with the recipe above. It loads cleanly via `Lightning.load_from_checkpoint` or `eval.partial_load_ckpt: True` in the EDLM trainer.