latent-space-language-diffusion-model

<!-- SEO: keywords for search discoverability --> <!-- keywords: masked diffusion language model, MDLM, parallel text generation, diffusion model, non-autoregressive generation, bidirectional transformer, AR oracle validation, speculative decoding, Qwen3, SplatsDB, open weights, open source LLM, torch, safetensors description: A 201M parameter masked diffusion language model that generates text by predicting all tokens in parallel. Validated by an autoregressive oracle (Qwen3-0.6B). Open weights, open code, fully reproducible. -->

💻 Full source code, experiments, benchmarks, and research lineage on GitHub — 8 phases of experiments, training scripts, ablation studies, and honest negative results.

Latent Space Language Diffusion Model

Open weights. Open code. Fully reproducible.

A 201M parameter masked diffusion language model that generates text by predicting all tokens in parallel. Validated by an autoregressive oracle (Qwen3-0.6B). Trained on a single consumer GPU (RTX 3090) in ~7 hours.

📊 Key Results

🔓 What's Included (Everything Open)

No black boxes. Every line of model code, training code, and inference code is included. The checkpoint is the exact state after 49,779 optimizer steps — nothing hidden, no cherry-picking.

🏗️ Architecture

MDLM-BPE v3 (201M params)
├── Token embedding: 16K BPE vocab → 1024 dims
├── 10× Transformer blocks
│   ├── RoPE positional encoding (non-learned)
│   ├── AdaLN timestep conditioning (FiLM-style)
│   └── Flash attention (non-causal / bidirectional)
├── LayerNorm + output projection (weight-tied)
└── Full-parallel unmasking (all positions simultaneously)

Design philosophy: Unlike autoregressive models that generate tokens left-to-right (one at a time), this model predicts ALL positions simultaneously via iterative diffusion. This enables 1.8-3.9× faster forward passes and 2.4× faster generation.

Generation Pipeline

1. MDLM full-parallel diffusion (32 steps)
   All masked positions predicted simultaneously each step

2. Adaptive guidance (during generation):
   • Frequency penalty (0.4): penalize overused tokens
   • Repetition penalty (1.3): reduce already-used token logits
   • No-repeat bigram: hard ban on repeated 2-grams
   • Top-p (0.95): nucleus sampling, filter tail noise

3. Qwen3-0.6B validation (optional, 1 forward pass):
   Teacher-forcing score per segment

4. Segment regeneration (optional):
   Bad segments regenerated by Qwen3 AR model

🔬 Training Details (Full Transparency)

Data

• Source: Ultra-FineWeb by OpenBMB
• Documents: 1,000,000 (from 2M downloaded, 470,122 rejected by quality filter)
• Tokens: 272M (packed into 2.12M sequences of 128 tokens)
• Language: English
• Filtering: Length >50 chars, deduplicated

Hyperparameters

Training Curve

Step      0:  loss=7.97  PPL=1341    (random init)
Step   1000:  loss=7.20  PPL=711     (learning structure)
Step  23000:  loss=5.00  PPL=137.6   (coherent phrases)
Step  38000:  loss=4.84  PPL=126.7   (refining)
Step  44000:  loss=4.63  PPL=102.6   ← BEST checkpoint
Step  49779:  loss=4.80  PPL=112.6   (final, cosine decay end)

Hardware

• GPU: NVIDIA GeForce RTX 3090 (24GB VRAM)
• System RAM: 8GB (constrained — required memmap data loading)
• Training time: 6 hours 45 minutes
• Throughput: 33,500 tokens/second, 2.0 optimizer steps/second
• No distributed training: Single workstation

📏 Benchmark [MEASURED]

All metrics measured on RTX 3090. Oracle log-prob = mean per-token log-probability under Qwen3-0.6B teacher forcing (higher = more coherent).

Generation: Speed vs Quality

Forward-Pass Throughput (Raw Inference)

⚖️ Honest Limitations

This model is a research artifact, not a production system. The quality gap vs production models is real:

The quality gap is driven by scale (parameters + data), not architecture. This model was trained on 272M tokens on one GPU in 7 hours. Production models train on trillions of tokens across clusters. The throughput advantage (1.8-3.9× faster) is architecture-level and would compound at scale.

What Works

• ✅ Repetition elimination: Adaptive guidance (0.79 → 0.99) at zero cost
• ✅ Parallel speed: 2.4× faster generation than sequential decoding
• ✅ Oracle validation: Qwen3 segment regeneration improves coherence (+0.22 log-prob)

What Doesn't Work (Documented Failures)

• ❌ Embedding-based drift detection: 201M model's embeddings too weak
• ❌ Token-level oracle replacement: Cross-vocabulary mapping breaks coherence
• ❌ Guided MDLM regeneration: Model too weak to generate good replacements even with oracle bias

🚀 Usage

import torch
from safetensors.torch import load_file
from modeling_mdlm import MDLMConfig, MDLMBPEV3, BPETokenizer, sample_semi_ar

# Load model
config = MDLMConfig(d_model=1024, n_heads=16, n_layers=10, vocab_size=16000)
model = MDLMBPEV3(config)
state_dict = load_file("model.safetensors")
model.load_state_dict(state_dict)
model.eval().cuda()

# Load tokenizer
tokenizer = BPETokenizer("tokenizer.json")

# Generate
prompt_ids = tokenizer.encode("The future of artificial intelligence", add_special=False)
result = sample_semi_ar(
    model, tokenizer,
    prompt_ids=prompt_ids,
    seq_len=64,
    block_size=4,
    temperature=0.7,
)
print(result[0])

Full-parallel generation with adaptive guidance

from logit_guidance import sample_with_guidance

result = sample_with_guidance(
    model, tokenizer,
    prompt_ids=tokenizer.encode("Climate change is", add_special=False),
    seq_len=64,
    repetition_penalty=1.3,
    no_repeat_ngram=2,
    frequency_penalty=0.4,
)
print(result[0])

🔄 Reproduce Training

# All code is in this repo. To train from scratch:
# 1. Download data
python -c "from datasets import load_dataset; ..."

# 2. Train (see train.py for full args)
python train.py --epochs 3 --batch-size 32 --seq-len 128

# Expected: ~7 hours on RTX 3090, PPL ~100-110

📚 Citation

@misc{schwabauer2026latentspace,
  title={Latent Space Language Diffusion Model: Parallel Text Generation via Masked Diffusion with AR-Oracle Validation},
  author={Schwabauer, Brian Tomas},
  year={2026},
  publisher={HuggingFace},
  url={https://huggingface.co/brianschwabauer/latent-space-language-diffusion-model}
}

📜 License

MIT — Full open source. Use, modify, distribute freely.

🔗 Related Projects

• GitHub Repository — Full source code, experiments, benchmarks
• SplatsDB — Vector memory infrastructure
• MDLM Logit Guidance — Adaptive guidance module
• Ultra-FineWeb — Training data
• Qwen3-0.6B — AR oracle validation model

References

• Sahoo et al. (2024). Simple and Effective Masked Diffusion Language Models. arXiv:2406.03709
• Hinton (2022). The Forward-Forward Algorithm. arXiv:2212.13345 (explored and refuted for generation in Phase 1)