MetaDiffusion-150M-ChatBase

MetaDiffusion-150M-ChatBase is a masked-diffusion language model converted from an autoregressive base model and chat-tuned for downstream experimentation. It uses bidirectional attention with timestep conditioning and generates text through iterative masked denoising with left-to-right block commitment. Built by chat tuning CodeSoft/MetaDiffusion-150M-exp on smol-smoltalk (460K conversations, Apache-2.0) + no_robots (9.5K, Apache-2.0).

Architecture graph

• Architecture: 16L x 768W, 169.5M params, bidirectional attention, timestep conditioning, 32,010 vocab (ChatML tokens added), context 5120
• Weights: fp16 safetensors (339 MB), untied lm_head
• Generation: masked denoising (left-to-right block commit, 128 steps default)

Intended use

Downstream fine-tuning for specific tasks. This is NOT designed for production, instruction following is weak and coherent output is limited to roughly 60-100 tokens.

Quickstart (chat)

pip install -r scripts/requirements.txt

# Interactive, with live denoising view (--watch):
python scripts/chat.py --model-path . --watch

# One-shot:
python scripts/chat.py --model-path . --prompt "What is the capital of France?"

The --watch flag shows the response denoising in real time: step count, noise level t, masks remaining, and the partial text building into place. Generation defaults live in generation_config.json (128 steps, 96-token block, temperature 0.7, repetition penalty 1.5).

Fine-tune it

# 1. Your data:
#    a) HF dataset names (no_robots, alpaca, dolly, smol-smoltalk, math)
python scripts/prepare_data.py --model-path . \
    --data-dir my_data --datasets no_robots

#    b) Local ChatML files: .jsonl, .json, .parquet (messages/instruction shapes)
python scripts/convert_data.py --model-path . --input my_chat.jsonl \
    --output my_data
python scripts/convert_data.py --model-path . --input ./data_folder \
    --output my_data

# 2. Train
python scripts/train_chat.py --model-path . \
    --data-dir my_data --output-dir my_checkpoints \
    --lr 7e-5 --epochs 3 --patience 6

# 3. Chat with your model
python scripts/chat.py --model-path my_checkpoints/best.pt \
    --tokenizer my_data/tokenizer --watch

# 4. Export a new release artifact (self-packages the scripts too)
python scripts/export_hf.py --checkpoint my_checkpoints/best.pt \
    --tokenizer my_data/tokenizer --output ./MetaDiffusion-150M-MyTask --fp16

The whole pipeline runs on a single consumer GPU.

Behavior

The model returns coherent sentences but loses coherence and factuality over multi-turn conversations.

Generation notes

• Left-to-right (semi-autoregressive) block commit: confidence-based unmasking lets <|im_end|> win at position 0 and produced empty responses on this architecture at 150M; left-to-right fixed it (verified).
• Denoising stops early once <|im_end|> is committed.

Model lineage

MetaDiffusion-150M-ChatBase is derived through the following process:

1. Start from SupraLabs/Supra-1.5-50M-Base-exp.
2. Convert the autoregressive model into a masked-diffusion architecture.
3. Continue pre-training on FineWeb-EDU.
4. Expand the model to 150M parameters.
5. Continue pretraining on The Pile.
6. Chat-tune on HuggingFaceTB/smol-smoltalk and HuggingFaceH4/no_robots.

The resulting model uses bidirectional attention and timestep conditioning rather than conventional causal attention.

License details

Apache-2.0