Caffin/bert-dllm
BERT dLLM · Masked Diffusion
An interactive reproduction of Nathan Barry's "BERT is just a Single Text Diffusion Step". It repurposes a RoBERTa masked-language model as a discrete language diffusion model (dLLM): instead of generating left-to-right, the model repeatedly predicts a whole masked canvas in parallel.
What is reproduced
- Backbone:
roberta-base, a BERT-style bidirectional masked-language model. - Training canvas: 256 tokens, with the first 16 tokens preserved as a fixed prompt prefix.
- Noise objective: one mask probability sampled per batch from
1.0, 0.9, ..., 0.1; labels exist only at masked positions. - Reverse process: start with 240 masks, fill masked positions in parallel, then re-mask exactly 90% down to 0% over ten passes.
- Trace: the UI renders the actual masks from the same inference run—no second model call or illustrative simulation.
The default Article reproduction decoder uses random re-masking. The optional confidence-aware mode preserves the model's most confident predictions, which is a standard improvement to parallel masked decoding but is deliberately separate from the reproduction mode.
Training the checkpoint
train.py contains the self-contained training pipeline. It starts from roberta-base, fine-tunes on WikiText-2 by default, and optionally pushes the resulting checkpoint to the model repository used by the app:
python -m pip install torch -r requirements-train.txt
python train.py \
--max-steps 1000 \
--hub-model-id Caffin/bert-dllm-wikitext2The original article reports a short H200 fine-tuning run; this implementation keeps the same core objective while making the data source and run length explicit. It is a research demo, not a replacement for a production autoregressive language model.