Fix model card: escape ~ so it renders literally (was rendering as strikethrough)
Browse files
README.md
CHANGED
|
@@ -1,160 +1,160 @@
|
|
| 1 |
-
---
|
| 2 |
-
license: apache-2.0
|
| 3 |
-
language:
|
| 4 |
-
- en
|
| 5 |
-
datasets:
|
| 6 |
-
- Skylion007/openwebtext
|
| 7 |
-
tags:
|
| 8 |
-
- masked-diffusion
|
| 9 |
-
- diffusion-language-model
|
| 10 |
-
- mamba
|
| 11 |
-
- mamba-2
|
| 12 |
-
- state-space-model
|
| 13 |
-
- mdlm
|
| 14 |
-
- text-generation
|
| 15 |
-
- pytorch-lightning
|
| 16 |
-
---
|
| 17 |
-
|
| 18 |
-
# DiffMamba β Checkpoints
|
| 19 |
-
|
| 20 |
-
Training checkpoints for **DiffMamba**, a small-scale independent study of
|
| 21 |
-
**bidirectional Mamba-2 (state-space) denoisers for Masked Diffusion Language
|
| 22 |
-
Models (MDLM)**. The Transformer/DiT denoiser in MDLM is replaced with a
|
| 23 |
-
**bidirectional Mamba-2 backbone**, and a matched set of models is trained from
|
| 24 |
-
scratch on OpenWebText for a controlled quality / scaling / efficiency comparison.
|
| 25 |
-
|
| 26 |
-
> **Code, full technical report, and documentation:**
|
| 27 |
-
> π **https://github.com/shivnarainms22/DiffMamba**
|
| 28 |
-
>
|
| 29 |
-
> This repo holds **weights only**. The GitHub repository is the source of
|
| 30 |
-
> truth for architecture, training recipe, evaluation, and the honest write-up
|
| 31 |
-
> of results and limitations.
|
| 32 |
-
|
| 33 |
-
This work builds directly on **MDLM** (Sahoo et al., NeurIPS 2024) and is a
|
| 34 |
-
small-scale reproduction of the research direction introduced by
|
| 35 |
-
**DiffuApriel / DiffuMamba** (arXiv 2511.15927). It is **not** claimed as a
|
| 36 |
-
novel architecture.
|
| 37 |
-
|
| 38 |
-
---
|
| 39 |
-
|
| 40 |
-
## What's in this repo
|
| 41 |
-
|
| 42 |
-
Six training runs, each in its own folder. Within a folder you'll find periodic
|
| 43 |
-
snapshots `step_<N>.ckpt` (every 5000 steps; every 3000 for the 50M run) and
|
| 44 |
-
`last.ckpt` (the final-step weights). These are **PyTorch Lightning
|
| 45 |
-
checkpoints** from the MDLM codebase β they bundle model weights *and* EMA
|
| 46 |
-
shadow parameters (EMA decay 0.9999), optimizer state, and config. They are
|
| 47 |
-
**not** `transformers`-loadable via `from_pretrained`; load them with the
|
| 48 |
-
training repo (see *How to use* below).
|
| 49 |
-
|
| 50 |
-
| Folder | Backbone | Params | LR | Steps | Tokens | Val PPL β |
|
| 51 |
-
|--------|----------|:------:|:--:|:-----:|:------:|:---------:|
|
| 52 |
-
| `runB_transformer_130m` | Transformer (DiT) | ~130M | 3e-4 | 76k | ~5B | **70.5** |
|
| 53 |
-
| `runD_130m_seed1` | BiMamba-2 (SSM) | ~130M | 3e-4 | 76k | ~5B | 85.9 |
|
| 54 |
-
| `runD_130m_seed2` | BiMamba-2 (SSM) | ~130M | 3e-4 | 76k | ~5B | 83.5 |
|
| 55 |
-
| `runD_130m_lr1e3_seed1` | BiMamba-2 (SSM) | ~130M | **1e-3** | 76k | ~5B | **79.3** |
|
| 56 |
-
| `scaling_100m` | BiMamba-2 (SSM) | ~100M | 3e-4 | 60k* | ~4B | 97.5 |
|
| 57 |
-
| `scaling_50m` | BiMamba-2 (SSM) | ~50M | 3e-4 | 30k | ~2B | 136.3 |
|
| 58 |
-
|
| 59 |
-
Val PPL = MDLM ELBO-bound validation perplexity on the OpenWebText validation
|
| 60 |
-
split, measured on each run's **best EMA checkpoint** (lower is better).
|
| 61 |
-
`*` the 100M run's valid final checkpoint is `step_60000`/`last.ckpt`
|
| 62 |
-
(see the GitHub report for why 61k looped).
|
| 63 |
-
|
| 64 |
-
### Results at a glance
|
| 65 |
-
|
| 66 |
-
- **Quality.** With the MDLM (Transformer-tuned) recipe at matched 130M / ~5B
|
| 67 |
-
tokens, the Transformer denoiser (70.5) is modestly but consistently stronger
|
| 68 |
-
than pure BiMamba-2. BiMamba prefers a **~3.3Γ higher learning rate**; a
|
| 69 |
-
50M LR sweep found `1e-3` best, and retraining 130M at `1e-3` (the
|
| 70 |
-
`runD_130m_lr1e3_seed1` checkpoints) closes **~43%** of the gap (85.9 β 79.3)
|
| 71 |
-
but does not close it.
|
| 72 |
-
- **Scaling** (BiMamba, lr 3e-4): 50M β 136.3, 100M β 97.5, 130M β 84.7 β
|
| 73 |
-
clean, monotonic, seed-stable (Ξβ2.4 between seeds).
|
| 74 |
-
- **Efficiency.** Forward-pass latency is **textbook-linear** in sequence
|
| 75 |
-
length for BiMamba vs. empirically O(L^1.55) for DiT (with FlashAttention);
|
| 76 |
-
crossover at ~3K tokens, **3.12Γ faster at 32K**.
|
| 77 |
-
- **Honest finding:** *pure* BiMamba-2 trades quality for long-context
|
| 78 |
-
throughput β consistent with DiffuApriel, where a *hybrid* Mamba+attention
|
| 79 |
-
model is what recovers quality.
|
| 80 |
-
|
| 81 |
-
Full numbers, caveats, and the LR-fairness analysis are in the
|
| 82 |
-
[technical report on GitHub](https://github.com/shivnarainms22/DiffMamba/blob/master/DiffMamba_Report.md).
|
| 83 |
-
|
| 84 |
-
---
|
| 85 |
-
|
| 86 |
-
## Model details
|
| 87 |
-
|
| 88 |
-
- **Framework:** MDLM β absorbing-state discrete diffusion, SUBS
|
| 89 |
-
parameterization, loglinear noise schedule, continuous time (T=0).
|
| 90 |
-
- **Tokenizer:** GPT-2 BPE (vocab 50257 + 1 mask token).
|
| 91 |
-
- **Sequence length:** 1024.
|
| 92 |
-
- **BiMamba-2 backbone** (`models/dimamba.py`): forward + flipped-reverse
|
| 93 |
-
Mamba-2 with weight-tied projections and **AdaLN** noise-level conditioning,
|
| 94 |
-
Mamba-2 defaults `d_state=64`, `headdim=64`, `cond_dim=128`, dropout 0.1.
|
| 95 |
-
- 130M = hidden 768 / 12 blocks Β· 100M = hidden 640 / 10 blocks Β·
|
| 96 |
-
50M = hidden 512 / 8 blocks.
|
| 97 |
-
- **Transformer baseline** (`models/dit.py`): DiT, hidden 768 / 12 blocks /
|
| 98 |
-
12 heads.
|
| 99 |
-
- **Training:** AdamW (wd 0.01, Ξ²=(0.9, 0.999), eps 1e-8), constant LR with
|
| 100 |
-
warmup, gradient clip 1.0, `bf16-mixed`, global batch 64 (micro-batch 16 Γ
|
| 101 |
-
grad-accum 4), EMA 0.9999, single A100 per run on an academic SLURM cluster
|
| 102 |
-
with 8-hour-wall checkpoint/resume job-chaining.
|
| 103 |
-
- **Data:** OpenWebText (`Skylion007/openwebtext`), GPT-2-tokenized,
|
| 104 |
-
~40:1 tokens-per-parameter recipe.
|
| 105 |
-
|
| 106 |
-
## How to use
|
| 107 |
-
|
| 108 |
-
These are Lightning checkpoints for the [DiffMamba / MDLM codebase](https://github.com/shivnarainms22/DiffMamba),
|
| 109 |
-
not `from_pretrained`-loadable. To evaluate or resume:
|
| 110 |
-
|
| 111 |
-
```bash
|
| 112 |
-
git clone https://github.com/shivnarainms22/DiffMamba
|
| 113 |
-
cd DiffMamba
|
| 114 |
-
# set up the environment (see requirements.yaml / scripts/)
|
| 115 |
-
|
| 116 |
-
# download a checkpoint, e.g. the LR-tuned BiMamba-130M
|
| 117 |
-
huggingface-cli download Shiv-22/diffmamba-checkpoints \
|
| 118 |
-
runD_130m_lr1e3_seed1/last.ckpt --local-dir ./ckpts
|
| 119 |
-
|
| 120 |
-
# validation perplexity (EMA), matching the table above
|
| 121 |
-
python main.py mode=ppl_eval +experiment=runD_130m \
|
| 122 |
-
eval.checkpoint_path=./ckpts/runD_130m_lr1e3_seed1/last.ckpt \
|
| 123 |
-
data.cache_dir=<path>/data loader.eval_batch_size=32
|
| 124 |
-
|
| 125 |
-
# generate samples
|
| 126 |
-
python main.py mode=sample_eval +experiment=runD_130m \
|
| 127 |
-
eval.checkpoint_path=./ckpts/runD_130m_lr1e3_seed1/last.ckpt \
|
| 128 |
-
loader.eval_batch_size=4
|
| 129 |
-
```
|
| 130 |
-
|
| 131 |
-
Use the matching `+experiment=` for each folder: `runD_130m` (BiMamba-130M and
|
| 132 |
-
its LR-tuned variant), `runB_transformer_130m` (DiT-130M), `scaling_100m`,
|
| 133 |
-
`scaling_50m`.
|
| 134 |
-
|
| 135 |
-
## Limitations
|
| 136 |
-
|
| 137 |
-
Small scale (50β130M, β€5B tokens), single-GPU academic compute, forward-pass-only
|
| 138 |
-
efficiency benchmark, and a Transformer-tuned training recipe that BiMamba is
|
| 139 |
-
shown to be undertuned for. Pure BiMamba-2 does **not** match the Transformer on
|
| 140 |
-
quality at this scale. Treat these as a reproduction/portfolio artifact, not a
|
| 141 |
-
production model. See the GitHub report for the full limitations section.
|
| 142 |
-
|
| 143 |
-
## Citation & attribution
|
| 144 |
-
|
| 145 |
-
Built on **MDLM** (Sahoo et al., *Simple and Effective Masked Diffusion Language
|
| 146 |
-
Models*, NeurIPS 2024; [code](https://github.com/kuleshov-group/mdlm)) and
|
| 147 |
-
reproduces the direction of **DiffuApriel / DiffuMamba**
|
| 148 |
-
(*High-Throughput Diffusion LMs with Mamba Backbone*, arXiv 2511.15927, 2025).
|
| 149 |
-
|
| 150 |
-
```bibtex
|
| 151 |
-
@inproceedings{sahoo2024simple,
|
| 152 |
-
title={Simple and Effective Masked Diffusion Language Models},
|
| 153 |
-
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},
|
| 154 |
-
booktitle={The Thirty-eighth Annual Conference on Neural Information Processing Systems},
|
| 155 |
-
year={2024},
|
| 156 |
-
url={https://openreview.net/forum?id=L4uaAR4ArM}
|
| 157 |
-
}
|
| 158 |
-
```
|
| 159 |
-
|
| 160 |
-
License: Apache-2.0 (inherited from MDLM).
|
|
|
|
| 1 |
+
---
|
| 2 |
+
license: apache-2.0
|
| 3 |
+
language:
|
| 4 |
+
- en
|
| 5 |
+
datasets:
|
| 6 |
+
- Skylion007/openwebtext
|
| 7 |
+
tags:
|
| 8 |
+
- masked-diffusion
|
| 9 |
+
- diffusion-language-model
|
| 10 |
+
- mamba
|
| 11 |
+
- mamba-2
|
| 12 |
+
- state-space-model
|
| 13 |
+
- mdlm
|
| 14 |
+
- text-generation
|
| 15 |
+
- pytorch-lightning
|
| 16 |
+
---
|
| 17 |
+
|
| 18 |
+
# DiffMamba β Checkpoints
|
| 19 |
+
|
| 20 |
+
Training checkpoints for **DiffMamba**, a small-scale independent study of
|
| 21 |
+
**bidirectional Mamba-2 (state-space) denoisers for Masked Diffusion Language
|
| 22 |
+
Models (MDLM)**. The Transformer/DiT denoiser in MDLM is replaced with a
|
| 23 |
+
**bidirectional Mamba-2 backbone**, and a matched set of models is trained from
|
| 24 |
+
scratch on OpenWebText for a controlled quality / scaling / efficiency comparison.
|
| 25 |
+
|
| 26 |
+
> **Code, full technical report, and documentation:**
|
| 27 |
+
> π **https://github.com/shivnarainms22/DiffMamba**
|
| 28 |
+
>
|
| 29 |
+
> This repo holds **weights only**. The GitHub repository is the source of
|
| 30 |
+
> truth for architecture, training recipe, evaluation, and the honest write-up
|
| 31 |
+
> of results and limitations.
|
| 32 |
+
|
| 33 |
+
This work builds directly on **MDLM** (Sahoo et al., NeurIPS 2024) and is a
|
| 34 |
+
small-scale reproduction of the research direction introduced by
|
| 35 |
+
**DiffuApriel / DiffuMamba** (arXiv 2511.15927). It is **not** claimed as a
|
| 36 |
+
novel architecture.
|
| 37 |
+
|
| 38 |
+
---
|
| 39 |
+
|
| 40 |
+
## What's in this repo
|
| 41 |
+
|
| 42 |
+
Six training runs, each in its own folder. Within a folder you'll find periodic
|
| 43 |
+
snapshots `step_<N>.ckpt` (every 5000 steps; every 3000 for the 50M run) and
|
| 44 |
+
`last.ckpt` (the final-step weights). These are **PyTorch Lightning
|
| 45 |
+
checkpoints** from the MDLM codebase β they bundle model weights *and* EMA
|
| 46 |
+
shadow parameters (EMA decay 0.9999), optimizer state, and config. They are
|
| 47 |
+
**not** `transformers`-loadable via `from_pretrained`; load them with the
|
| 48 |
+
training repo (see *How to use* below).
|
| 49 |
+
|
| 50 |
+
| Folder | Backbone | Params | LR | Steps | Tokens | Val PPL β |
|
| 51 |
+
|--------|----------|:------:|:--:|:-----:|:------:|:---------:|
|
| 52 |
+
| `runB_transformer_130m` | Transformer (DiT) | \~130M | 3e-4 | 76k | \~5B | **70.5** |
|
| 53 |
+
| `runD_130m_seed1` | BiMamba-2 (SSM) | \~130M | 3e-4 | 76k | \~5B | 85.9 |
|
| 54 |
+
| `runD_130m_seed2` | BiMamba-2 (SSM) | \~130M | 3e-4 | 76k | \~5B | 83.5 |
|
| 55 |
+
| `runD_130m_lr1e3_seed1` | BiMamba-2 (SSM) | \~130M | **1e-3** | 76k | \~5B | **79.3** |
|
| 56 |
+
| `scaling_100m` | BiMamba-2 (SSM) | \~100M | 3e-4 | 60k* | \~4B | 97.5 |
|
| 57 |
+
| `scaling_50m` | BiMamba-2 (SSM) | \~50M | 3e-4 | 30k | \~2B | 136.3 |
|
| 58 |
+
|
| 59 |
+
Val PPL = MDLM ELBO-bound validation perplexity on the OpenWebText validation
|
| 60 |
+
split, measured on each run's **best EMA checkpoint** (lower is better).
|
| 61 |
+
`*` the 100M run's valid final checkpoint is `step_60000`/`last.ckpt`
|
| 62 |
+
(see the GitHub report for why 61k looped).
|
| 63 |
+
|
| 64 |
+
### Results at a glance
|
| 65 |
+
|
| 66 |
+
- **Quality.** With the MDLM (Transformer-tuned) recipe at matched 130M / \~5B
|
| 67 |
+
tokens, the Transformer denoiser (70.5) is modestly but consistently stronger
|
| 68 |
+
than pure BiMamba-2. BiMamba prefers a **\~3.3Γ higher learning rate**; a
|
| 69 |
+
50M LR sweep found `1e-3` best, and retraining 130M at `1e-3` (the
|
| 70 |
+
`runD_130m_lr1e3_seed1` checkpoints) closes **\~43%** of the gap (85.9 β 79.3)
|
| 71 |
+
but does not close it.
|
| 72 |
+
- **Scaling** (BiMamba, lr 3e-4): 50M β 136.3, 100M β 97.5, 130M β 84.7 β
|
| 73 |
+
clean, monotonic, seed-stable (Ξβ2.4 between seeds).
|
| 74 |
+
- **Efficiency.** Forward-pass latency is **textbook-linear** in sequence
|
| 75 |
+
length for BiMamba vs. empirically O(L^1.55) for DiT (with FlashAttention);
|
| 76 |
+
crossover at \~3K tokens, **3.12Γ faster at 32K**.
|
| 77 |
+
- **Honest finding:** *pure* BiMamba-2 trades quality for long-context
|
| 78 |
+
throughput β consistent with DiffuApriel, where a *hybrid* Mamba+attention
|
| 79 |
+
model is what recovers quality.
|
| 80 |
+
|
| 81 |
+
Full numbers, caveats, and the LR-fairness analysis are in the
|
| 82 |
+
[technical report on GitHub](https://github.com/shivnarainms22/DiffMamba/blob/master/DiffMamba_Report.md).
|
| 83 |
+
|
| 84 |
+
---
|
| 85 |
+
|
| 86 |
+
## Model details
|
| 87 |
+
|
| 88 |
+
- **Framework:** MDLM β absorbing-state discrete diffusion, SUBS
|
| 89 |
+
parameterization, loglinear noise schedule, continuous time (T=0).
|
| 90 |
+
- **Tokenizer:** GPT-2 BPE (vocab 50257 + 1 mask token).
|
| 91 |
+
- **Sequence length:** 1024.
|
| 92 |
+
- **BiMamba-2 backbone** (`models/dimamba.py`): forward + flipped-reverse
|
| 93 |
+
Mamba-2 with weight-tied projections and **AdaLN** noise-level conditioning,
|
| 94 |
+
Mamba-2 defaults `d_state=64`, `headdim=64`, `cond_dim=128`, dropout 0.1.
|
| 95 |
+
- 130M = hidden 768 / 12 blocks Β· 100M = hidden 640 / 10 blocks Β·
|
| 96 |
+
50M = hidden 512 / 8 blocks.
|
| 97 |
+
- **Transformer baseline** (`models/dit.py`): DiT, hidden 768 / 12 blocks /
|
| 98 |
+
12 heads.
|
| 99 |
+
- **Training:** AdamW (wd 0.01, Ξ²=(0.9, 0.999), eps 1e-8), constant LR with
|
| 100 |
+
warmup, gradient clip 1.0, `bf16-mixed`, global batch 64 (micro-batch 16 Γ
|
| 101 |
+
grad-accum 4), EMA 0.9999, single A100 per run on an academic SLURM cluster
|
| 102 |
+
with 8-hour-wall checkpoint/resume job-chaining.
|
| 103 |
+
- **Data:** OpenWebText (`Skylion007/openwebtext`), GPT-2-tokenized,
|
| 104 |
+
\~40:1 tokens-per-parameter recipe.
|
| 105 |
+
|
| 106 |
+
## How to use
|
| 107 |
+
|
| 108 |
+
These are Lightning checkpoints for the [DiffMamba / MDLM codebase](https://github.com/shivnarainms22/DiffMamba),
|
| 109 |
+
not `from_pretrained`-loadable. To evaluate or resume:
|
| 110 |
+
|
| 111 |
+
```bash
|
| 112 |
+
git clone https://github.com/shivnarainms22/DiffMamba
|
| 113 |
+
cd DiffMamba
|
| 114 |
+
# set up the environment (see requirements.yaml / scripts/)
|
| 115 |
+
|
| 116 |
+
# download a checkpoint, e.g. the LR-tuned BiMamba-130M
|
| 117 |
+
huggingface-cli download Shiv-22/diffmamba-checkpoints \
|
| 118 |
+
runD_130m_lr1e3_seed1/last.ckpt --local-dir ./ckpts
|
| 119 |
+
|
| 120 |
+
# validation perplexity (EMA), matching the table above
|
| 121 |
+
python main.py mode=ppl_eval +experiment=runD_130m \
|
| 122 |
+
eval.checkpoint_path=./ckpts/runD_130m_lr1e3_seed1/last.ckpt \
|
| 123 |
+
data.cache_dir=<path>/data loader.eval_batch_size=32
|
| 124 |
+
|
| 125 |
+
# generate samples
|
| 126 |
+
python main.py mode=sample_eval +experiment=runD_130m \
|
| 127 |
+
eval.checkpoint_path=./ckpts/runD_130m_lr1e3_seed1/last.ckpt \
|
| 128 |
+
loader.eval_batch_size=4
|
| 129 |
+
```
|
| 130 |
+
|
| 131 |
+
Use the matching `+experiment=` for each folder: `runD_130m` (BiMamba-130M and
|
| 132 |
+
its LR-tuned variant), `runB_transformer_130m` (DiT-130M), `scaling_100m`,
|
| 133 |
+
`scaling_50m`.
|
| 134 |
+
|
| 135 |
+
## Limitations
|
| 136 |
+
|
| 137 |
+
Small scale (50β130M, β€5B tokens), single-GPU academic compute, forward-pass-only
|
| 138 |
+
efficiency benchmark, and a Transformer-tuned training recipe that BiMamba is
|
| 139 |
+
shown to be undertuned for. Pure BiMamba-2 does **not** match the Transformer on
|
| 140 |
+
quality at this scale. Treat these as a reproduction/portfolio artifact, not a
|
| 141 |
+
production model. See the GitHub report for the full limitations section.
|
| 142 |
+
|
| 143 |
+
## Citation & attribution
|
| 144 |
+
|
| 145 |
+
Built on **MDLM** (Sahoo et al., *Simple and Effective Masked Diffusion Language
|
| 146 |
+
Models*, NeurIPS 2024; [code](https://github.com/kuleshov-group/mdlm)) and
|
| 147 |
+
reproduces the direction of **DiffuApriel / DiffuMamba**
|
| 148 |
+
(*High-Throughput Diffusion LMs with Mamba Backbone*, arXiv 2511.15927, 2025).
|
| 149 |
+
|
| 150 |
+
```bibtex
|
| 151 |
+
@inproceedings{sahoo2024simple,
|
| 152 |
+
title={Simple and Effective Masked Diffusion Language Models},
|
| 153 |
+
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},
|
| 154 |
+
booktitle={The Thirty-eighth Annual Conference on Neural Information Processing Systems},
|
| 155 |
+
year={2024},
|
| 156 |
+
url={https://openreview.net/forum?id=L4uaAR4ArM}
|
| 157 |
+
}
|
| 158 |
+
```
|
| 159 |
+
|
| 160 |
+
License: Apache-2.0 (inherited from MDLM).
|