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

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.

---

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