Shiv-22 commited on
Commit
2e70ce6
Β·
verified Β·
1 Parent(s): a3dcde9

Fix model card: escape ~ so it renders literally (was rendering as strikethrough)

Browse files
Files changed (1) hide show
  1. README.md +160 -160
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).