Bigjay101/mdlm-african-lang

Research fork of MDLM (Sahoo et al., 2024) exploring masked diffusion language modelling for low-resource African languages. Training and evaluation on Swahili. # Simple and Effective Masked Diffusion Language Models (NeurIPS 2024) By Subham Sekhar Sahoo, Marianne Arriola, Yair Schiff, Aaron Gokaslan, Edgar Marroquin, Justin T Chiu, Alexander Rush, Volodymyr Kuleshov **Update April 14, 2025: An improved implementation is available here:** [DUO Github repo.] **Update Jun 3, 2025: MDMs with KV caching:** [Eso-LMs Github repo.] We introduce *MDLM*, a **M**asked discrete **D**iffusion **L**anguage **M**odel that features a novel (SUBS)titution based parameterization which simplifies the absorbing state diffusion loss to a mixture of classical masked language modeling losses. In doing so, we achieve SOTA perplexity numbers on LM1B and OpenWebText among diffusion models while achiving competitive zero-shot perplexity with SOTA AR models on numerous datasets. We provide a demo in this notebook or and a video tutorial here: In this repo, we release: * **The MDLM framework.** 1. SUBStitution based parameterization 2. Simplified loss calculation for masked diffusion processes * **Baseline implementations** [[Examples]](#baselines): 1. Autoregressive model that matches the SOTA AR performance on LM1B. 2. Score Entropy Based Discrete Diffusion SEDD. 3. An efficient implementation of the absorbing state D3PM that beats the previous SOTA text diffusion model SEDD on LM1B. * **Samplers** 1. Ancestral sampling as proposed in D3PM. 2. Analytic sampler as proposed in SEDD. 3. Our proposed efficient sampler that - makes MDLM **~3-4x** faster than the existing diffusion models. [[Example]](#sample-gen) - supports semi-autoregressive (SAR) generation. [[Example]](#semi-ar-gen) ## Code Organization 1. ```main.py```: Routines for training and evaluation 2. ```noise_schedule.py```: Noise schedules 3. ```diffusion.py```: Forward/reverse diffusion 4. ```dataloader.py```: Dataloaders 5. ```utils.py```: LR scheduler, logging, `fsspec` handling 6. ```models/```: Denoising network architectures. Supports DiT, AR transformer, an …