Crimson-LaDa

A masked diffusion language model based on the LLaDA paper ("LLaDA: Large Language Diffusion with Masking"), trained on general instruction-following datasets.

Unlike autoregressive models (GPT-style) that generate left-to-right, Crimson-LaDa uses iterative denoising: starting from fully masked text, it progressively recovers tokens over multiple steps, allowing the model to revise earlier decisions.

Model Details

Parameter	Value
Layers	8
Hidden dim	384
Attention heads	8
Parameters	15.3M
Vocab size	2,685 (character-level)
Sequence length	256
Training steps	5,000
Best val loss	1.63
Schedule	WSD (warmup-stable-decay)
Mask strategy	t ~ U[ε, 1-ε] (LLaDA-style)

Training Data

Trained on 6,000 instruction-following examples from 4 datasets:

• databricks/databricks-dolly-15k: 3,000 examples of human-generated instruction/response pairs
• timdettmers/openassistant-guanaco: 3,000 examples of assistant-style conversations

Usage

import torch
import json
from model import MDLM, generate

# Load
model = MDLM()
ckpt = torch.load("model_final.pt", map_location="cpu")
model.load_state_dict(ckpt, strict=False)
model.eval()

# Generate
output = generate(
    model, tokenizer,
    prompt="user: What is the capital of France?\n\nassistant:",
    max_len=96, steps=64, temperature=0.8
)
print(output)

Architecture

Standard decoder-only transformer with:

• Pre-norm (LayerNorm before attention/MLP)
• Weight tying (input embeddings = output head)
• GELU activations
• Learned positional embeddings
• No causal masking (bidirectional attention)

Inference

Uses the LLaDA low-confidence remasking strategy:

1. Start with all non-prompt tokens masked
2. At each step, predict all positions and sample
3. Re-mask the 50% lowest-confidence positions
4. Repeat for N steps until all tokens are unmasked

Limitations

• Character-level tokenizer: Limits semantic understanding. A BPE tokenizer would improve quality.
• Small model (15M): Appropriate for experimentation but not production use.
• Limited training data: Only 6K examples from 4 datasets.
• CPU-friendly: Small enough to run inference efficiently on CPU.