Architecture
MetaDiffusion-150M-exp is a diffusion language model created by converting the Supra-1.5-50M-Base-exp autoregressive model into a diffusion language model, then expanding it to approximately 150 million parameters through layer duplication before continued diffusion training.
The model retains the original tokenizer while introducing diffusion-specific components, including timestep conditioning and a learned mask token.
Architecture graph
| Specification | Value |
|---|---|
| Architecture | MetaDiffusionForCausalLM |
| Parameters | ~169.5M |
| Vocabulary Size | 32,001 (32,000 + MASK token) |
| Hidden Size | 768 |
| Intermediate Size | 2,112 |
| Layers | 16 |
| Attention Heads | 12 |
| KV Heads | 6 |
| Head Dimension | 64 |
| Context Length | 5,120 tokens |
| Tokenizer | Original Supra byte-level BPE tokenizer |
| Activation | SiLU |
| RoPE θ | 10,000 |
| Model Type | Diffusion Language Model |
Data
Training was performed in two stages:
- Autoregressive-to-diffusion conversion using the Supra-1.5-50M-Base-exp checkpoint
- Model expansion from approximately 50M to 150M parameters via layer duplication, followed by continued diffusion pretraining.
Training schedule:
| Stage | Dataset | Steps |
|---|---|---|
| Initial diffusion training | FineWeb-EDU | 100,000 |
| Continued diffusion training | The Pile | 50,000 |
Benchmarks
Evaluation: lm-evaluation-harness (0-shot)
| Benchmark | Samples | Accuracy | Normalized Accuracy |
|---|---|---|---|
| ARC-Easy | 2,376 | 38.30% | 35.40% |
| ARC-Challenge | 1,172 | 18.86% | 21.84% |
| ArithMark-3 | 1,000 | 31.30% | 31.60% |
Average Scores
| Metric | Score |
|---|---|
| Average Accuracy | 29.49% |
| Average Normalized Accuracy | 29.61% |
Running the Model
This repository includes an inference.py script for sampling from the model.
Example:
python inference.py \
--model-path ./model.safetensors \
--prompt "The cat sat on the" \
--num-steps 256
Command Line Arguments
| Flag | Description | Default |
|---|---|---|
--model-path | Path to a local checkpoint or Hugging Face repository | Required |
--prompt | Input prompt | "The cat sat on the" |
--seq-len | Number of generated tokens | 256 |
--num-steps | Number of diffusion denoising steps | 256 |
--temperature | Sampling temperature | 0.6 |
--repetition-penalty | Repetition penalty | 1.5 |
--device | cuda or cpu | cuda |
--watch | Display intermediate denoising progress | Disabled |
--watch-every | Display every N denoising iterations | 4 |
--base-model | Tokenizer source | SupraLabs/Supra-1.5-50M-Base-exp |
Generation Defaults
| Setting | Value |
|---|---|
| Denoising Steps | 512 |
| Temperature | 0.6 |
| Repetition Penalty | 1.5 |
| Re-mask Ratio | 0.1 |
| Max New Tokens | 512 |
Intended Use
MetaDiffusion-150M-exp is intended for:
- Research on diffusion language models.
- Experiments involving autoregressive-to-diffusion conversion.
- Benchmarking novel diffusion language model architectures.
- Further finetuning and experimentation.
This release is not instruction tuned and is not intended for production deployments.
Acknowledgements
MetaDiffusion is derived from the Supra-1.5-50M-Base-exp model. Credit goes to the Supra authors for the original autoregressive checkpoint that served as the initialization for this work. MetaDiffusion is released under the Apache-2.0 license in accordance with the licensing terms of the original Supra checkpoint.