FLustc/multimdm-toy-reproduction

MultiMDM Toy Reproduction

This repository contains a minimal, low-resource reproduction of the core multi-mask forward corruption and clean-token reconstruction idea from “Multi-Mask Diffusion Language Models for Few-Step Generation.”

This repository is a minimal/toy reproduction of the core MultiMDM mechanism. It does not reproduce the full-scale experiments or reported paper numbers.

Implemented

• Eight configurable mask tokens
• Designated token-to-mask assignment: mask_start + ((clean_id - clean_start) % num_masks)
• The paper's multi-mask forward corruption with alpha(t) = beta(t) = 1 - t
• A time-conditioned Transformer encoder
• Masked clean-token reconstruction training
• An optional mask-class auxiliary objective
• Checkpoint save/load, reconstruction, approximate sampling, and evaluation

The auxiliary mask-class objective is an engineering approximation for the paper's intra-mask identification term, not an exact implementation of the full closed-form ELBO.

Not implemented

• Full-scale 170M training
• Exact full ELBO from Eq. (3.5)
• Shared-Gumbel consistency distillation
• Exact reverse-posterior sampling
• Original paper benchmark reproduction
• LLaDA-8B adaptation

The sampler is a simplified approximate denoising procedure and is not the exact posterior sampler or distilled few-step sampler from the paper.

Install

python -m venv .venv

Windows:

.venv\Scripts\activate

Linux/macOS:

source .venv/bin/activate

pip install -r requirements.txt

Train

python train.py

Low-resource smoke test:

python train.py --device cpu --d-model 64 --num-layers 2 \
  --batch-size 8 --seq-len 32 --max-steps 300

The trainer automatically reduces batch size, model width, sequence length, and finally falls back to CPU after a CUDA out-of-memory error.

Reconstruct

python generate.py --checkpoint outputs/model.pt \
  --text "multi mask diffusion preserves information" --time 0.7

Approximate iterative denoising

python generate.py --checkpoint outputs/model.pt \
  --mode sample --steps 8 --length 32

Evaluate

python evaluate.py --checkpoint outputs/model.pt

Upload

Authenticate without putting the token in source control:

hf auth login
python upload_to_hf.py --repo-id YOUR_USERNAME/multimdm-toy-reproduction

Intended use and limitations

Educational and pipeline-verification purposes only.

This checkpoint is trained on a tiny corpus for pipeline verification. It should not be treated as a competitive language model.