Shiv-22/diffmamba-checkpoints

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.

Full numbers, caveats, and the LR-fairness analysis are in the technical report on GitHub.

---

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, not from_pretrained-loadable. To evaluate or resume:

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) and reproduces the direction of DiffuApriel / DiffuMamba (High-Throughput Diffusion LMs with Mamba Backbone, arXiv 2511.15927, 2025).

@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).