Add model card

README.md

@@ -0,0 +1,160 @@
+---
+license: apache-2.0
+language:
+- en
+datasets:
+- Skylion007/openwebtext
+tags:
+- masked-diffusion
+- diffusion-language-model
+- mamba
+- mamba-2
+- state-space-model
+- mdlm
+- text-generation
+- pytorch-lightning
+---
+
+# DiffMamba — Checkpoints
+
+Training checkpoints for **DiffMamba**, a small-scale independent study of
+**bidirectional Mamba-2 (state-space) denoisers for Masked Diffusion Language
+Models (MDLM)**. The Transformer/DiT denoiser in MDLM is replaced with a
+**bidirectional Mamba-2 backbone**, and a matched set of models is trained from
+scratch on OpenWebText for a controlled quality / scaling / efficiency comparison.
+
+> **Code, full technical report, and documentation:**
+> 👉 **https://github.com/shivnarainms22/DiffMamba**
+>
+> This repo holds **weights only**. The GitHub repository is the source of
+> truth for architecture, training recipe, evaluation, and the honest write-up
+> of results and limitations.
+
+This work builds directly on **MDLM** (Sahoo et al., NeurIPS 2024) and is a
+small-scale reproduction of the research direction introduced by
+**DiffuApriel / DiffuMamba** (arXiv 2511.15927). It is **not** claimed as a
+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 ↓ |
+|--------|----------|:------:|:--:|:-----:|:------:|:---------:|
+| `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 |
+
+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).
+
+### Results at a glance
+
+- **Quality.** With the MDLM (Transformer-tuned) recipe at matched 130M / ~5B
+  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)
+  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).
+- **Efficiency.** Forward-pass latency is **textbook-linear** in sequence
+  length for BiMamba vs. empirically O(L^1.55) for DiT (with FlashAttention);
+  crossover at ~3K tokens, **3.12× faster at 32K**.
+- **Honest finding:** *pure* BiMamba-2 trades quality for long-context
+  throughput — consistent with DiffuApriel, where a *hybrid* Mamba+attention
+  model is what recovers quality.
+
+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).
+
+---
+
+## Model details
+
+- **Framework:** MDLM — absorbing-state discrete diffusion, SUBS
+  parameterization, loglinear noise schedule, continuous time (T=0).
+- **Tokenizer:** GPT-2 BPE (vocab 50257 + 1 mask token).
+- **Sequence length:** 1024.
+- **BiMamba-2 backbone** (`models/dimamba.py`): forward + flipped-reverse
+  Mamba-2 with weight-tied projections and **AdaLN** noise-level conditioning,
+  Mamba-2 defaults `d_state=64`, `headdim=64`, `cond_dim=128`, dropout 0.1.
+  - 130M = hidden 768 / 12 blocks · 100M = hidden 640 / 10 blocks ·
+    50M = hidden 512 / 8 blocks.
+- **Transformer baseline** (`models/dit.py`): DiT, hidden 768 / 12 blocks /
+  12 heads.
+- **Training:** AdamW (wd 0.01, β=(0.9, 0.999), eps 1e-8), constant LR with
+  warmup, gradient clip 1.0, `bf16-mixed`, global batch 64 (micro-batch 16 ×
+  grad-accum 4), EMA 0.9999, single A100 per run on an academic SLURM cluster
+  with 8-hour-wall checkpoint/resume job-chaining.
+- **Data:** OpenWebText (`Skylion007/openwebtext`), GPT-2-tokenized,
+  ~40:1 tokens-per-parameter recipe.
+
+## How to use
+
+These are Lightning checkpoints for the [DiffMamba / MDLM codebase](https://github.com/shivnarainms22/DiffMamba),
+not `from_pretrained`-loadable. To evaluate or resume:
+
+```bash
+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
+huggingface-cli download Shiv-22/diffmamba-checkpoints \
+  runD_130m_lr1e3_seed1/last.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 \
+  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 \
+  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`.
+
+## Limitations
+
+Small scale (50–130M, ≤5B tokens), single-GPU academic compute, forward-pass-only
+efficiency benchmark, and a Transformer-tuned training recipe that BiMamba is
+shown to be undertuned for. Pure BiMamba-2 does **not** match the Transformer on
+quality at this scale. Treat these as a reproduction/portfolio artifact, not a
+production model. See the GitHub report for the full limitations section.
+
+## Citation & attribution
+
+Built on **MDLM** (Sahoo et al., *Simple and Effective Masked Diffusion Language
+Models*, NeurIPS 2024; [code](https://github.com/kuleshov-group/mdlm)) and
+reproduces the direction of **DiffuApriel / DiffuMamba**
+(*High-Throughput Diffusion LMs with Mamba Backbone*, arXiv 2511.15927, 2025).
+
+```bibtex
+@inproceedings{sahoo2024simple,
+  title={Simple and Effective Masked Diffusion Language Models},
+  author={Subham Sekhar Sahoo and Marianne Arriola and Aaron Gokaslan and Edgar Mariano Marroquin and Alexander M Rush and Yair Schiff and Justin T Chiu and Volodymyr Kuleshov},
+  booktitle={The Thirty-eighth Annual Conference on Neural Information Processing Systems},
+  year={2024},
+  url={https://openreview.net/forum?id=L4uaAR4ArM}
+}
+```
+
+License: Apache-2.0 (inherited from MDLM).