| ---
|
| 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
|
|
|
| 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 (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.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
|
|
|
| - **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_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).
|
| - **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.
|
| - **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).
|
|
|
| ---
|
|
|
| ## 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 (best.ckpt reproduces
|
| # the table PPL; --local-dir preserves the runs/β¦/checkpoints/ path)
|
| huggingface-cli download Shiv-22/diffmamba-checkpoints \
|
| 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/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/runs/runD_lr1e3/checkpoints/best.ckpt \
|
| loader.eval_batch_size=4
|
| ```
|
|
|
| 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
|
|
|
| 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).
|
|
|
